在这篇博文中，我们介绍了第一篇经过同行评审的研究论文，该论文探讨了加速 PyTorch DDP (torch.nn.parallel.DistributedDataParallel) [1] 和 Pipeline (torch.distributed.pipeline) 的混合 - PipeTransformer：大规模模型分布式训练的自动化弹性流水线（Transformer，例如 BERT [2] 和 ViT [3]），发表于 ICML 2021。
PipeTransformer 利用自动化弹性流水线，实现 Transformer 模型的高效分布式训练。在 PipeTransformer 中，我们设计了一种自适应的即时冻结算法，可以在训练期间逐步识别和冻结某些层，以及一个弹性流水线系统，可以动态分配资源来训练剩余的活动层。更具体地说，PipeTransformer 自动将冻结层从流水线中排除，将活动层打包到更少的 GPU 中，并派生更多副本以增加数据并行宽度。我们使用 Vision Transformer (ViT) 在 ImageNet 上，以及 BERT 在 SQuAD 和 GLUE 数据集上评估 PipeTransformer。我们的结果表明，与最先进的基线相比，PipeTransformer 在不损失准确率的情况下，速度提升高达 2.83 倍。我们还提供了各种性能分析，以便更全面地理解我们的算法和系统级设计。
接下来，我们将介绍背景、动机、我们的想法、设计，以及如何使用 PyTorch Distributed 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。
- 关于 Transformer 增长以进行渐进式 BERT 训练。NACCL 2021
图 2. 可解释的冻结训练：DNN 自下而上收敛（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，这些 API 在冻结算法、模型定义和训练加速之间提供了清晰的分离，从而可以转移到其他需要类似冻结策略的算法。
总体设计
假设我们的目标是在分布式训练系统中训练一个大型模型，其中流水线模型并行和数据并行的混合用于针对以下场景：单个 GPU 设备的内存无法容纳模型，或者如果加载了模型，则批量大小太小以至于不会耗尽内存。更具体地说，我们将我们的设置定义如下
训练任务和模型定义。 我们训练 Transformer 模型（例如，Vision Transformer、BERT）在大型图像或文本数据集上。Transformer 模型 mathcal{F} 具有 L 层，其中第 i 层由前向计算函数 f_i 和相应的参数集组成。
训练基础设施。 假设训练基础设施包含一个 GPU 集群，该集群具有 N 个 GPU 服务器（即节点）。每个节点有 I 个 GPU。我们的集群是同构的，这意味着每个 GPU 和服务器都具有相同的硬件配置。每个 GPU 的内存容量为 M_\text{GPU}。服务器通过高带宽网络接口（例如 InfiniBand 互连）连接。
流水线并行。 在每台机器中，我们将模型 \mathcal{F} 加载到具有 K 个分区（K 也代表流水线长度）的流水线 \mathcal{P} 中。第 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 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) 相应地优化小批量分块大小。
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) 倍。因此，很明显，较短的流水线具有较小的气泡大小。
微批量的动态数量
先前的流水线并行系统对每个小批量使用固定数量的微批量 (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。我们使用 PyTorch 按照其 TensorFlow 实现实现了 Vision Transformer。更多详细信息可以在我们的源代码中找到。
模型和数据集。 实验采用 CV 和 NLP 中两个有代表性的 Transformer：Vision Transformer (ViT) 和 BERT。ViT 在图像分类任务上运行，使用在 ImageNet21K 上预训练的权重进行初始化，并在 ImageNet 和 CIFAR-100 上进行微调。BERT 在两个任务上运行，文本分类在来自通用语言理解评估 (GLUE) 基准测试的 SST-2 数据集上，问题解答在 SQuAD v1.1 数据集（斯坦福问答）上，该数据集是 10 万个众包问题/答案对的集合。
训练方案。 鉴于大型模型通常需要数千个 GPU 日 {\emph{例如}，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. 速度提升分解（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 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 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., 和 Su, B. Y. 使用参数服务器扩展分布式机器学习。 第11届 {USENIX} 操作系统设计与实现研讨会 ({OSDI} 14), 页码 583–598, 2014年。
[7] Jiang, Y., Zhu, Y., Lan, C., Yi, B., Cui, Y., 和 Guo, C. 用于加速异构 GPU/CPU 集群中分布式 DNN 训练的统一架构。 第14届 USENIX 操作系统设计与实现研讨会 (OSDI 20), 页码 463–479。 USENIX Association, 2020年11月。 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., 和 Chun, B. G. Parallax: 深度神经网络的稀疏感知数据并行训练。 第十四届 EuroSys Conference 2019 会议论文集, 页码 1–15, 2019年。
[9] Kim, C., Lee, H., Jeong, M., Baek, W., Yoon, B., Kim, I., Lim, S., 和 Kim, S. TorchGPipe: 用于训练巨型模型的即时流水线并行。
[10] Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, M. X., Chen, D., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., 等. Gpipe: 使用流水线并行高效训练巨型神经网络。
[11] Park, J. H., Yun, G., Yi, C. M., Nguyen, N. T., Lee, S., Choi, J., Noh, S. H., 和 ri Choi, Y. Hetpipe: 通过集成流水线模型并行和数据并行，在（微弱的）异构 GPU 集群上实现大型 DNN 训练。 2020 USENIX Annual Technical Conference (USENIX ATC 20), 页码 307–321。 USENIX Association, 2020年7月。 ISBN 978-1-939133- 14-4.
[12] Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N. R., Ganger, G. R., Gibbons, P. B., 和 Zaharia, M. Pipedream: 用于 DNN 训练的通用流水线并行。 第27届 ACM 操作系统原理研讨会论文集, SOSP ’19, 页码 1–15, 纽约, NY, 美国, 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., 和 Chen, 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年。
