在这篇博文中，我们将介绍首篇经同行评审的研究论文，该论文探讨了如何加速 PyTorch DDP (torch.nn.parallel.DistributedDataParallel) [1] 与流水线并行 (torch.distributed.pipeline) 的混合训练模式——PipeTransformer：大规模模型（如 BERT [2] 和 ViT [3]）分布式训练的自动化弹性流水线，该论文已发表于 ICML 2021。

PipeTransformer 利用自动化弹性流水线技术实现 Transformer 模型的高效分布式训练。在 PipeTransformer 中，我们设计了一种自适应即时冻结算法，能够在训练过程中逐步识别并冻结部分层，同时设计了一个弹性流水线系统，可以动态分配资源以训练剩余的活跃层。更具体地说，PipeTransformer 会自动将冻结层从流水线中剔除，将活跃层打包到更少的 GPU 上，并创建更多的副本以增加数据并行宽度。我们使用 ImageNet 上的视觉 Transformer (ViT) 以及 SQuAD 和 GLUE 数据集上的 BERT 对 PipeTransformer 进行了评估。结果表明，与当前最先进的基准方法相比，PipeTransformer 在不损失精度的情况下实现了高达 2.83 倍的加速。我们还提供了多项性能分析，以更全面地理解我们的算法和系统设计。

接下来，我们将介绍其背景、动机、核心理念、设计方案，以及如何利用 PyTorch 分布式 API 实现该算法和系统。

论文： https://pmlr.com.cn/v139/he21a.html
源代码： https://DistML.ai
幻灯片： https://docs.google.com/presentation/d/1t6HWL33KIQo2as0nSHeBpXYtTBcy0nXCoLiKd0EashY/edit?usp=sharing

引言

图 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]。

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

冻结训练 (Freeze Training)：深度学习动力学与可解释性的奇异向量典型相关分析。NeurIPS 2017
通过渐进式堆叠高效训练 BERT。ICML 2019
利用渐进式层丢弃加速基于 Transformer 的语言模型训练。NeurIPS 2020。
关于渐进式 BERT 训练的 Transformer 增长研究。NAACL 2021

图 2. 可解释的冻结训练：DNN 从底层向上收敛（基于 ResNet 在 CIFAR10 上的结果）。每个方格显示了使用 SVCCA [17][18] 进行的逐层相似度分析

例如，在冻结训练 [17][18] 中，神经网络通常从底层向上收敛（即并非所有层都需要在整个训练过程中持续训练）。图 2 展示了这种方法中权重在训练过程中如何逐渐趋于稳定的示例。这一观察启发我们将冻结训练应用于 Transformer 模型的分布式训练，通过动态分配资源专注于不断缩减的活跃层集合来加速训练。这种层冻结策略对于流水线并行尤为重要，因为从流水线中剔除连续的底层可以减少计算、内存和通信开销。

图 3. PipeTransformer 的自动化和弹性流水线加速 Transformer 模型分布式训练的过程

我们提出了 PipeTransformer，这是一个弹性流水线训练加速框架，能够通过动态转换流水线模型的作用域和流水线副本数量来自动响应冻结层。据我们所知，这是首篇在流水线并行和数据并行训练背景下研究层冻结的论文。图 3 展示了这种组合的优势。首先，通过从流水线中剔除冻结层，相同的模型可以被打包到更少的 GPU 上，从而减少跨 GPU 通信并缩减流水线气泡。其次，将模型打包到更少的 GPU 上后，同一集群可以容纳更多的流水线副本，从而增加了数据并行宽度。更重要的是，这两种优势带来的加速是乘法而非加法关系，进一步加速了训练。

PipeTransformer 的设计面临四大挑战。首先，冻结算法必须做出即时的自适应冻结决策；然而，现有工作 [17][18] 仅提供事后分析工具。其次，流水线重新分区的效率受多种因素影响，包括分区粒度、跨分区激活大小以及 mini-batch 中的分块（micro-batch 的数量），这些需要在大解空间中进行推理和搜索。第三，为了动态引入额外的流水线副本，PipeTransformer 必须克服集合通信的静态特性，并避免在加入新进程时出现复杂的跨进程消息传递协议（一个流水线由一个进程处理）。最后，缓存虽然可以节省冻结层重复前向传播的时间，但它必须在现有流水线和新增流水线之间共享，因为系统无法负担为每个副本创建和预热专用缓存的成本。

图 4：展示 PipeTransformer 动态过程的动画

正如动画（图 4）所示，PipeTransformer 设计了四个核心构建模块来应对上述挑战。首先，我们设计了一种可调的自适应算法（冻结算法），用于生成指导在不同迭代中选择冻结层的信号。一旦被这些信号触发，我们的弹性流水线模块（AutoPipe）就会通过考虑异构分区（冻结层和活跃层）的激活大小和工作负载方差，将剩余的活跃层打包到更少的 GPU 上。然后，它根据之前针对不同流水线长度的分析结果，将 mini-batch 拆分为最佳数量的 micro-batch。我们的下一个模块 AutoDP 会生成额外的流水线副本以占用空出的 GPU，并维护分层通信进程组，以实现集合通信的动态成员资格。我们的最后一个模块 AutoCache 可以高效地在现有和新的数据并行进程之间共享激活，并在转换期间自动替换陈旧的缓存。

总而言之，PipeTransformer 结合了冻结算法、AutoPipe、AutoDP 和 AutoCache 模块，提供了显著的训练加速。我们使用 ImageNet 上的视觉 Transformer (ViT) 和 GLUE 及 SQuAD 数据集上的 BERT 对 PipeTransformer 进行了评估。结果表明，PipeTransformer 在不损失精度的情况下实现了高达 2.83 倍的加速。我们还提供了多项性能分析，以更全面地理解我们的算法和系统设计。最后，我们还为 PipeTransformer 开发了开源的灵活 API，实现了冻结算法、模型定义和训练加速之间的清晰分离，从而支持迁移到其他需要类似冻结策略的算法。

总体设计

假设我们的目标是在一个使用流水线模型并行和数据并行混合模式的分布式训练系统中训练大规模模型，以应对单个 GPU 设备内存无法容纳模型，或者即使装载后 batch size 过小导致内存不足的场景。更具体地，我们定义设置如下：

训练任务与模型定义：我们训练 Transformer 模型（例如视觉 Transformer、BERT），处理大规模图像或文本数据集。Transformer 模型 \mathcal{F} 有 L 层，其中第 i 层由前向计算函数 f_i 和一组对应的参数组成。

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

流水线并行：在每台机器中，我们将模型 \mathcal{F} 加载到一个拥有 K 个分区的流水线 \mathcal{P} 中（K 也代表流水线长度）。第 k 个分区 p_k 由连续的层组成。我们假设每个分区由单个 GPU 设备处理。1 \leq K \leq I，这意味着我们可以在单台机器内为多个模型副本构建多个流水线。我们假设流水线中的所有 GPU 设备属于同一台机器。我们的流水线是同步流水线，不涉及陈旧梯度，micro-batch 的数量为 M。在 Linux 操作系统中，每个流水线由单个进程处理。详情请参阅 GPipe [10]。

数据并行：DDP 是一个跨机器的分布式数据并行进程组，包含 R 个并行工作节点。每个工作节点是一个流水线副本（单个进程）。第 r 个工作节点的索引（ID）为 rank r。对于 DDP 中的任意两个流水线，它们可以属于同一 GPU 服务器或不同 GPU 服务器，并且可以通过 AllReduce 算法交换梯度。

在这些设置下，我们的目标是通过利用冻结训练来加速训练，这种方法不需要所有层在整个训练过程中都进行训练。此外，它可以通过连续冻结层来节省计算、通信、内存成本，并可能防止过拟合。然而，这些优势只有在克服了介绍中讨论的自适应冻结算法、动态流水线重分区、高效资源重分配和跨进程缓存这四个设计挑战后才能实现。

图 5. PipeTransformer 训练系统概览

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

使用 PyTorch API 实现

从图 5 可以看出，PipeTransformer 包含四个组件：冻结算法、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) 相应地优化 mini-batch 分块大小。

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. 分区边界位于跳跃连接（skip connection）中间

跨分区通信开销：将分区边界放置在跳跃连接中间会导致额外的通信，因为跳跃连接中的张量现在必须复制到不同的 GPU。例如，图 6 中的 BERT 分区，分区 k 必须接收来自分区 k-2 和分区 k-1 的中间输出。相反，如果边界放置在加法层之后，分区 k-1 和 k 之间的通信开销明显更小。我们的测量表明，跨设备通信比轻微不平衡的分区更昂贵（参见我们论文的附录）。因此，我们不考虑断开跳跃连接（在算法 1 的第 7 行中以绿色突出显示为一个完整的注意力和 MLP 层）。
冻结层的内存占用：在训练过程中，AutoPipe 必须多次重新计算分区边界，以平衡两种不同类型的层：冻结层和活跃层。冻结层的内存成本仅为非活动层的一小部分，因为冻结层不需要反向激活图、优化器状态和梯度。我们没有启动侵入式分析器来获取内存和计算成本的详尽指标，而是定义了一个可调成本因子 \lambda_{\text{frozen}} 来估计冻结层相对于同等活跃层的内存占用比例。根据我们实验硬件的经验测量，我们将其设置为 \frac{1}{6}。

基于上述两点考虑，AutoPipe 根据参数大小平衡流水线分区。具体而言，AutoPipe 使用贪心算法来分配所有冻结层和活跃层，以便将分区后的子层均匀分布到 K 个 GPU 设备中。伪代码描述为算法 1 中的 load\_balance() 函数。冻结层从原始模型中提取，并保存在流水线第一个设备中的独立模型实例 \mathcal{F}_{\text{frozen}} 中。

注意，本文采用的分区算法并非唯一选择；PipeTransformer 采用了模块化设计，可以与任何替代方案配合使用。

流水线压缩

流水线压缩有助于释放 GPU 以容纳更多的流水线副本，并减少分区之间跨设备的通信次数。为了确定压缩时机，我们可以估计压缩后最大分区的内存成本，然后将其与时间步 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 个微批次如何在 4 设备流水线 K = 4 中运行。总的来说，总气泡大小是 (K-1) 倍的微批次前向和反向成本。因此，很明显较短的流水线具有较小的气泡尺寸。

微批次的动态数量

先前的流水线并行系统对每个 mini-batch 使用固定数量的微批次 (M)。GPipe 建议 M \geq 4 \times K，其中 K 是分区数（流水线长度）。然而，考虑到 PipeTransformer 动态配置 K，我们发现训练过程中保持静态的 M 是次优的。此外，当与 DDP 集成时，M 的值也会对 DDP 梯度同步的效率产生影响。由于 DDP 必须等待最后一个微批次完成参数上的反向计算才能启动梯度同步，更细粒度的微批次会导致计算和通信之间的重叠更少。因此，PipeTransformer 不使用静态值，而是在混合 DDP 环境中通过枚举从 K 到 6K 的 M 值来实时搜索最佳 M。对于特定的训练环境，分析只需执行一次（参见算法 1 第 35 行）。

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

AutoDP：生成更多的流水线副本

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

尽管概念上很简单，但对通信和状态的微妙依赖需要仔细设计。挑战主要体现在三个方面：

DDP 通信：PyTorch DDP 中的集合通信需要静态成员资格，这阻碍了新流水线与现有流水线的连接；
状态同步：新激活的进程必须与现有流水线在训练进度（例如 epoch 数和学习率）、权重和优化器状态、冻结层边界以及流水线 GPU 范围上保持一致；
数据集重分配：应重新平衡数据集以匹配动态数量的流水线。这不仅避免了落后者，还确保来自所有 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 实例，活跃进程使用缓存的流水线模型构建新的 DDP 训练组（AutoPipe 分别存储冻结模型和缓存模型）。

我们使用以下 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 实现了视觉 Transformer。更多细节可以在我们的源代码中找到。

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

训练方案：鉴于大型模型如果从头开始训练通常需要数千个 GPU 天（例如 GPT-3），使用预训练模型微调下游任务已成为 CV 和 NLP 社区的趋势。此外，PipeTransformer 是一个涉及多个核心组件的复杂训练系统。因此，对于 PipeTransformer 系统开发和算法研究的第一个版本，从头开始进行大规模预训练开发和评估并不划算。因此，本节展示的实验集中在预训练模型上。注意，由于预训练和微调的模型架构相同，PipeTransformer 可以兼顾两者。我们在附录中讨论了预训练结果。

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

超参数：实验使用 ViT-B/16（12 层 transformer，16x16 输入块大小）用于 ImageNet 和 CIFAR-100，BERT-large-uncased（24 层）用于 SQuAD 1.1，以及 BERT-base-uncased（12 层）用于 SST-2。使用 PipeTransformer，ViT 和 BERT 训练可以将每个流水线的批处理大小分别设置为约 400 和 64。所有实验的其他超参数（如 epoch、学习率）均列于附录中。

总体训练加速

我们在上表中总结了总体实验结果。请注意，我们报告的加速比基于保守的 \alpha 1/3 值，该值可以获得相当甚至更高的精度。更激进的 \alpha (\frac{2}{5}, \frac{1}{2}) 可以获得更高的加速比，但可能会导致精度的轻微损失。请注意，BERT（24 层）的模型大小比 ViT-B/16（12 层）大，因此通信耗时更多。

性能分析

加速分解

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

图 9. 加速分解（ImageNet 上的 ViT）

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

主要的加速源于弹性流水线，这是通过 AutoPipe 和 AutoDP 的联合使用实现的；
AutoCache 的贡献随着 AutoDP 的使用而放大；
单纯的冻结训练而不进行系统层面的调整甚至会降低训练速度。

调整冻结算法中的 \alpha

图 10. 调整冻结算法中的 \alpha

我们进行了实验来展示冻结算法中的 \alpha 如何影响训练速度。结果清楚地表明，较大的 \alpha（过度冻结）会导致更大的加速，但性能会有轻微下降。在图 10 所示的情况下，当 \alpha=1/5 时，冻结训练优于普通训练并获得了 2.04 倍的加速。我们在附录中提供了更多结果。

弹性流水线中的最佳分块

图 11. 弹性流水线中的最佳分块数量

我们针对不同的流水线长度 K 分析了最佳微批次数量 M。结果汇总在图 11 中。正如我们所见，不同的 K 值会导致不同的最佳 M，且不同 M 值之间的吞吐量差距很大（如 K=8 时所示），这证实了在弹性流水线中进行前序分析的必要性。

了解缓存的时机

图 12. 缓存的时机

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

有关更详细的实验分析，请参考我们的论文。

总结

本博文介绍了 PipeTransformer，这是一个结合了弹性流水线并行和数据并行的整体解决方案，使用 PyTorch 分布式 API 进行分布式训练。更具体地，PipeTransformer 在流水线中递增地冻结层，将剩余的活跃层打包到更少的 GPU 上，并生成更多的流水线副本以增加数据并行宽度。在 ViT 和 BERT 模型上的评估表明，与当前最先进的基准相比，PipeTransformer 在不损失精度的情况下实现了高达 2.83 倍的加速。

参考

[1] Li, S., Zhao, Y., Varma, R., Salpekar, O., Noordhuis, P., Li,T., Paszke, A., Smith, J., Vaughan, B., Damania, P., et al. Pytorch Distributed: Experiences on Accelerating Dataparallel Training. Proceedings of the VLDB Endowment, 13(12), 2020

[2] Devlin, J., Chang, M. W., Lee, K., and Toutanova, K. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In NAACL-HLT, 2019

[3] Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. An image is Worth 16×16 words: Transformers for Image Recognition at Scale.

[4] Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language Models are Few-shot Learners.

[5] Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z. Gshard: Scaling Giant Models with Conditional Computation and Automatic Sharding.

[6] Li, M., Andersen, D. G., Park, J. W., Smola, A. J., Ahmed, A., Josifovski, V., Long, J., Shekita, E. J., and Su, B. Y. Scaling Distributed Machine Learning with the Parameter Server. In 11th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 14), pp. 583–598, 2014.

[7] Jiang, Y., Zhu, Y., Lan, C., Yi, B., Cui, Y., and Guo, C. A Unified Architecture for Accelerating Distributed DNN Training in Heterogeneous GPU/CPU Clusters. In 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20), pp. 463–479. USENIX Association, November 2020. ISBN 978-1-939133-19- 9.

[8] Kim, S., Yu, G. I., Park, H., Cho, S., Jeong, E., Ha, H., Lee, S., Jeong, J. S., and Chun, B. G. Parallax: Sparsity-aware Data Parallel Training of Deep Neural Networks. In Proceedings of the Fourteenth EuroSys Conference 2019, pp. 1–15, 2019.

[9] Kim, C., Lee, H., Jeong, M., Baek, W., Yoon, B., Kim, I., Lim, S., and Kim, S. TorchGPipe: On-the-fly Pipeline Parallelism for Training Giant Models.

[10] Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, M. X., Chen, D., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., et al. Gpipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.

[11] Park, J. H., Yun, G., Yi, C. M., Nguyen, N. T., Lee, S., Choi, J., Noh, S. H., and ri Choi, Y. Hetpipe: Enabling Large DNN Training on (whimpy) Heterogeneous GPU Clusters through Integration of Pipelined Model Parallelism and Data Parallelism. In 2020 USENIX Annual Technical Conference (USENIX ATC 20), pp. 307–321. USENIX Association, July 2020. ISBN 978-1-939133- 14-4.

[12] Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N. R., Ganger, G. R., Gibbons, P. B., and Zaharia, M. Pipedream: Generalized Pipeline Parallelism for DNN Training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles, SOSP ’19, pp. 1–15, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450368735. doi: 10.1145/3341301.3359646.

[13] Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z. Gshard: Scaling Giant Models with Conditional Computation and Automatic Sharding.

[14] Shazeer, N., Cheng, Y., Parmar, N., Tran, D., Vaswani, A., Koanantakool, P., Hawkins, P., Lee, H., Hong, M., Young, C., Sepassi, R., and Hechtman, B. Mesh-Tensorflow: Deep Learning for Supercomputers. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 31, pp. 10414–10423. Curran Associates, Inc., 2018.

[15] Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-LM: Training Multi-billion Parameter Language Models using Model Parallelism.

[16] Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. ZERO: Memory Optimization towards Training a Trillion Parameter Models.

[17] Raghu, M., Gilmer, J., Yosinski, J., and Sohl Dickstein, J. Svcca: Singular Vector Canonical Correlation Analysis for Deep Learning Dynamics and Interpretability. In NIPS, 2017.

[18] Morcos, A., Raghu, M., and Bengio, S. Insights on Representational Similarity in Neural Networks with Canonical Correlation. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 31, pp. 5732–5741. Curran Associates, Inc., 2018.

PipeTransformer：用于大规模模型分布式训练的自动化弹性流水线

引言