在这篇博文中，我们描述了第一篇经过同行评审的研究论文，该论文探讨了加速 PyTorch DDP (torch.nn.parallel.DistributedDataParallel) [1] 和 Pipeline (torch.distributed.pipeline) 混合使用的技术——PipeTransformer：大规模模型分布式训练的自动化弹性流水线（如 BERT [2] 和 ViT [3] 等 Transformer 模型），该论文发表在 ICML 2021 上。
PipeTransformer 利用自动化弹性流水线技术，实现 Transformer 模型的高效分布式训练。在 PipeTransformer 中，我们设计了一种自适应的即时冻结算法，可以在训练过程中逐步识别并冻结某些层，以及一个弹性流水线系统，可以动态地为剩余的活动层分配资源进行训练。更具体地说，PipeTransformer 会自动将已冻结的层从流水线中排除，将活动层打包到更少的 GPU 中，并复制更多副本来增加数据并行宽度。我们在 ImageNet 数据集上使用 Vision Transformer (ViT)，在 SQuAD 和 GLUE 数据集上使用 BERT 对 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：可解释的冻结训练：DNN 从底部向上收敛（使用 ResNet 在 CIFAR10 上的结果）。每个面板显示使用 SVCCA [17][18] 的逐层相似性。
例如，在冻结训练 [17][18] 中，神经网络通常从底部向上收敛（即并非所有层都需要在整个训练过程中一直训练）。图 2 显示了在这种方法中权重如何在训练过程中逐渐稳定下来的示例。这一观察结果促使我们利用冻结训练来进行 Transformer 模型的分布式训练，通过动态分配资源以专注于逐渐缩小的活动层集合来加速训练。这种层冻结策略对于流水线并行尤为相关，因为将连续的底部层从流水线中排除可以减少计算、内存和通信开销。
图 3：PipeTransformer 的自动化弹性流水线加速 Transformer 模型分布式训练的过程。
我们提出了 PipeTransformer，这是一种弹性流水线训练加速框架，它通过动态改变流水线模型的范围和流水线副本的数量，自动响应已冻结的层。据我们所知，这是第一篇在流水线并行和数据并行训练背景下研究层冻结的论文。图 3 展示了这种组合的好处。首先，通过将已冻结的层从流水线中排除，可以将相同的模型打包到更少的 GPU 中，从而减少跨 GPU 通信并缩小流水线气泡。其次，将模型打包到更少的 GPU 后，同一个集群可以容纳更多的流水线副本，从而增加数据并行宽度。更重要的是，从这两个好处中获得的加速是乘法效应而非加法效应，进一步加速了训练。
PipeTransformer 的设计面临四个主要挑战。首先，冻结算法必须做出即时且自适应的冻结决策；然而，现有工作 [17][18] 仅提供事后分析工具。其次，流水线重新划分结果的效率受多种因素影响，包括划分粒度、跨划分激活大小以及 mini-batch 中的分块（微批次数量），这需要在大型解空间中进行推理和搜索。第三，为了动态引入额外的流水线副本，PipeTransformer 必须克服集合通信的静态特性，并避免在引入新进程（一个流水线由一个进程处理）时可能出现的复杂跨进程消息协议。最后，缓存可以节省已冻结层重复前向传播的时间，但必须在现有流水线和新添加的流水线之间共享，因为系统无法为每个副本创建和预热专用的缓存。
图 4：展示 PipeTransformer 动态过程的动画。
正如动画（图 4）所示，PipeTransformer 设计有四个核心构建块来解决上述挑战。首先，我们设计了一个可调且自适应的算法，生成信号以指导不同迭代中要冻结的层选择（冻结算法）。一旦这些信号触发，我们的弹性流水线模块 (AutoPipe) 会考虑异构划分（冻结层和活动层）的激活大小和工作负载方差，将剩余的活动层打包到更少的 GPU 中。然后，它根据先前针对不同流水线长度的性能分析结果，将一个 mini-batch 分割成最优数量的微批次。我们的下一个模块 AutoDP 会生成额外的流水线副本以占用释放的 GPU，并维护分层通信进程组以实现集合通信的动态成员资格。我们的最后一个模块 AutoCache，可以有效地在现有和新的数据并行进程之间共享激活，并在转换过程中自动替换旧缓存。
总的来说，PipeTransformer 结合了冻结算法、AutoPipe、AutoDP 和 AutoCache 模块，显著提高了训练速度。我们在 ImageNet 上使用 Vision Transformer (ViT)，在 GLUE 和 SQuAD 数据集上使用 BERT 对 PipeTransformer 进行了评估。我们的结果表明，PipeTransformer 在不损失准确性的前提下，实现了高达 2.83 倍的加速。我们还提供了各种性能分析，以便更全面地理解我们的算法和系统设计。最后，我们还为 PipeTransformer 开发了开源的灵活 API，这些 API 在冻结算法、模型定义和训练加速之间提供了清晰的分离，使其能够迁移到需要类似冻结策略的其他算法。
总体设计
假设我们旨在分布式训练系统中训练一个大型模型，该系统采用流水线模型并行和数据并行的混合方案，旨在解决单个 GPU 设备内存不足以容纳模型，或者即使能加载，batch size 也太小以至于不会耗尽内存的场景。更具体地说，我们将我们的设置定义如下：
训练任务和模型定义。 我们在大规模图像或文本数据集上训练 Transformer 模型（例如，Vision Transformer，BERT）。Transformer 模型 mathcal{F} 有 L 层，其中第 i 层由一个前向计算函数 f_i 和一组相应的参数组成。
训练基础设施。 假设训练基础设施包含一个 GPU 集群，该集群有 N 个 GPU 服务器（即节点）。每个节点有 I 个 GPU。我们的集群是同构的，意味着每个 GPU 和服务器具有相同的硬件配置。每个 GPU 的内存容量为 M_\text{GPU}。服务器通过高带宽网络接口（如 InfiniBand 互连）连接。
流水线并行。 在每台机器中，我们将模型 \mathcal{F} 加载到流水线 \mathcal{P} 中，该流水线有 K 个划分（K 也表示流水线长度）。第 kth 个划分 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 Pipeline。对于数据并行，我们使用 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) 相应优化 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：划分边界位于跳跃连接中间。
-
跨划分通信开销。 将划分边界置于跳跃连接中间会导致额外的通信，因为跳跃连接中的张量现在必须复制到不同的 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 环境的混合模式下，通过枚举 M 从 K 到 6K 的值来即时搜索最优 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 实例，并且活动进程使用缓存的流水线模型（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 进行评估。
硬件。 实验在 2 台通过 InfiniBand CX353A (5GB/s) 连接的相同机器上进行，每台机器配备 8 块 NVIDIA Quadro RTX 5000 (16GB GPU 内存)。机器内部的 GPU 到 GPU 带宽（PCI 3.0，16 通道）为 15.754GB/s。
实现。 我们使用 PyTorch Pipe 作为构建块。BERT 模型定义、配置和相关的分词器来自 HuggingFace 3.5.0。我们参考 Vision Transformer 的 TensorFlow 实现，使用 PyTorch 实现了它。更多细节可以在我们的源代码中找到。
模型和数据集。 实验采用了 CV 和 NLP 中两个具有代表性的 Transformer 模型：Vision Transformer (ViT) 和 BERT。ViT 在图像分类任务上运行，使用在 ImageNet21K 上预训练的权重进行初始化，并在 ImageNet 和 CIFAR-100 上进行微调。BERT 在两个任务上运行：在 General Language Understanding Evaluation (GLUE) 基准测试中的 SST-2 数据集上进行文本分类，以及在 SQuAD v1.1 数据集（Stanford Question Answering）上进行问答，该数据集包含 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 输入 patch 大小），SQuAD 1.1 使用 BERT-large-uncased（24 层），SST-2 使用 BERT-base-uncased（12 层）。使用 PipeTransformer，ViT 和 BERT 训练可以分别将每个流水线的 batch size 设置为大约 400 和 64。所有实验的其他超参数（例如 epoch、学习率）在附录中介绍。
总体训练加速
我们在上表中总结了总体实验结果。请注意，我们报告的加速比是基于保守的 \alpha \frac{1}{3} 值，该值可以获得可比甚至更高的准确率。更激进的 \alpha（\frac{2}{5}, \frac{1}{2}）可以获得更高的加速比，但可能会导致准确率略有下降。请注意，BERT 模型的大小（24 层）大于 ViT-B/16（12 层），因此通信需要更多时间。
性能分析
加速分解
本节介绍评估结果，并分析 \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 Distributed APIs 进行分布式训练的整体解决方案。更具体地说，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 16x16 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.
