在这篇博客文章中，我们描述了第一篇经过同行评审的研究论文，该论文探讨了加速 PyTorch DDP（torch.nn.parallel.DistributedDataParallel）[1] 和流水线（torch.distributed.pipeline）混合训练的方法——PipeTransformer：大型模型（如 BERT [2] 和 ViT [3] 等 Transformer）分布式训练的自动化弹性流水线，该论文发表于 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

引言

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

大型 Transformer 模型[4][5] 在自然语言处理和计算机视觉领域都推动了准确性突破。GPT-3[4] 在几乎所有 NLP 任务上都达到了新的最高准确率。Vision 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。
关于渐进式 BERT 训练的 Transformer 增长。NACCL 2021

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

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

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

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

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

图 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，它在冻结算法、模型定义和训练加速之间提供了清晰的分离，从而可以将其转移到需要类似冻结策略的其他算法。

总体设计

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

训练任务和模型定义。我们训练 Transformer 模型（例如 Vision 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 设备都属于同一台机器。我们的流水线是同步流水线，不涉及陈旧的梯度，微批次数量为 M。在 Linux 操作系统中，每个流水线由一个进程处理。更多细节请参考 GPipe [10]。

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

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

图 5. PipeTransformer 训练系统概述

PipeTransformer 共同设计了实时冻结算法和自动化弹性流水线训练系统，能够动态转换流水线模型的范围和流水线副本的数量。整个系统架构如图 5 所示。为了支持 PipeTransformer 的弹性流水线，我们维护了一个自定义版本的 PyTorch 流水线。对于数据并行，我们使用 PyTorch DDP 作为基准。其他库是操作系统的标准机制（例如多进程），因此避免了专门的软件或硬件定制要求。为了确保我们框架的通用性，我们将训练系统解耦为四个核心组件：冻结算法、AutoPipe、AutoDP 和 AutoCache。冻结算法（灰色）从训练循环中采样指标并做出逐层冻结决策，这些决策将与 AutoPipe（绿色）共享。AutoPipe 是一个弹性流水线模块，通过将冻结层排除在流水线之外并将活跃层打包到更少的 GPU 中来加速训练（粉色），从而减少跨 GPU 通信和缩小流水线气泡。随后，AutoPipe 将流水线长度信息传递给 AutoDP（紫色），然后 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) 相应地优化微批次块大小。

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

跨分区通信开销。 将分区边界放置在跳跃连接的中间会导致额外的通信，因为跳跃连接中的张量现在必须复制到不同的 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” 分别表示微批次 b 在设备 d 上的前向传播、反向传播和优化器更新。每次迭代的总气泡大小是 K-1 乘以每个微批次的前向和反向成本。

此外，这种技术还可以通过缩小流水线气泡的大小来加速训练。为了解释流水线中的气泡大小，图 7 描述了 4 个微批次如何通过一个 4 设备流水线 K = 4 运行。通常，总气泡大小是 (K-1) 乘以每个微批次的前向和反向成本。因此，很明显，较短的流水线具有较小的气泡大小。

动态微批次数

先前的流水线并行系统对每个迷你批次使用固定数量的微批次（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 中的集合通信需要静态成员资格，这会阻止新流水线与现有流水线连接；
状态同步：新激活的进程必须与现有流水线在训练进度（例如，周期数和学习率）、权重和优化器状态、冻结层边界和流水线 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。

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

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

模型和数据集。 实验采用了计算机视觉和自然语言处理领域中两种代表性的 Transformer 模型：Vision Transformer (ViT) 和 BERT。ViT 在图像分类任务上运行，使用 ImageNet21K 上的预训练权重进行初始化，并在 ImageNet 和 CIFAR-100 上进行微调。BERT 在两个任务上运行：来自通用语言理解评估 (GLUE) 基准的 SST-2 数据集上的文本分类，以及 SQuAD v1.1 数据集（斯坦福问答数据集，包含 10 万个众包问答对）上的问答。

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

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

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

整体训练加速

上表总结了整体实验结果。请注意，我们报告的加速是基于保守的 \alpha \frac{1}{3} 值，该值可以获得可比甚至更高的准确率。更激进的 \alpha (\frac{2}{5}, \frac{1}{2}) 可以获得更高的加速，但可能会导致准确率略有下降。请注意，BERT（24 层）的模型大小大于 ViT-B/16（12 层），因此通信需要更多时间。

性能分析

加速分解

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

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

为了了解所有四个组件的效率及其对训练速度的影响，我们尝试了不同的组合，并以它们的训练样本吞吐量（样本/秒）和加速比作为衡量指标。结果如图 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：用于大规模模型分布式训练的自动化弹性流水线

引言