由朝阳河、李申、马赫迪·索尔坦奥利科塔比和萨曼·阿维斯提梅尔著

在这篇博客文章中，我们描述了第一篇经过同行评审的研究论文，该论文探讨了加速 PyTorch DDP（ torch.nn.parallel.DistributedDataParallel ）[1]和 Pipeline（ torch.distributed.pipeline ）- PipeTransformer：用于大规模模型（如 BERT [2]和 ViT [3]）分布式训练的自动弹性管道的混合，该论文发表于 ICML 2021。

PipeTransformer 利用自动弹性管道化，以高效地分布式训练 Transformer 模型。在 PipeTransformer 中，我们设计了一种自适应的即时冻结算法，可以在训练过程中逐步识别并冻结一些层，以及一个弹性管道化系统，可以动态分配资源以训练剩余的活跃层。更具体地说，PipeTransformer 会自动将冻结层排除在管道之外，将活跃层打包到更少的 GPU 上，并通过更多副本来增加数据并行宽度。我们使用 Vision Transformer (ViT)在 ImageNet 上以及 BERT 在 SQuAD 和 GLUE 数据集上评估 PipeTransformer。我们的结果表明，与最先进的基线相比，PipeTransformer 在保持准确性的同时，速度提高了高达 2.83 倍。我们还提供了各种性能分析，以更全面地理解我们的算法和系统设计。

接下来，我们将介绍背景、动机、我们的想法、设计，以及如何使用 PyTorch 分布式 API 实现算法和系统。

  • 论文：http://proceedings.mlr.press/v139/he21a.html
  • 源代码：https://DistML.ai.
  • 演示文稿：https://docs.google.com/presentation/d/1t6HWL33KIQo2as0nSHeBpXYtTBcy0nXCoLiKd0EashY/edit?usp=sharing

引言

Model Size
图 1：Transformer 模型的参数数量急剧增加。

大型 Transformer 模型[4][5]推动了自然语言处理和计算机视觉领域的精度突破。GPT-3[4]在几乎所有 NLP 任务上达到了新的记录高精度。视觉 Transformer（ViT）[3]在 ImageNet 上也实现了 89%的 top-1 精度，超过了最先进的卷积网络 ResNet-152 和 EfficientNet。为了应对模型尺寸的增长，研究人员提出了各种分布式训练技术，包括参数服务器[6][7][8]、流水线并行[9][10][11][12]、层内并行[13][14][15]和零冗余数据并行[16]。

然而，现有的分布式训练解决方案仅研究了在整个训练过程中所有模型权重都需要被优化的场景（即，计算和通信开销在不同迭代中相对静态）。最近关于渐进式训练的研究表明，神经网络中的参数可以动态训练：

  • 冻结训练：用于深度学习动态和可解释性的单向量正交相关分析。NeurIPS 2017
  • 通过逐步堆叠高效训练 BERT。ICML 2019
  • 使用渐进层丢弃加速基于 Transformer 的语言模型训练。NeurIPS 2020.
  • 关于 Transformer 增长对渐进 BERT 训练的影响。NACCL 2021

Freeze Training

图 2. 可解释的冻结训练：DNN 自下而上收敛（使用 ResNet 在 CIFAR10 上的结果）。每个面板展示了层与层之间的相似性，使用 SVCCA [17][18]

例如，在冻结训练[17][18]中，神经网络通常自下而上收敛（即并非所有层都需要在整个训练过程中进行训练）。图 2 展示了在这种方法中权重如何逐渐稳定。这一观察结果促使我们利用冻结训练来加速 Transformer 模型的分布式训练，通过动态分配资源以关注不断缩小的活动层集来加速训练。这种层冻结策略对于管道并行特别相关，因为排除管道中的连续底层可以减少计算、内存和通信开销。


图 3. PipeTransformer 的自动化和弹性管道过程，以加速 Transformer 模型的分布式训练

我们提出 PipeTransformer，一个弹性管道训练加速框架，该框架能够通过动态调整管道模型范围和管道副本数量来自动响应冻结层。据我们所知，这是第一篇研究在管道和数据并行训练背景下层冻结的论文。图 3 展示了这种组合的优势。首先，通过将冻结层排除在管道之外，相同的模型可以打包到更少的 GPU 中，从而减少跨 GPU 通信和更小的管道气泡。其次，将模型打包到更少的 GPU 后，相同的集群可以容纳更多的管道副本，增加数据并行的宽度。更重要的是，从这两个优势中获得的速度提升是乘法而非加法，进一步加速了训练。

管道转换器的设计面临四大挑战。首先，冻结算法必须实时和自适应地做出冻结决策；然而，现有工作[17][18]仅提供后验分析工具。其次，管道重新分区结果的效率受多个因素影响，包括分区粒度、跨分区激活大小以及微批处理中的分块（微批次的数量），这需要在庞大的解决方案空间中进行推理和搜索。第三，为了动态引入额外的管道副本，PipeTransformer 必须克服集体通信的静态性质，并在添加新进程时避免可能复杂的跨进程消息协议（一个管道由一个进程处理）。最后，缓存可以节省冻结层重复前向传播的时间，但必须共享现有管道和新添加的管道，因为系统无法为每个副本创建和预热专用缓存。

Freeze Training
图 4：展示 PipeTransformer 动态的动画

如图动画（图 4）所示，PipeTransformer 设计有四个核心构建块来解决上述挑战。首先，我们设计了一种可调自适应算法来生成信号，引导在不同迭代中选择冻结的层（冻结算法）。一旦被这些信号触发，我们的弹性管道模块（AutoPipe）就会通过考虑不同分区（冻结层和活动层）的工作负载的激活大小和方差，将剩余的活动层打包到更少的 GPU 上。然后，根据不同管道长度的先前分析结果，将一个迷你批次分割成最佳数量的微批次。我们的下一个模块 AutoDP 会启动额外的管道副本以占用释放的 GPU，并维护分层通信进程组以实现集体通信的动态成员资格。我们的最后一个模块 AutoCache，在现有和新数据并行进程之间有效地共享激活，并在转换期间自动替换过时的缓存。

总体而言，PipeTransformer 结合了冻结算法、AutoPipe、AutoDP 和 AutoCache 模块，以提供显著的训练速度提升。我们使用 Vision Transformer（ViT）在 ImageNet 上以及 BERT 在 GLUE 和 SQuAD 数据集上评估 PipeTransformer。我们的结果表明，PipeTransformer 在不损失精度的前提下，实现了高达 2.83 倍的加速。我们还提供了各种性能分析，以更全面地理解我们的算法和系统设计。最后，我们还开发了针对 PipeTransformer 的开源灵活 API，这些 API 在冻结算法、模型定义和训练加速之间提供了清晰的分离，允许将其转移到需要类似冻结策略的其他算法。

总体设计

假设我们旨在在一个分布式训练系统中训练一个大型模型，其中混合了流水线模型并行和数据并行，以针对以下场景：单个 GPU 设备的内存无法容纳模型，或者如果加载，批次大小足够小，以避免内存不足。更具体地说，我们定义以下设置：

训练任务和模型定义。我们训练 Transformer 模型（例如，视觉 Transformer、BERT 在大型图像或文本数据集上。Transformer 模型 F_具有 L 层，其中第 i 层由前向计算函数 f_i 和相应的参数集组成。

训练基础设施。假设训练基础设施包含一个具有 N 个 GPU 服务器的 GPU 集群（即节点）。每个节点有 I 个 GPU。我们的集群是同构的，意味着每个 GPU 和服务器具有相同的硬件配置。每个 GPU 的内存容量为 M_\text{GPU}。服务器通过高速网络接口（如 InfiniBand 互连）连接。

管道并行。在每台机器中，我们将模型 \mathcal{F} 装载到管道 \mathcal{P} 中，该管道有 K 个分区（K 也代表管道长度）。第 k 个分区 p_k 由连续的层组成。我们假设每个分区由单个 GPU 设备处理。1 \leq K \leq I，这意味着我们可以在一台机器中构建多个管道以适应多个模型副本。我们假设同一机器中的管道属于同一 GPU 设备。我们的管道是同步管道，不涉及过时的梯度，并且微批次的数量为 M。在 Linux 操作系统中，每个管道由单个进程处理。更多详情请参阅 GPipe [10]。

数据并行。DDP 是 R 并行工作员中的跨机分布式数据并行进程组。每个工作员是一个管道副本（单个进程）。第 r 个工作员的索引（ID）是 rank r。对于 DDP 中的任何两个管道，它们可以属于同一 GPU 服务器或不同的 GPU 服务器，并且可以使用 AllReduce 算法交换梯度。

在这些设置下，我们的目标是利用冻结训练来加速训练，这种训练不需要在整个训练过程中训练所有层。此外，它还可能有助于节省计算、通信、内存成本，并通过连续冻结层来防止过拟合。然而，只有克服设计自适应冻结算法、动态管道重新分区、高效资源重新分配和跨进程缓存这四个挑战，才能实现这些好处，正如引言中所讨论的。

Overview
图 5. PipeTransformer 训练系统概述

PipeTransformer 共同设计了一种即时冻结算法和自动弹性管道训练系统，该系统能够动态地转换管道模型的范围和管道副本的数量。整体系统架构如图 5 所示。为了支持 PipeTransformer 的弹性管道，我们维护了一个定制的 PyTorch Pipeline 版本。对于数据并行性，我们使用 PyTorch DDP 作为基线。其他库是操作系统的标准机制（例如，多进程）因此避免了专用软件或硬件定制需求。为了确保我们框架的通用性，我们将训练系统解耦为四个核心组件：冻结算法、AutoPipe、AutoDP 和 AutoCache。冻结算法（灰色）从训练循环中采样指标并做出逐层冻结决策，这些决策将与 AutoPipe（绿色）共享。AutoPipe 是一个弹性管道模块，通过排除管道中的冻结层并将活动层打包到更少的 GPU（粉色）中来加速训练，从而减少跨 GPU 通信和更小的管道气泡。 随后，AutoPipe 将管道长度信息传递给 AutoDP（紫色），然后如果可能的话，AutoDP 会启动更多的管道副本以增加数据并行宽度。插图还包括一个 AutoDP 引入新副本（紫色）的示例。AutoCache（橙色边缘）是一个跨管道缓存模块，如图中管道之间的连接所示。源代码架构与图 5 对齐，以提高可读性和通用性。

使用 PyTorch API 实现

如图 5 所示，PipeTransformers 包含四个组件：冻结算法、AutoPipe、AutoDP 和 AutoCache。其中，AutoPipe 和 AutoDP 分别依赖于 PyTorch DDP（ torch.nn.parallel.DistributedDataParallel ）[1]和 Pipeline（ torch.distributed.pipeline ）。在本博客中，我们仅突出 AutoPipe 和 AutoDP 的关键实现细节。有关冻结算法和 AutoCache 的详细信息，请参阅我们的论文。

AutoPipe：弹性管道

AutoPipe 可以通过排除冻结层从管道中，并将活动层打包到更少的 GPU 中来加速训练。本节详细介绍了 AutoPipe 的关键组件，这些组件可以动态地 1）划分管道，2）最小化管道设备数量，3）相应地优化小批量块大小。

PyTorch Pipeline 的基本用法

在深入了解 AutoPipe 的细节之前，让我们先了解一下 PyTorch Pipeline 的基本用法（ torch.distributed.pipeline.sync.Pipe ，请参阅此教程）。更具体地说，我们提供了一个简单的示例来理解 Pipeline 在实际中的设计：

# Step 1: build a model including two linear layers
fc1 = nn.Linear(16, 8).cuda(0)
fc2 = nn.Linear(8, 4).cuda(1)

# Step 2: wrap the two layers with nn.Sequential
model = nn.Sequential(fc1, fc2)

# Step 3: build Pipe (torch.distributed.pipeline.sync.Pipe)
model = Pipe(model, chunks=8)

# do training/inference
input = torch.rand(16, 16).cuda(0)
output_rref = model(input)

在这个基本示例中，我们可以看到在初始化 Pipe 之前，我们需要将模型 nn.Sequential 划分到多个 GPU 设备上，并设置最优的块数量（ chunks ）。平衡各分区间的计算时间对于管道训练速度至关重要，因为阶段间的工作负载分布不均可能导致延迟，并迫使负载较轻的设备等待。块数量也可能对管道的吞吐量产生非同寻常的影响。

平衡管道分区

在 PipeTransformer 等动态训练系统中，从参数数量角度保持最优平衡分区并不能保证最快的训练速度，因为其他因素也起着至关重要的作用：


图 6. 分区边界位于跳跃连接的中间

  1. 跨分区通信开销。将分区边界放置在跳跃连接的中间会导致额外的通信，因为跳跃连接中的张量现在必须复制到不同的 GPU。例如，在图 6 中的 BERT 分区中，分区 k 必须从分区 k-2 和分区 k-1 获取中间输出。相比之下，如果边界放置在加法层之后，分区 k-1 和 k 之间的通信开销将明显较小。我们的测量结果表明，跨设备通信的成本比略微不平衡的分区更高（参见论文附录）。因此，我们不考虑打断跳跃连接（在算法 1 的第 7 行以绿色突出显示作为一个完整的注意力层和 MLP 层）。

  2. 冻结层内存占用。在训练过程中，AutoPipe 必须多次重新计算分区边界，以平衡两种不同类型的层：冻结层和活动层。由于冻结层不需要反向激活图、优化器状态和梯度，因此其内存成本仅为非活动层的一小部分。为了获取内存和计算成本的详细指标，我们不需要启动侵入式分析器，而是定义一个可调的成本因子λ_{\text{frozen}}来估计冻结层相对于相同活动层的内存占用比。根据我们在实验硬件上的经验测量，我们将其设置为\frac{1}{6}。


基于上述两个考虑因素，AutoPipe 根据参数大小平衡管道分区。更具体地说，AutoPipe 使用贪婪算法将所有冻结层和活动层分配到 K 个 GPU 设备，以均匀分配分区子层。伪代码在算法 1 中的 load\_balance() 函数中描述。冻结层从原始模型中提取出来，并保存在管道的第一个设备中的单独模型实例\mathcal{F}_{\text{frozen}}中。

注意，本文中采用的分区算法并非唯一选择；PipeTransformer 被模块化设计，可以与任何替代方案协同工作。

管道压缩

管道压缩有助于释放 GPU 资源，以便容纳更多管道副本，并减少分区之间的跨设备通信次数。为了确定压缩时机，我们可以估算压缩后最大分区的内存成本，并将其与 timestep T=0 时管道的最大分区内存成本进行比较。为了避免进行广泛的内存分析，压缩算法使用参数大小作为训练内存占用大小的代理。基于这种简化，管道压缩的标准如下：


接收到冻结通知后，AutoPipe 将始终尝试将管道长度 K 除以 2（例如，从 8 变为 4，然后变为 2）。使用 \frac{K}{2} 作为输入，压缩算法可以验证结果是否满足方程（1）中的标准。伪代码在算法 1 的第 25-33 行中给出。请注意，这种压缩在训练期间会指数级增加加速比，这意味着如果 GPU 服务器有更多的 GPU（例如，超过 8 个），加速比将进一步放大。


图 7. 管道气泡：F_{d,b}、U_d" 分别表示设备 d 上微批 b 的正向、反向和优化器更新。每个迭代中每个微批正向和反向的总气泡大小为 K-1 倍。

此外，这种技术还可以通过缩小管道气泡的大小来加速训练。为了解释管道中的气泡大小，图 7 描述了 4 个微批如何在 K = 4 的 4 设备管道中运行。一般来说，每个微批正向和反向的总气泡大小为 (K-1) 倍。因此，很明显，较短的管道具有更小的气泡大小。

动态微批数量

之前的流水线并行系统使用每个小批量（M）固定的微批量数。GPipe 建议 M≥4×K，其中 K 是分区数（流水线长度）。然而，鉴于 PipeTransformer 动态配置 K，我们发现保持静态 M 在训练过程中是不理想的。此外，当与 DDP 集成时，M 的值也会影响 DDP 梯度同步的效率。由于 DDP 必须在启动梯度同步之前等待最后一个微批量完成对参数的反向计算，因此更细的微批量会导致计算和通信之间的重叠更小。因此，PipeTransformer 不是使用静态值，而是在 DDP 环境的混合中动态搜索最佳 M 值，范围从 K 到 6K。对于特定的训练环境，只需要进行一次配置（见算法 1 的第 35 行）。

完整的源代码请参阅 https://github.com/Distributed-AI/PipeTransformer/blob/master/pipe_transformer/pipe/auto_pipe.py 。

AutoDP：启动更多流水线副本

由于 AutoPipe 将相同的流水线压缩到更少的 GPU 上，AutoDP 可以自动启动新的流水线副本以增加数据并行宽度。

尽管概念上简单，但通信和状态上的微妙依赖关系需要精心设计。挑战有三：

  1. DDP 通信：PyTorch DDP 中的集体通信需要静态成员资格，这阻止了新管道与现有管道连接；

  2. 状态同步：新激活的进程必须与现有管道的训练进度（例如，epoch 数量和学习率）、权重和优化器状态、冻结层的边界以及管道 GPU 范围保持一致；

  3. 数据集重新分配：数据集应重新平衡以匹配动态的管道数量。这不仅避免了落后者，还确保了所有 DDP 进程的梯度权重相等。


图 8. AutoDP：通过双进程组之间的消息处理动态数据并行（0-7 号进程属于机器 0，而 8-15 号进程属于机器 1）

为了应对这些挑战，我们为 DDP 创建了双通信进程组。如图 8 所示，消息进程组（紫色）负责轻量级控制消息，覆盖所有进程，而活跃训练进程组（黄色）仅包含活跃进程，在训练期间作为重型张量通信的载体。消息组保持静态，而训练组则根据活跃进程进行拆解和重建。在 T0 阶段，只有 0 号和 8 号进程活跃。在过渡到 T1 阶段时，0 号进程激活 1 号和 9 号进程（新添加的管道副本）并使用消息组同步上述必要信息。这四个活跃进程随后形成一个新训练组，允许静态集体通信适应动态成员。为了重新分配数据集，我们实现了一种 DistributedSampler 的变体，可以无缝调整数据样本以匹配活跃管道副本的数量。

上述设计也有助于自然地减少 DDP 通信开销。更具体地说，当从 T0 过渡到 T1 时，进程 0 和 1 销毁现有的 DDP 实例，而活跃进程使用缓存的流水线模型（AutoPipe 分别存储冻结模型和缓存模型）构建一个新的 DDP 训练组。

我们使用以下 API 来实现上述设计。

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# initialize the process group (this must be called in the initialization of PyTorch DDP)
dist.init_process_group(init_method='tcp://' + str(self.config.master_addr) + ':' +
str(self.config.master_port), backend=Backend.GLOO, rank=self.global_rank, world_size=self.world_size)
...

# create active process group (yellow color)
self.active_process_group = dist.new_group(ranks=self.active_ranks, backend=Backend.NCCL, timeout=timedelta(days=365))
...

# create message process group (yellow color)
self.comm_broadcast_group = dist.new_group(ranks=[i for i in range(self.world_size)], backend=Backend.GLOO, timeout=timedelta(days=365))
...

# create DDP-enabled model when the number of data-parallel workers is changed. Note:
# 1. The process group to be used for distributed data all-reduction.
If None, the default process group, which is created by torch.distributed.init_process_group, will be used.
In our case, we set it as self.active_process_group
# 2. device_ids should be set when the pipeline length = 1 (the model resides on a single CUDA device).

self.pipe_len = gpu_num_per_process
if gpu_num_per_process > 1:
    model = DDP(model, process_group=self.active_process_group, find_unused_parameters=True)
else:
    model = DDP(model, device_ids=[self.local_rank], process_group=self.active_process_group, find_unused_parameters=True)

# to broadcast message among processes, we use dist.broadcast_object_list
def dist_broadcast(object_list, src, group):
    """Broadcasts a given object to all parties."""
    dist.broadcast_object_list(object_list, src, group=group)
    return object_list

完整的源代码请参考 https://github.com/Distributed-AI/PipeTransformer/blob/master/pipe_transformer/dp/auto_dp.py 。

实验

本节首先总结实验设置，然后使用计算机视觉和自然语言处理任务评估 PipeTransformer。

硬件。实验在通过 InfiniBand CX353A（5GB/s）连接的 2 台相同机器上进行，每台机器配备 8 块 NVIDIA Quadro RTX 5000（16GB GPU 内存）。机器内部的 GPU 到 GPU 带宽（PCI 3.0，16 条通道）为 15.754GB/s。

实现。我们使用 PyTorch Pipe 作为构建块。BERT 模型定义、配置和相关分词器来自 HuggingFace 3.5.0。我们通过遵循其 TensorFlow 实现来使用 PyTorch 实现 Vision Transformer。更多详细信息可以在我们的源代码中找到。

模型和数据集。实验使用了 CV 和 NLP 中的两个代表性 Transformer：视觉 Transformer（ViT）和 BERT。ViT 在图像分类任务上运行，使用 ImageNet21K 上的预训练权重初始化，并在 ImageNet 和 CIFAR-100 上进行微调。BERT 在两个任务上运行，在 GLUE 基准中的 SST-2 数据集上进行文本分类，以及在 SQuAD v1.1 数据集（斯坦福问答）上进行问答，这是一个包含 10 万个众包问答对的集合。

训练方案。鉴于大型模型通常需要数千个 GPU 天（例如，GPT-3）从头开始训练，因此使用预训练模型进行下游任务的微调已成为计算机视觉和自然语言处理领域的趋势。此外，PipeTransformer 是一个复杂的训练系统，涉及多个核心组件。因此，对于 PipeTransformer 系统开发的第一版和算法研究，从头开始使用大规模预训练进行开发和评估并不经济。因此，本节中展示的实验专注于预训练模型。请注意，由于预训练和微调中的模型架构相同，PipeTransformer 可以同时服务于两者。我们已在附录中讨论了预训练结果。

基准。本节中的实验将 PipeTransformer 与最先进的框架进行了比较，该框架是 PyTorch Pipeline（PyTorch 对 GPipe 的实现）和 PyTorch DDP 的混合方案。由于这是第一篇研究通过冻结层加速分布式训练的论文，因此尚无完全对齐的对应解决方案。

超参数。实验使用 ViT-B/16（12 个转换器层，16x16 输入块大小）进行 ImageNet 和 CIFAR-100，BERT-large-uncased（24 层）进行 SQuAD 1.1，BERT-base-uncased（12 层）进行 SST-2。使用 PipeTransformer，ViT 和 BERT 训练可以将每个管道的批处理大小设置为约 400 和 64。所有实验的其他超参数（例如，epoch、学习率）见附录。

总体训练加速


我们在上表中总结了总体实验结果。请注意，我们报告的加速是基于保守的 α 1/3 值，可以获得相当甚至更高的精度。更激进的 α（2/5，1/2）可以获得更高的加速，但可能会导致精度略有损失。请注意，BERT（24 层）的模型大小大于 ViT-B/16（12 层），因此通信需要更多时间。

性能分析

加速分解

本节展示了评估结果，并分析了\autopipe 中不同组件的性能。更多实验结果可在附录中找到。


图 9.速度提升分解（ViT 在 ImageNet 上的应用）

为了了解所有四个组件的有效性和它们对训练速度的影响，我们进行了不同的组合实验，并使用它们的训练样本吞吐量（样本/秒）和加速比作为指标。结果如图 9 所示。这些实验结果的关键要点是：

  1. 主要的加速是弹性管道的结果，这是通过联合使用 AutoPipe 和 AutoDP 实现的。
  2. AutoCache 的贡献通过 AutoDP 得到放大；
  3. 即使没有系统性的调整，冻结训练本身也会降低训练速度。

调整冻结算法中的α值


图 10. 调整冻结算法中的α值

我们进行了实验，以展示冻结算法中的 \alpha 对训练速度的影响。结果表明，较大的 \alpha（过度冻结）会导致更大的加速，但会略微降低性能。在图 10 所示的情况下，当 \alpha=1/5 时，冻结训练优于正常训练，实现了 2.04 倍的加速。我们还在附录中提供了更多结果。

弹性管道中的最佳块大小


图 11. 弹性管道中的最佳块数量

我们对不同管道长度 K 下的最佳微批大小 M 进行了分析。结果总结在图 11 中。如图所示，不同的 K 值会导致不同的最佳 M 值，不同 M 值之间的吞吐量差距很大（如图 K=8 时所示），这证实了在弹性管道中需要一个前置分析器的必要性。

理解缓存时机


图 12. 缓存时机

为了评估 AutoCache，我们比较了从第 0 个 epoch 激活 AutoCache 的训练样本吞吐量（蓝色）与没有 AutoCache 的训练作业（红色）。图 12 显示，过早启用缓存可能会减慢训练速度，因为对于少量冻结层，缓存可能比前向传播更昂贵。随着更多层的冻结，缓存激活明显优于相应的正向传播。因此，AutoCache 使用分析器来确定启用缓存的正确时机。在我们的系统中，对于 ViT（12 层），缓存从第 3 个冻结层开始，而对于 BERT（24 层），缓存从第 5 个冻结层开始。

如需更详细的实验分析，请参阅我们的论文。

摘要

本博客介绍了 PipeTransformer，这是一个结合弹性管道并行和数据并行的整体解决方案，用于使用 PyTorch 分布式 API 进行分布式训练。更具体地说，PipeTransformer 会逐步冻结管道中的层，将剩余的活跃层打包到更少的 GPU 中，并通过更多管道副本来增加数据并行宽度。在 ViT 和 Bert 模型上的评估表明，与最先进的基线相比，PipeTransformer 在不损失精度的前提下，速度提升了高达 2.83 倍。

参考文献

[1] 李，S.，赵，Y.，瓦拉，R.，萨尔帕卡尔，O.，诺德豪斯，P.，李，T.，帕斯克，A.，史密斯，J.，沃恩，B.，达米亚尼，P.等. Pytorch 分布式：加速数据并行训练的经验。VLDB Endowment，13(12)，2020

[2] 德夫林，J.，张，M. W.，李，K.，和图托瓦诺娃，K. BERT：用于语言理解的深度双向变换器预训练。NAACL-HLT，2019

[3] Dosovitskiy, A.，Beyer, L.，Kolesnikov, A.，Weissenborn, D.，Zhai, X.，Unterthiner, T.，Dehghani, M.，Minderer, M.，Heigold, G.，Gelly, S.，等. 一张图片胜过 16x16 个单词：大规模图像识别的 Transformer。

[4] Brown, T. B.，Mann, B.，Ryder, N.，Subbiah, M.，Kaplan, J.，Dhariwal, P.，Neelakantan, A.，Shyam, P.，Sastry, G.，Askell, A.，等. 语言模型是少样本学习者。

[5] Lepikhin, D.，Lee, H.，Xu, Y.，Chen, D.，Firat, O.，Huang, Y.，Krikun, M.，Shazeer, N.，Chen, Z. Gshard：通过条件计算和自动分片扩展巨型模型。

[6] Li, M.，Andersen, D. G.，Park, J. W.，Smola, A. J.，Ahmed, A.，Josifovski, V.，Long, J.，Shekita, E. J.，Su, B. Y. 使用参数服务器扩展分布式机器学习。在 11 届{USENIX}操作系统设计与实现研讨会({OSDI} 14)，第 583–598 页，2014 年。

[7] 姜毅，朱宇，蓝晨，易波，崔宇，郭晨。针对异构 GPU/CPU 集群的分布式深度神经网络训练加速统一架构。在第十四届 USENIX 操作系统设计与实现研讨会（OSDI 20），第 463-479 页。USENIX 协会，2020 年 11 月。ISBN 978-1-939133-19-9。

[8] 金善，余国义，朴亨，赵锡，郑锡，韩浩，李硕，郑钟秀，崔钟国。Parallax：深度神经网络稀疏感知数据并行训练。在第十四届 EuroSys 会议论文集，第 1-15 页，2019 年。

[9] 金成，李河，郑明，白完，尹钟，金伊，林硕，金善。TorchGPipe：训练巨型模型的即时管道并行。

[10] 黄宇，程宇，阿帕纳，奥菲尔，陈明轩，陈德，李浩，吴恩达，等。Gpipe：使用管道并行高效训练巨型神经网络。

[11] 帕克，J. H.，云，G.，尹，C. M.，阮，N. T.，李，S.，崔，J.，吴，S. H.，以及崔，Y. Hetpipe：通过集成流水线模型并行和数据并行，在（瘦弱）异构 GPU 集群上实现大型深度神经网络训练。在 2020 年 USENIX 年度技术会议（USENIX ATC 20），第 307-321 页。USENIX 协会，2020 年 7 月。ISBN 978-1-939133-14-4。

[12] 纳拉亚南，D.，哈拉普，A.，潘尼沙耶，A.，塞沙德里，V.，德瓦努尔，N. R.，甘格尔，G. R.，吉本斯，P. B.，以及扎哈里亚，M. Pipedream：用于深度神经网络训练的通用管道并行。在 2019 年第 27 届 ACM 操作系统原理研讨会（SOSP '19）论文集中，第 1-15 页，纽约，纽约，美国。ACM 协会。ISBN 9781450368735。doi: 10.1145/3341301.3359646。

[13] 列皮欣，D.，李，H.，徐，Y.，陈，D.，菲拉特，O.，黄，Y.，克里昆，M.，沙泽尔，N.，以及陈，Z. Gshard：通过条件计算和自动分片扩展巨型模型。

[14] Shazeer, N.，Cheng, Y.，Parmar, N.，Tran, D.，Vaswani, A.，Koanantakool, P.，Hawkins, P.，Lee, H.，Hong, M.，Young, C.，Sepassi, R.，和 Hechtman, B. Mesh-Tensorflow：超级计算机的深度学习。在 Bengio, S.，Wallach, H.，Larochelle, H.，Grauman, K.，Cesa-Bianchi, N.，和 Garnett, R. (编者)，神经信息处理系统进展，第 31 卷，第 10414–10423 页。Curran Associates, Inc.，2018 年。

[15] Shoeybi, M.，Patwary, M.，Puri, R.，LeGresley, P.，Casper, J.，和 Catanzaro, B. Megatron-LM：使用模型并行训练数十亿参数语言模型。

[16] Rajbhandari, S.，Rasley, J.，Ruwase, O.，和 He, Y. ZERO：训练万亿参数模型时的内存优化。

[17] Raghu, M.，Gilmer, J.，Yosinski, J.，和 Sohl Dickstein, J. Svcca：用于深度学习动态和可解释性的奇异向量典型相关分析。在 NIPS，2017 年。

[18] Morcos, A.，Raghu, M.，和 Bengio, S. 在神经网络中关于表示相似性的见解。在 Bengio, S.，Wallach, H.，Larochelle, H.，Grauman, K.，Cesa-Bianchi, N.，和 Garnett, R.（编者），《神经信息处理系统 31 卷进展》，第 5732-5741 页。Curran Associates, Inc.，2018 年。