零气泡流水线并行

介绍

本文是Zero Bubble (Almost) Pipeline Parallelism的笔记及译文。

Notes

摘要

流水线并行是大规模分布式训练的关键组成部分之一，但其效率一直受到流水线气泡的影响，而这些气泡此前被认为是不可避免的。本文提出了一种调度策略，据我们所知，这是首次在同步训练语义下成功实现零流水线气泡。该改进的核心思想是将反向计算拆分为两部分：一部分用于计算输入的梯度，另一部分用于计算参数的梯度。基于这一思想，我们手工设计了新的流水线调度方案，在性能上显著优于基线方法。进一步地，我们还提出了一种算法，可根据具体的模型配置和内存限制自动搜索最优调度方案。此外，为了真正实现零气泡，我们引入了一种新的技术，在优化器步骤中绕过同步开销。实验评估表明，在相近的内存限制下，我们的方法在吞吐量上相比 1F1B 调度最高可提升 15%；当放宽内存约束时，这一提升可进一步达到 30%。我们相信，这些结果标志着在挖掘流水线并行真实潜力方面迈出了重要一步。基于 Megatron-LM 的源码已在以下地址公开：https://github.com/sail-sg/zero-bubble-pipeline-parallelism。

引言

分布式模型训练已成为深度学习社区的核心议题，尤其是在模型规模和复杂度不断增长的背景下。训练这些庞大的模型通常需要大量 GPU，并通过多种拓扑结构互连。过去几年中，人们提出了多种用于训练深度神经网络（DNN）的并行化技术。数据并行（DP）由于其简单性，通常是中小规模模型的默认选择。然而，当模型规模进一步增大时，已无法将全部模型参数放入单个 GPU 中，此时模型并行便成为解决方案。模型并行主要包括两种形式：张量并行（TP）和流水线并行（PP）。TP 将单层中的矩阵乘法拆分到多个设备上执行，而 PP 则将整个模型划分为多个阶段，在不同设备上依次处理。值得注意的是，ZeRO 通过在设备间切分参数，为模型并行提供了一种有力的替代方案，同时保持了 DP 的简洁性。

近期研究表明，在大规模训练场景下获得最优性能，需要 DP、TP 和 PP 策略之间进行非平凡的协同。在互连资源充足的情况下，例如同一计算节点内 GPU 之间通过 NVLink 互连，DP、TP 与 ZeRO 的混合策略能够高效运行。而大量实证研究表明，PP 在跨服务器连接的利用上具有显著优势，尤其是在数千 GPU 的规模下。这突出了我们工作的主要目标：提升 PP 的效率。

进一步深入 PP 的实现细节，其效率在很大程度上取决于设备空闲时间的多少，这种空闲时间通常被称为流水线气泡。由于层之间存在依赖关系，气泡似乎不可避免。早期一个重要的工作是 GPipe，其通过增加流水线中并发批次的数量来降低气泡比例，但直接后果是峰值内存需求的增加。为缓解这一问题，GPipe 在前向过程中丢弃部分中间激活，并在反向传播时重新计算它们，然而这引入了约 20% 的额外计算开销。

另一类改进工作聚焦于异步流水线并行，包括 PipeDream 和 PipeMare。从理论上看，异步 PP 可以实现无气泡，从而大幅提升流水线效率，但代价是牺牲了精确的优化语义。与此同时，在同步设置下也有相应的改进方法。其中一种著名的调度策略是一次前向一次反向（1F1B）。该方法最早在异步设置中提出，随后被引入到同步训练中。1F1B 通过更早地调度反向传播，加快了内存释放速度。在相同微批次数量下，它具有相似的气泡比例，但在峰值内存方面具有明显优势。在此基础上，又提出了 1F1B interleaved 策略，通过在同一设备上分配多个阶段，进一步减少气泡大小，但代价是更高的通信开销和更大的峰值内存需求。

尽管已有诸多努力，但在同步训练语义下，残余的气泡至今仍是 PP 面临的最大问题。本文发现，通过以更细粒度来表示和调度计算图，PP 仍有进一步优化的空间。传统深度学习框架通常以层为粒度进行设计，而现代深度学习编译器则采用多种中间表示，在不同层级上进行优化。尽管更细的粒度意味着更大的搜索空间，但由于缺乏有效的优化工具来探索这一空间，往往难以实施。因此，选择一种合适的粒度至关重要。

图 1：用于 MLP 的计算图

传统上，神经网络通常被划分为层的堆叠结构。每一层都关联两个函数：前向和反向。在前向过程中，输入 $x$ 通过带参数的映射 $f(x, W)$ 被转换为输出 $y$。反向过程对训练至关重要，包含两项计算： $\nabla_x f(x, W)^\top \dfrac{d\ell}{dy}$ 和 $\nabla_W f(x, W)^\top \dfrac{d\ell}{dy}$。 它们分别计算关于输入 $x$ 以及该层参数 $W$ 的梯度。为方便起见，我们分别用字母 B 和 W 表示这两项计算，并用 F 表示前向过程（如图所示）。传统上，B 和 W 被组合在一起，作为单一的反向函数提供给用户。这种设计在概念上对用户友好，并且在 DP 中运行良好，因为第 $i$ 层权重梯度的通信可以与第 $i-1$ 层的反向计算重叠。然而，在 PP 中，这种设计不必要地增加了顺序依赖的计算，即第 $i-1$ 层的 B 依赖于第 $i$ 层的 W，这通常会损害流水线效率。

基于将 B 和 W 拆分的思想，我们提出了新的流水线调度方案，大幅提升了流水线效率。本文其余部分组织如下：在第二节中，我们在一个理想化假设下介绍手工设计的调度方案，该假设认为 F、B 和 W 的执行时间相同。随后，在第三节中，我们去除这一假设，并提出一种在更现实条件下工作的自动调度算法。为实现零气泡，第四节详细描述了一种在优化器步骤中绕过同步需求的方法，同时保持同步训练语义。最后，我们在多种设置下，将所提出的方法与基线方法进行了实证评估。

需要指出的是，我们并不旨在探索大规模分布式训练中的通用混合策略。相反，我们的目标是专注于提升流水线调度的效率，并在严格可比的条件下与基线方法进行比较。我们的方法与 DP、TP 和 ZeRO 策略是正交的，可以作为大规模训练中 PP 部分的并行替代方案。

手工设计的流水线调度

基于将 B 和 W 拆分可以减少顺序依赖、从而提升效率这一关键观察，我们从广泛使用的 1F1B 调度出发重新设计流水线。如图所示，1F1B 以一个预热阶段开始。在该阶段中，不同的工作节点执行不同数量的前向计算，通常每个阶段比其后继阶段多执行一次前向计算。预热阶段之后，每个工作节点进入稳态阶段，在该阶段中交替执行一次前向和一次反向计算，从而在各阶段之间保持均衡的工作负载。在最后阶段，每个工作节点处理仍在流水线中的微批次的反向计算，完成整个批次。

在我们改进的版本中，反向过程被拆分为 B 和 W 两个阶段。需要强调的是，同一微批次的 F 和 B 在不同流水线阶段之间仍然必须保持顺序依赖。然而，W 可以在同一阶段对应的 B 之后灵活地调度。这使得我们能够策略性地安排 W 来填补流水线气泡。实际上，存在多种相较于 1F1B 更优的调度方案，它们在气泡大小与内存占用之间提供了不同的权衡。在本节中，我们介绍两种特别有代表性的手工设计调度方案，以展示更细粒度在减少流水线气泡方面的巨大潜力（见图）。为便于初始设计的阐述，我们假设 F、B 和 W 的时间开销相同，这一假设也被早期研究所采用。然而，在第三节中，我们将重新审视这一假设，以在真实场景下进一步优化调度效率。

图 2：1F1B 流水线调度

图 3：手工设计的流水线调度，上：ZB-H1；下：ZB-H2

内存高效调度

我们提出的第一个手工设计调度称为 ZB-H1，其目标是确保所有工作节点的最大峰值内存使用量不超过 1F1B。ZB-H1 在整体上遵循 1F1B 的调度方式，但会根据预热阶段中微批次的数量调整 W 的启动时间。这一设计保证了所有工作节点保持相同数量的在途微批次。因此，如图所示，其流水线气泡大小被降低到 1F1B 的 三分之一。这一缩减源于这样一个事实：与 1F1B 相比，B 在所有工作节点上都更早启动，而尾部的气泡则由稍后开始的 W 填充。由于 W 通常比 B 占用更少的内存（见表），第一个工作节点具有最大的峰值内存使用量，这一点与 1F1B 保持一致。

零气泡调度

当允许比 1F1B 更大的内存占用，并且微批次数量足够多时，就可以实现零气泡调度，我们将其记为 ZB-H2。如图所示，在预热阶段我们引入了更多的前向计算 F，以填补首次 B 之前的气泡。同时，我们重新排列了尾部的 W，将整体布局从梯形变为平行四边形，从而消除了流水线中的所有气泡。需要特别指出的是，在该调度中，优化器步骤之间的同步被移除了；我们将在第四节中讨论这一操作如何在保持同步训练语义的前提下安全实现。

定量分析

我们用 $p$ 表示流水线阶段数，用 $b$ 表示每个微批次的大小。对于 Transformer 架构，记注意力头数为 $a$，序列长度为 $s$，隐藏维度大小为 $h$。我们使用 MB/MW 表示一次 B/W 计算所需存储激活的内存，用 TF / TB / TW 表示一次 F/B/W 计算的运行时间。为简化分析，我们仅对 Transformer 架构进行定量分析，采用与 GPT-3 类似的典型设置，其中前馈网络内部的隐藏维度为 $4h$，每个注意力头的维度为 $h/a$。

与先前工作一致，在计算 FLOPs 时我们仅考虑矩阵乘法操作，因为它们占据了 Transformer 层中绝大多数的计算量。对于前向过程中的每一次矩阵乘法，在相应的反向过程中都会有两次具有相同 FLOPs 的矩阵乘法（见图），它们分别属于 B 或 W。每个 Transformer 层的 FLOPs 近似计算公式见表。可以看出，$TW < TF < TB$，并且满足 $$ TB + TW = 2TF $$ 我们采用与先前工作相同的方法来估计 B 所需的激活内存。当 B 完成后，它会释放部分不再使用的激活，但仍会保留一些额外的梯度（如图中的 $\nabla z_L$）供 W 使用。因此，如表所示，W 所需的总内存小于 B。

表 1：每个 Transformer 层在不同计算阶段所需的 FLOPs 与激活内存

在不再假设 $TF = TB = TW$ 的情况下，ZB-H1 和 ZB-H2 的峰值激活内存与气泡大小如表所示。具体而言，工作节点 $i$ 的激活内存需求为：

• 对于 ZB-H1：$(p - i + 1)MB + (i - 1)MW$
• 对于 ZB-H2：$(2p - 2i + 1)MB + (2i - 2)MW$

如表所示，W 所需的激活内存小于 B，因此 ZB-H1 和 ZB-H2 的峰值激活内存分别为 $pMB$ 和 $(2p - 1)MB$。

表 2：1F1B 与手工设计调度方案的对比

自动流水线调度

尽管手工设计的调度方案具有简单性和可理解性，但在实际应用中仍面临若干问题。首先，在假设 $TF = TB = TW$ 的条件下进行调度，往往会引入不必要的气泡，尤其是在这些时间差异显著的模型中。其次，手工调度通常忽略了阶段之间传输激活或梯度所需的通信时间（记为 $T_{comm}$），这会在流水线执行中引入明显的延迟。最后，当可用内存不足以容纳足够多的微批次以实现无气泡调度时，在最小化气泡大小与满足内存限制之间取得平衡将变得尤为困难。

为了解决上述挑战，并确保在实际场景中的泛化能力，我们提出了一种算法，在给定流水线阶段数 $p$、微批次数 $m$、激活内存上限 $M_{limit}$ 以及运行时间估计 $TF$、$TB$、$TW$ 和 $T_{comm}$ 的情况下，自动搜索最优调度方案。我们设计了一种启发式策略，当 $m$ 足够大时，该策略通常能够生成最优或近似最优的解。此外，我们还将该问题系统地形式化为整数线性规划（更多细节见附录），在问题规模受限时，可使用现成的 ILP 求解器进行求解。这两种方法可以结合使用：先以启发式解作为初始化，再通过 ILP 进一步优化。

启发式算法

我们的启发式算法包含以下步骤：

• 在预热阶段，在不超过内存限制的前提下，尽可能调度更多的前向计算 F，以最小化首次 B 之前的气泡。如果在未达到内存上限的情况下继续调度一个 F 会延迟后续 B，则可能在首次 B 之前仍然存在一个小于 $TF$ 的气泡。我们通过一个二值超参数来控制是否执行该操作。

• 在预热阶段之后，我们遵循交替调度一个 F 和一个 B 的模式。当出现大于 $TW$ 的空隙时，插入 W 来填补气泡。当气泡大小小于 $TW$ 时，如果当前气泡会使所有阶段中的最大累计气泡增大，我们仍会插入一个 W。当达到内存上限时，也会插入 W 以回收部分内存。通常情况下，该启发式策略会进入遵循 1F-1B-1W 模式的稳态阶段。

• 在整个过程中，始终保证流水线阶段 $i$ 在任意时刻所调度的 F 数量至少比阶段 $i+1$ 多一个，直到 F 被用尽。当这一差值超过 1 时，我们使用另一个二值超参数来决定是否在阶段 $i$ 跳过一个 F，前提是不引入更多气泡。我们通过网格搜索来寻找最佳的超参数组合。

• 在每个阶段，当 F 和 B 执行完毕后，依次调度所有剩余的 W。

绕过优化器同步

在大多数 PP 的实践中，为了保证数值稳定性，通常会在优化器步骤中对所有流水线阶段进行同步。例如，在梯度裁剪中需要计算全局梯度范数；在混合精度训练中，需要进行全局的 NAN 和 INF 检查。这些操作都依赖于跨所有阶段的 all-reduce 通信。然而，在优化器步骤中的同步会破坏平行四边形结构（见图），使得零气泡变得不可能。

基于此，我们提出了一种替代机制，在保持同步优化语义的同时，绕过这些同步操作。在现有实现中，通常先发起一次 all-reduce 来收集全局状态，然后根据该全局状态执行优化器步骤。我们注意到，在大多数情况下，这些全局状态实际上并不会产生影响：例如，在数值稳定的设置下，NAN/INF 检查很少触发；梯度裁剪在经验上发生的频率也较低，不足以支撑在每一次迭代中都进行全局同步。

基于这些观察，我们提出用更新后的验证来替代事前同步。其思想如图所示：在每个阶段执行优化器步骤之前，从前一阶段接收一个部分规约的全局状态，将其与当前阶段的本地状态合并后，再传递给下一阶段。每个阶段的优化器步骤由该部分规约状态控制，例如在检测到 NAN 或部分规约的梯度范数超过裁剪阈值时跳过更新。在下一次迭代的预热阶段，完全规约的全局状态会从最后一个阶段反向传播回第一个阶段。各阶段在接收到全局状态后，对上一轮优化器步骤进行验证；如果需要对梯度进行修正，则触发一次回滚，随后基于完全规约的全局状态重新执行优化器步骤。

图 4：用于替代优化器同步的事后验证策略

实验

实验设置

我们的实现基于开源的 Megatron-LM 项目，并使用与 GPT-3 类似的模型进行性能评估，如表所示。在实验过程中，我们首先运行一定数量的迭代以进行性能分析，收集 $TF$、$TB$、$TW$ 以及 $T_{comm}$ 的经验测量值。获得这些数值后，将其输入到我们的自动流水线调度算法中，以确定最优调度方案。需要注意的是，流水线的首尾阶段相比中间阶段各少一个 Transformer 层。这样设计是为了补偿首尾阶段额外的嵌入查找和损失计算，从而避免它们成为瓶颈并向其他阶段引入气泡。

表 3：实验中使用的模型与固定设置

对比方法包括：

• ZB-1p：自动搜索得到的调度方案，其激活内存限制为 $pMB$，在理论上与 1F1B 具有相同的峰值内存占用。
• ZB-2p：自动搜索得到的调度方案，其激活内存限制为 $2pMB$，这是在经验上实现接近零气泡所需的最小内存量。
• 1F1B 与 1F1B-I：分别为 Harlap et al. 和 Narayanan et al. 提出的 1F1B 及其交错版本，采用 Megatron-LM 中的实现。对于交错 1F1B，整个模型被划分为一系列块，并由各阶段循环处理，形成交错流水线。在我们的交错实验中，始终使用最大块数以尽量减少气泡，即每个 Transformer 层作为一个块。

我们的实验最多使用 32 张 NVIDIA A100 SXM 80G GPU，分布在 4 个节点上，并通过 RoCE RDMA 网络互连。在若干预热迭代之后记录每次迭代的运行时间。得益于 Megatron-LM 实现所提供的可复现性，我们无需将模型训练至收敛即可验证 ZB-1p 和 ZB-2p 的正确性。我们使用固定的随机种子初始化模型，记录 ZB-1p、ZB-2p 和 1F1B 在每次迭代后的损失值，并验证它们在逐位意义上完全一致。

主要结果

图 5：不同流水线调度方案下的吞吐量对比

我们在图中展示了所有方法的吞吐量，并在表中给出了各实验设置的更多细节。实验结果表明，ZB-2p 在各种设置下都能稳定地优于其他方法。值得注意的是，1F1B、1F1B-I 和 ZB-1p 的吞吐量与微批次数量呈现出显著的正相关关系；而 ZB-2p 即使在微批次数较少的情况下，仍能保持较高效率。这是因为 ZB-2p 的气泡率几乎降为零，其吞吐量已经接近理论上限。这里的上限是通过将 1F1B 的吞吐量乘以 $$ \frac{1}{1-\text{气泡率}_{\text{1F1B}}} $$ 进行近似估计的。

如前所述，ZB-2p 的效率提升是以比 1F1B 更高的内存消耗为代价的。我们还在附录中比较了在相同内存消耗下的 ZB-2p 与 1F1B，实验结果表明，即使 ZB-2p 的微批次大小仅为 1F1B 的一半，其吞吐量仍然更高。

相比之下，ZB-1p 的设计目标是使峰值内存开销与 1F1B 基线相当。在 8 GPU 的设置中，它表现出与 1F1B-I 相近的吞吐量。在多节点设置下，由于通信带宽更容易成为瓶颈，ZB-1p 明显优于 1F1B-I，这凸显了其在不引入额外通信成本的前提下减少流水线气泡的优势。

在大多数实验设置中，我们将微批次数 $m$ 设置为大于阶段数 $p$，因为这是流水线并行中更常见的使用场景。此外，我们还在附录中给出了 $m \le p$ 情况下的实验结果，在相近的内存消耗下，吞吐量提升可达到 20%–30%。

自动调度效率分析

表 5：在不同设置下，1F1B、1F1B-I、ZB-H1、ZB-H2、ZB-1p、ZB-2p 的气泡率

图 6：ZB-2p 生成的调度（上）及其实际执行过程的性能分析结果（下）

我们研究了由自动调度算法生成的调度方案的效率。实验设置与主要实验相同，但由于本节旨在分析自动调度算法本身的效率，这里的结果基于理论计算而非真实实验。为了量化流水线调度的效率，我们引入了气泡率这一概念，其计算方式为： $$ \frac{\text{cost} - m(TF + TB + TW)}{\text{cost}} $$ 其中，$\text{cost}$ 定义为所有阶段中最大执行时间，该值基于已分析得到的 $TF$、$TB$、$TW$ 和 $T_{comm}$ 对每一种调度方案进行计算。$m(TF + TB + TW)$ 表示在所有通信都与计算完全重叠、流水线中不存在任何气泡时的最优执行时间。

不同调度方案的气泡率如表所示。我们将手工设计的 ZB-H1 和 ZB-H2 作为自动搜索调度的对照基线。在大多数设置下，ZB-2p 的气泡率低于 1%，在所有方案中表现最佳。相比之下，ZB-H2 的表现始终劣于 ZB-2p。这为我们的自动调度算法能够利用更精确的 $TF$、$TB$、$TW$ 和 $T_{comm}$ 估计，从而更好地适应真实场景，提供了有力证据。相反，在 ZB-1p 与 ZB-H1 之间并未观察到类似改进，推测其原因在于内存限制成为了主导因素。值得注意的是，我们提出的所有方法都显著优于 1F1B。

我们还在 16 GPU 的设置下，将 ZB-2p 的调度方案与其实际执行过程的性能分析结果进行可视化对比，以直观证明其确实为零气泡调度。如图所示，自动生成的 ZB-2p 调度几乎不存在气泡；实际执行过程中虽然出现了少量气泡，但整体对齐性依然良好。

图 7：使用启发式算法时，内存上限与气泡率之间的关系

内存上限分析

为了更好地理解内存上限的影响，我们研究了气泡率与 $M_{limit}$ 之间的关系。我们在一系列不同的 $M_{limit}$ 下运行启发式算法，并将结果绘制于图中。最初，随着 $M_{limit}$ 的增大，气泡率呈现出近似线性的下降趋势。从理论上讲，该曲线应在 $$ \frac{(p-1)(TB + 2T_{comm}) + pTF}{TF \cdot MB} $$ 附近趋于平稳。经验结果表明，当 $TF \approx TB$ 且 $T_{comm}$ 相对较小时，$2pMB$ 是实现接近零气泡率的一个良好阈值。

在拐点之后，尽管足够大的内存上限在理论上能够实现零气泡率，但通常代价大于收益。更多细节请参见附录。

结论与讨论

本文提出了一种新的策略，通过在反向计算中拆分激活梯度与参数梯度来提升流水线并行的效率，并设计了一种自动流水线调度算法，能够在不同内存预算下最小化流水线气泡率。该算法生成的调度方案在各类设置下均持续优于 1F1B，并且能够实现接近零的气泡率。

从实验角度看，实现零气泡通常需要相较于 1F1B 约 两倍的激活内存，这可能引发显存不足的担忧。然而，根据附录中的结果，我们认为在大模型训练中，用一定的内存开销来换取零气泡流水线调度是值得的。诸如 ZeRO 和张量并行等策略可用于缓解额外的内存需求。

零气泡调度的另一优势在于，它可以在更少的微批次（通常约为 $3p$）下达到最优效率，这意味着可以将更多微批次划分到数据并行维度上，从而为大模型训练带来更好的可扩展性。