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

内容导读
本文围绕一篇论文展开，探讨了 PyTorch DDP (torch.nn.parallel) 以及 Pipeline (torch.Distributed.Pipeline) 的加速混合。
本文首发自微信公众号「PyTorch 开发者社区」

论文题目：

PipeTransformer: Automated Elastic Pipelining for Distributed Training of Large-scale Models（PipeTransformer: 用于大规模模型分布式训练的自动化弹性管线）

Pipeptransformer 利用自动化弹性管线，对 Transformer 模型进行高效的分布式训练。在 PipeTransformer 中，我们设计了一个自适应的动态冻结算法，可以在训练过程中逐步识别和冻结某些层，并设计了一个弹性管线系统，可以动态分配资源来训练剩余的活动层。

具体来说，PipeTransformer 自动从管线中排除冻结层，将活动层打包到更少的 GPU 中，并分支更多的副本以增加数据并行宽度。

对 ViT（使用 ImageNet 数据集）和 BERT（用 SQuAD 和 GLUE 数据集）的评估表明，与最先进基线相比，PipeTransformer 在没有精度损失的情况下，实现了高达 2.83 倍的速度提升。

论文中还包括多种性能分析，以方便用户更全面地理解算法和系统设计。

接下来，本文将详细介绍该系统的研究背景、动机、设计思想、设计方案，以及如何用 PyTorch 分布式 API 实现该算法和系统。

论文地址

源代码

幻灯片

简介

大型 Transformer 模型在自然语言处理和计算机视觉方面都取得了准确性突破。GPT-3 为绝大多数 NLP 任务，创造了新的高精度记录。在 ImageNet 中，Vision Transformer（简称 ViT）也达到了89% 的 top-1 准确率，性能优于最先进的卷积网络 ResNet-152 和 EfficientNet。

为了应对模型不断变大的问题，研究人员提出了各种分布式训练技术，包括参数服务器、管线并行、intra-layer 并行以及和 zero redundancy data-parallel。

然而，现有分布式训练的解决方案只是研究场景，所有模型权重都要在训练过程中进行优化（即，在不同迭代过程中，计算和通信 overhead 都要保持相对稳定）。近期关于渐进式训练的研究表明，神经网络中的参数可以动态地进行训练：

• Singular Vector Canonical Correlation Analysis for Deep Learning Dynamics and Interpretability. NeurIPS 2017

论文地址

• Efficient Training of BERT by Progressively Stacking. ICML 2019

论文地址

• Accelerating Training of Transformer-Based Language Models with Progressive Layer Dropping. NeurIPS 2020

论文地址

• On the Transformer Growth for Progressive BERT Training. NACCL 2021

论文地址

图 2：可判读的冻结训练：DNN 自下而上收敛，（用 ResNet 测试在 CIFAR10 上的结果）每个 pane 通过 SVCCA 显示每一层的相似性

例如，在冻结训练中，神经网络通常是自下而上地收敛（即并非所有层都需要通过训练来得到某些结果）。

上图显示了借助类似方法的训练过程中，权重如何逐渐稳定的示例。基于此，我们利用冻结训练对 Transformer 模型进行分布式训练，通过动态分配资源集中于一组缩小的活动层来加速训练。

这种层冻结策略特别适用于管线并行，因为从管线中排除连续的底层可以减少计算、内存和通信开销 (overhead)。

图 3：PipeTransformer 的自动化和弹性管线的流程，进而加速 Transformer 模型的分布式训练

PipeTransformer 是一个弹性管线训练加速框架，可以通过动态转换管线模型的范围和管线副本的数量，自动对冻结层做出反应。

据我们所知，这是第一篇在管线和数据并行训练的背景下，研究层冻结的论文。

图 3 展示了这种组合的优势。

首先，通过从管线中排除冻结层，同样的模型可以打包到更少的 GPU 中，从而减少跨 GPU 通信和更小的管线停顿 (pipeline bubble)。

其次，在将模型封装到更少的 GPU 之后，同一个集群可以容纳更多的管线副本，从而增加数据并行的宽度。

更重要的一点是，两个优势是相乘而不是相加，因此可以进一步加快训练进度。

PipeTransformer 的设计面临四大挑战。

首先，冻结算法必须做出动态以及自适应的冻结决策; 然而，现有的工作只提供了一个后验分析工具。

其次，管线重划分 (pipeline repartitioning) 的效率受多种因素影响，包括分区粒度 (partition granularity)、交叉分区运动尺寸 (cross-partition activation size) 以及小批次组块数量等，这需要在较大的 solution space 中进行推理和搜索。

接着，为了动态地引入额外的管线副本，PipeTransformer 必须克服集体通信的静态特性，并在新进程上线时，避免潜在的复杂跨进程消息传递协议（一个管线只能由一个进程处理）。

最后，缓存可以为冻结层的重复正向传播 (forward propagation) 节省时间，但是它必须在现有的管线和新增加的管线之间共享，因为系统无法为每个副本创建和预热专用缓存。

图 4：PipeTransformer Dynamics 示意图

如图 4 所示，为了应对上述挑战，PipeTransformer 的设计包含了四个核心构建区块。

首先是一个可调节的自适应算法，它可以生成信号，引导在不同的迭代中选定冻结的层（冻结算法）。一旦被这些信号触发，弹性管线模块 (AutoPipe) 就会通过评估异构分区（冻结层和活动层）的活动尺寸和工作负载的变化，将剩余的活动层打包到更少的 GPU 中。

接下来，根据先前对不同管线长度的分析结果，将一个 mini-batch 分解为一系列更好的 micro-batch。

下一个模块 AutoDP 会产生额外的管线副本，来占用已释放的 GPU，并维护分层通信进程组，以获得集体通信的动态成员关系。

最后一个模块 AutoCache，能够有效地在现有和新增的数据并行进程之间共享激活，并在转换期间自动替换陈旧的缓存。

总的来说，PipeTransformer 结合了冻结算法、 AutoPipe、 AutoDP 以及 AutoCache 模块，提供了显著的训练加速。

我们用 ViT（使用 ImageNet 数据集）和 BERT（用 SQuAD 和 GLUE 数据集）模型对 PipeTransformer 进行了评估，结果表明与最先进基线相比，PipeTransformer 在没有精度损失的情况下，实现了高达 2.83 倍的速度提升。

我们还提供了各种性能分析，以便用户更全面地理解算法和系统方面的设计。最后，我们还为 PipeTransformer 开发了开放源码的灵活 API，能对冻结算法、模型定义和训练加速之间进行清晰的分离，允许向需要类似冻结策略的算法转移。

整体设计

假设我们的目标是在一个分布式训练系统中，训练一个大规模模型。这个系统融合了管线模型并行以及数据并行，可用于处理如下场景:

单个 GPU 设备的内存无法容纳模型，或加载时批尺寸很小，得以避免内存耗尽。具体来讲，定义的设置如下:

• 训练任务和模型定义。 在大规模的图像或文本数据集上训练 Transformer 模型（如Vision Transformer、BERT 等）。Transformer 模型 mathcalF 一共有 L 个层，其中第 i 层由一个正向计算函数 fi 和一组对应参数组成。

• 训练基础架构。 假设训练基础架构包含一个具有 N 个 GPU 服务器（即节点）的 GPU 集群。每个节点都有 I 个 GPU。集群是同构的，这意味着每个 GPU 和服务器的硬件配置都是相同的。每个 GPU 的内存容量是 MGPU。服务器通过高带宽的网络接口 ( 如 InfiniBand) 互相连结。

• 管线并行。 每台机器中，我们将一个模型 F 加载到一个具有 K 个分区（K 也表示管线长度）的管线中。第 k 个分区由 Pk 个连续的层组成。假设每个分区由一个 GPU 设备处理。1≤K≤I 表示我们可以在单个设备上，为多个模型副本构建多个管线。

假设一个管线上的所有 GPU 设备都归属于同一台机器，管线为同步管线，不涉及过期梯度，micro-batch 的数量为 M。在 Linux 操作系统中，每个管线都由一个进程处理。更多细节可查阅 GPipe 了解。

• 数据平行。 DDP 是 R 并行 Worker 内部的一个跨机器分布式数据并行处理组。每个 Worker 都是一个管线副本（单个进程）。第 r 个 Worker 的指数 (ID) 是等级 r。

对于 DDP 中的任意两个管线而言，它们既可以属于同一个 GPU 服务器，也可以属于不同的 GPU 服务器，也可以与 AllReduce 算法交换梯度。

在这些情况下，我们的目标是利用冻结训练来加速训练，这使得在整个训练过程中，无需对所有层进行训练。

此外，这还有助于节省计算、通信、内存损耗，并一定程度上避免连续冻结层引起的过拟合。

然而，要利用这些优势，必须要克服前文提到的四个挑战，即设计一个自适应的冻结算法、动态管线重新划分、高效的资源再分配，以及跨进程缓存。

图 5：PipeTransformer 训练系统概览

PipeTransformer 协同设计了一个即时冻结算法和一个自动弹性管线训练系统，可以动态地转换管线模型的范围和管线副本的数量。整体系统架构如图 5 所示。

为了支持 PipeTransformer 的弹性管线，我们维护了 PyTorch Pipeline 的一个定制版本。对于数据并行，我们使用 PyTorch DDP 作为基线。其他库是操作系统的标准机制（如 multi-processing），这也使得无需对软件或硬件进行定制化处理。

为了保证框架具有通用性，我们将训练系统解耦为四个核心组件：冻结算法、AutoPipe、AutoDP 和AutoCache。

冻结算法（灰色部分）从训练循环中取样指标，并做出逐层冻结的决定，这些决定将与 AutoPipe（绿色部分）共享。

AutoPipe 是一个弹性管线模块，通过将冻结层从管线中排除，并将活动层打包到更少的 GPU 中（粉色部分）来加快训练速度，从而减少跨 GPU 的通信，并保持更小的管线停顿。

随后，AutoPipe 将管线长度信息传递给 AutoDP（紫色部分），然后在可能的情况下产生更多的管线副本以增加数据并行的宽度。

图中还包括一个例子，即 AutoDP 引入了一个新的副本（紫色部分）。AutoCache（橙色划线部分）是一个跨管线的缓存模块。为了可读性和通用性，源代码架构与图 5 保持一致。

用 PyTorch API 实现

从图 5 可以看出，PipeTransformer 包括四个组件：Freeze Algorithm、 AutoPipe、 AutoDP 以及 AutoCache。

其中，AutoPipe 和 AutoDP 分别依赖于 PyTorch DDP (torch.nn.parallel.DistributedDataParallel) 以及管线 (torch.distributed.pipeline)。

在本博客中，我们只强调了 AutoPipe 和 AutoDP 的关键实现细节。有关冻结算法和 AutoCache 的详细信息，请参阅论文。

论文地址

AutoPipe：弹性管线

通过将冻结层从管线中排除，并将活动层压缩到较少 GPU 中，AutoPipe 可以加速训练。本节详细阐述了 AutoPipe 的关键组件：

1. 动态分区管线；

2. 减少管线设备的数量；

3. 对应地优化 mini-batch chunk size

PyTorch 管线的基本用法

深入研究 AutoPipe 细节之前，我们先来熟悉一下 PyTorch Pipeline (torch.distributed.pipeline.sync.Pipe) 的基本用法：

参考教程

理解实际中的管线设计，参考以下简单示例：

# 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 设备中，并设置最优 chunk 数量。

平衡各分区的计算时对管线训练速度至关重要，因为各阶段工作负载分布不均会导致滞后，迫使任务量少的设备等待。chunk 数量也可能对管线的吞吐量有非常大的影响。

平衡管线分区

在 PipeTransformer 等动态训练系统中，各个分区只做到参数数量一致，并不能保证训练速度最快，其他因素也会产生关键作用:

图 6：分区边界位于 skip connection 中间

1、跨分区通信 overhead。 将分区边界置于 skip connection 中间会产生额外的通信，因为这样一来，skip connection 中的张量就必须被复制到不同的 GPU 了。

例如对于图 6 中的 BERT 分区，分区 k 必须从分区 k-2 和分区 k-1 中获取中间输出。相比之下，如果边界被放置于 addition layer 之后，分区 k-1 和分区 k 之间的通信 overhead 明显变得更小了。

测量结果表明，跨设备通信比轻度不平衡的分区更昂贵，因此，我们不考虑打破 skip connection。

2、冻结层内存占用。 在训练过程中，AutoPipe 必须多次重新计算分区边界，以平衡两种不同类型的层：冻结层和活动层。

鉴于冻结层不需要 backward activation map、优化器状态和梯度，因此冻结层的内存成本只是非活动层的一小部分。

我们定义了一个可调控的成本因素 lambdafrozen 来评估冻结层对同一活动层的内存占用率，而不是启动侵入性的 profiler 来获取内存和计算成本的根本性指标。根据实验硬件的经验测量，我们将其设置为 1/6。

基于以上两点，AutoPipe 可以根据参数大小来平衡管线分区。具体来说，AutoPipe 借助一个贪心算法来分配冻结层和活动层，从而使得分区的子层能平均分配到 K 个 GPU 设备中。

Pseudocode 即算法 1 中的 load_balance() 函数。冻结层从原始模型中提取出来，保存在管线的首个设备中的一个单独的模型实例 Ffrozen 中。

请注意，本文采用的分割算法并不是唯一的选择；PipeTransformer 是模块化的，可以结合任何替代方案运行。

管线压缩

管线压缩有助于释放 GPU，以容纳更多的管线副本，并减少分区之间的跨设备通信数量。为了确定压缩的时长，我们可以估计压缩后最大分区的内存消耗，然后将其与 timestep T=0 时的管线最大分区的内存消耗进行比较。

为了避免大量的 memory profiling ，压缩算法使用参数大小作为训练内存占用的代理。基于这种简化，管线压缩的准则如下:

一旦收到冻结通知，AutoPipe 就会尝试将管线长度 K 除以 2（例如从 8 到 4，然后是 2）。通过输入 K/2，压缩算法可以验证压缩结果是否满足公式 (1) 中的准则。

伪代码在算法 1 的第 25-33 行中显示。注意，这种压缩使得加速比在训练期间呈指数增长，这意味着如果一个 GPU 服务器包含更多的 GPU（比如超过 8 个），加速比将进一步增大。
图 7：Pipeline Bubble

Fd, b 和 Ud 分别表示设备 d 上 micro=batch b 的前向、后向和优化器更新。

每次迭代中的总 bubble size 是每 micro=batch 前向和后向成本的 K-1 倍。

此外，这种技术还可以通过缩小 Pipeline Bubble 的尺寸来加快训练速度。为了解释管线中的 bubble size，图 7 描述了 4 个 micro-batche 是如何通过 4 个设备管线 K=4 运行的。

一般来说，总的 bubble size 是每 micro-batch 向前和向后耗费的 K-1 倍。因此，很明显，较短的管线 bubble size 也更小。

Micro-Batch 的动态数量

以前的管线并行系统用的是固定数量的 micro-batches per mini-batch(M)。GPipe 建议 M≥4 x K，其中 K 是分区数量（管线长度）。但是，考虑到 PipeTransformer 会对 K 进行动态配置，我们发现在训练期间保持静态 M 效果并不好。

此外，当与 DDP 集成时，M 的值也会影响 DDP 梯度同步的效率。由于 DDP 在梯度同步之前必须等待最后一个 micro-batch 完成对某一参数的后向计算，因此 micro-batch 越细，计算和通信的重叠就越小。

因此，PipeTransformer 没有使用静态值，而是通过枚举 K-6K 范围内 M 的值，在 DDP 环境的 hybrid 中动态地搜索 M 的最优值。对于特定的训练环境，profiling 只需要完成一次即可（参见算法 1 第 35 行）。

完整源码，请参考

AUTODP：生成更多的管线副本

考虑到 AutoPipe 可以把相同的管线压缩到更少的 GPU 中，AutoDP 可以自动生成新的管线副本，以增加数据并行的宽度。

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

1、DDP 通信： PyTorch DDP 中的集体通信要求静态成员身份，这就阻止了新管线与现有管线之间的连接;

2、状态同步: 新激活的进程必须在训练过程（如epoch 数量和学习率）、权重和优化器状态、冻结层边界以及管线 GPU 范围等方面，与现有的管线保持一致;

3、数据集再分配： 应该重新平衡数据集，以匹配动态的管线数量。这不仅可以避免掉队，而且可以确保所有 DDP 过程的梯度都是平等加权的。

图 8：AutoDP: 借助两个进程组之间的信息，处理动态数据并行

注：进程 0-7 属于 machine 0，进程 8-15 属于 machine 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

关于代码，请参考

实验部分

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

硬件。 实验是在两台由 InfiniBand CX353A (GB/s) 连接的相同的机器上进行的，每台机器配备 8 NVIDIA Quadro RTX 5000（16GB GPU 内存）。机器内 GPU 到 GPU 的带宽 (PCI 3.0,16 lanes) 为 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 数据集 (Stanford Question Answering) 上进行智能问答。SQuAD v1.1 数据集包括 10 万个众包问答组。

训练计划。 大型模型通常需要数千个 GPU-days {\emph{e.g.}, GPT-3)，如果从头开始训练，那么借助预训练模型对下游任务进行微调，已成为 CV 和 NLP 领域的一种趋势。此外，PipeTransformer 是一个复杂的训练系统，涉及多个核心组件。因此，对于第一版 PipeTransformer 的系统开发和算法研究来说，使用大规模的预训练从头开发和评估，成本效益很低。因此，本节介绍的实验侧重于预训练模型。请注意，由于预训练和微调中的模型体系结构是相同的，所以 PipeTransformer 可以同时满足这两个要求。我们在附录中讨论了预训练结果。

基线。 本节中的实验横向对比了 PipeTransformer 与最先进的框架 PyTorch Pipeline (PyTorch 的实现 GPipe) 与 PyTorch DDP 三种方案。由于这是第一篇研究通过冻结层来加速分布式训练的论文，因此还没有完全一致的对应解决方案。

超参数。 对于数据集 ImageNet 和 CIFAR-100，实验中使用的是 ViT-B/16（Transformer 层数为 12，input patch size 为 16 x16）。对于 SQuAD 1.1，实验中使用的是 BERT-large-uncased（24 层）。SST-2 用的是 BERT-base-uncased（12 层）。通过 PipeTransformer，ViT 和 BERT 训练可以将吗 per=pipeline 的批尺寸分别设置为 400 和 64 左右。其他超参数（如 epoch，learning rate 等）见附录。

整体加速训练

上表总结了整体实验结果。注意，此处的加速比是基于一个保守值 a*(1/3)，这个值可以获得相似甚至更高的准确度。如果取值 a*(2/5,1/2) 可以获得更高的加速比，但会对准确度造成轻微损失。此外，BERT (24 层) 比 ViT-B/16(12 层) 大 ，因此需要更多的通信时间。

性能分析

加速比分解

这部分介绍了评估结果，并分析了不同组件在 /AutoPipe 中的性能。

为了了解这四个组件的功效以及它们对训练速度的影响，我们用不同的组合进行了实验，并以它们的训练样本吞吐量 (samples/second) 和加速比作为度量指标。结果如图 9 所示。实验结果的关键要点包括：

1、主要加速比是 AutoPipe 和 AutoDP 联合实现的弹性管线的结果;

2、AutoCache 的作用被 AutoDP 放大了;

3、冻结训练是独立进行的，没有进行系统调整和训练降速。

冻结算法中调整 a

图 10：冻结算法中调整 a

我们进行了一些实验来说明冻结算法是如何影响训练速度的。结果表明，a (excessive freeze) 越大，加速比越大，但是会有轻微的性能下降。在图 10 所示的例子中，当 a=1/5 时，冻结训练效果优于普通训练，加速比达到 2.04。

弹性管线中的最优 Chunk 数量

图 11：弹性管线中的最佳 Chunk 数目

我们对不同管线长度 K 的最佳 micro-batch 数量 M 进行了分析。结果如图 11 所示。正如我们所看到的，K 的取值不同，最佳数量 M 也会相应地发生变化，当 M 取值不同时，吞吐量 gap 也会变大（如图所示当 K=8 时），这也证实了在弹性管线中使用 anterior profiler 的必要性。

理解缓存时序

图 12：缓存时序

为了评估 AutoCache，我们比较了从 epoch 0开始就是用 AutoCache 的训练任务（蓝色线），以及不使用 AutoCache 的训练任务（红色线）的样本吞吐量。

图 12 显示，过早启用 caching 会减慢训练速度，因为相较于在数量更少的冻结层上前向传播，这样的花费更高。在冻结更多层之后，缓存激活的性能明显优于相应的前向传播。因此，AutoCache 使用一个 Profiler 来确定启用缓存的正确时间。

在我们的系统中，对于 ViT (12 层) ，缓存从第 3 个冻结层开始；而对于 BERT (24 层) 而言，缓存则从第 5 个冻结层开始。

总结

本文介绍了 PipeTransformer，它是一个整体解决方案，结合了弹性管线并行和数据并行，使用 PyTorch 分布式 API 进行分布式训练。

具体来说，PipeTransformer 可以逐步冻结管线中的 layer，将剩余的活动层打包到更少的 GPU 中，并分叉更多的管线副本以增加数据并行宽度。对 ViT 和 BERT 模型的评估表明，与最先进的基线相比，PipeTransformer 使得速度提升了 2.83 倍，而且没有精度损失。

论文地址

开放源码

PPT

blog 地址