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

论文题目:

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 地址