文章标题：具有可控内存的流水线并行
作者/机构：Penghui Qi, Xinyi Wan, Nyamdavaa Amar, Min Lin (Sea AI Lab, National University of Singapore)

A1 主要贡献

本文旨在解决现有流水线并行（Pipeline Parallelism, PP）调度策略中存在的两个主要问题：流水线气泡（pipeline bubbles）和巨大的激活内存（activation memory）消耗。现有方法大多缺乏系统性的设计方法论，导致内存效率低下。

研究目标与核心思想：
本文提出了一个统一的框架，用于设计和分析流水线并行调度。该框架的核心思想是将任何流水线调度分解为重复一个基本构建块（building block）的过程。通过这个框架，作者发现了一个关键洞见：构建块中每个阶段（stage）的生命周期（lifespan）直接决定了该流水线调度的峰值激活内存。生命周期指的是从前向传播（Forward pass）开始到反向传播（Backward pass）结束的时间跨度。

主要创新点与贡献：
1. 提出一个构建流水线的四步框架：
* 构建块（Building Block）：为单个微批次（microbatch）设计计算流程。
* 重复（Repeating）：重复构建块以处理更多微批次，形成流水线。
* 压缩（Squeezing）：移除流水线中的冗余气泡。
* 重排序（Reordering）：可选步骤，优化预热（warm-up）和冷却（cool-down）阶段的吞吐量。
2. 揭示了构建块生命周期与峰值内存的直接关系：
* 通过公式 peak memory ≤ ⌈l/T⌉m（其中 l 是生命周期，T 是重复间隔，m 是单个微批次的激活内存），可以直接从构建块的设计计算出峰值内存，从而实现了对内存的可控设计。
3. 设计了一系列内存高效的V形（V-Shape）构建块：
* 针对现有方法中内存使用不均衡的问题（例如，1F1B中第一个设备内存压力最大），V形构建块通过将长生命周期的阶段与短生命周期的阶段搭配在同一设备上，实现了跨设备的内存负载均衡。
4. 基于V形构建块提出了三种新的流水线调度策略：
* V-Min：在吞吐量可比的情况下，将峰值激活内存渐进地降低到1F1B的1/3。
* V-Half：在吞吐量高于1F1B的情况下，将峰值激活内存渐进地降低到1F1B的1/2。
* V-ZB：在与1F1B相同的激活内存下，实现几乎零气泡的高吞吐量。
5. 实验验证：
* 在纯流水线并行设置中，新方法比1F1B的吞吐量高出7%到55%。
* 在混合并行（结合张量并行等）的实际场景中，通过超参数网格搜索，新方法比1F1B基线实现了16%的吞吐量提升。

该研究为流水线并行提供了一种系统化的设计和分析工具，并通过内存均衡的V形构建块，显著提升了内存效率和计算吞吐量，推进了大规模模型训练的帕累托前沿。

A3 如何构建一个流水线

我们提出了一个设计流水线调度的四步框架。

构建块（Building Block）：设计从为一个单一微批次（microbatch）规划计算过程开始，我们称之为构建块。例如，1F1B的构建块由一系列前向传播（forward passes）和随后的逆序反向传播（backward passes）组成。我们在图1a中用颜色高亮了1F1B的构建块。

重复（Repeating）：随后引入更多的微批次。构建块被重复并编织在一起形成一个流水线。在图1（上部）中，重复的构建块以不同深浅的灰色显示。值得注意的是，合法的构建块必须能够无冲突地重复，即来自两个构建块的计算过程不应相互重叠。

压缩（Squeezing）：根据构建块的不同，流水线中可能存在冗余的气泡，这些气泡可以通过压缩操作简单地移除，而无需改变计算过程的顺序。例如，图1b展示了一个通过压缩产生更高效流水线的案例。


图1: 一个流水线可以通过重复一个构建块，然后压缩冗余气泡来构建。

重排序（Reordering）（可选）：我们可以对预热（warm-up）和冷却（cool-down）阶段的计算过程进行重排序，以进一步提高计算吞吐量。直观地说，内存峰值发生在流水线的稳定阶段，而在预热和冷却阶段，RAM未被充分利用，这为在不改变峰值内存的情况下提高计算吞吐量留下了空间。我们在附录C中保留了细节。

大多数现有的流水线调度都可以在这个框架下得到解释。除了图1中展示的1F1B和eager 1F1B，我们还在一个更广泛的图集（见附录I）中展示了interleaved 1F1B【出处：Shoeybi et al., Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019】，ZB-H1【出处：Qi et al., Zero bubble pipeline parallelism, 2023】以及一系列知名的流水线。

2.1 构建块

构建块决定了计算和内存效率。各种流水线的计算和内存效率可以归因于它们的构建块。构建块的多样性主要来自三个因素：模型划分、设备放置以及计算过程之间的偏移量。我们沿用zero bubble PP【出处：Qi et al., Zero bubble pipeline parallelism, 2023】中的思想和符号，使用F表示前向传播（forward pass），B表示“为激活计算的反向传播”（backward for the activations），W表示“为权重计算的反向传播”（backward for the weights）。请注意，这种更细的粒度可以推广到像1F1B这样的先前方法，只需始终将B和W组合在一起。

模型划分。模型划分处理模型如何被分割成流水线阶段。最常见的模式是平均划分模型以匹配设备数量。一个突出的例子是1F1B调度（图1a）。这在interleaved 1F1B中得到了扩展，其中阶段的数量可以是设备数量的整数倍。

设备放置。设备放置是构建块设计中的另一个关键因素。虽然传统上每个流水线阶段按顺序放置在不同的设备上，但像在interleaved 1F1B（图18h）中那样将多个阶段放置在同一设备上也并不少见。另一个非常规设备放置的例子是Chimera，其中两个流水线以相反的设备顺序放置。

计算过程间的偏移量。最后但同样重要的是，F、B、W计算过程之间的偏移量在流水线的计算和内存效率中扮演着主要角色。通过简单地增大了1F1B构建块中后续F计算之间的偏移量，我们得到了eager 1F1B【出处：Zhuang et al., On optimizing the communication of model parallelism, 2023】（图1b），其中更多的F计算被急切地调度，导致更高的内存消耗（但有更好的通信重叠）。GPipe可以被看作是在1F1B构建块的最后一个F和第一个B之间增加了一个大的偏移量。关于偏移量影响的另一个例子是ZB-H1（图18c）和ZB-H2（图18d）调度的比较，可以看出，恰当选择的偏移量会产生像ZB-H2这样的零气泡调度。在这项工作中，我们假设每个F、B或W计算过程都占用一个单位的计算时间，并且只考虑整数单位的偏移量。虽然这可能会限制可行构建块的数量，但它极大地提高了分析的简便性。

2.2 计算峰值内存

通过构建块直接计算峰值内存。并非所有流水线生而平等，研究人员一直在寻找在计算和/或内存方面更高效的流水线。虽然通过枚举所有可能的构建块可以发现高效的流水线，但这无疑是成本高昂的。我们发现，一个流水线的峰值内存消耗可以通过其构建块，利用一个简单的公式来计算。这使我们能够设计具有可控峰值内存的流水线。

生命周期和重复间隔是关键。计算峰值内存需要两个关键量：一个阶段的生命周期（lifespan），以及构建块的重复间隔（repeating interval），这两者都在图1中进行了说明。一个阶段的生命周期定义为从F计算开始到B或W计算结束之间的时间量。一块激活内存（activation memory）在F计算开始时分配，并在整个生命周期内保留在RAM中，直到它被B和W两者消耗掉。峰值内存消耗可以通过找到其生命周期与其它所有微批次生命周期重叠的最大微批次数来计算。用l表示生命周期，T表示重复间隔，m表示单个微批次的激活内存大小，我们得到以下关系。

多阶段在单设备上的内存计算。当一个设备上有多个阶段时，例如interleaved 1F1B，它们对峰值内存的贡献是独立的。用Si表示分配给设备i的所有阶段，我们将每个阶段的贡献相加。
设备i的峰值内存 ≤ $ \sum_{s \in S_i} \lceil \frac{l_s}{T} \rceil m_s $

重复间隔的确定。另一个关键洞见是，重复间隔T可以从构建块中直接确定。在一个高效的流水线中，T应该等于构建块中每个阶段的计算单元数量。任何大于这个值的T都会在稳定阶段导致流水线气泡，而小于这个值的T则会导致冲突。一个微妙的例外是interleaved 1F1B，其重复间隔不统一。我们将讨论留到附录G。

2.3 无冲突重复

无冲突是构建块的合法性约束。在设计构建块时需要记住的一个约束是，一个合法的构建块必须能够无任何冲突地重复。如何设计满足这一约束的构建块可能看起来不直观。在实践中，我们首先设计构建块，然后进行事后验证。另一个有用的观察是，一个合法的构建块通常会在流水线中间产生一个稳定阶段，该阶段包含一个重复的d × T矩形，其中d是设备数量，T是重复间隔。这为约束构建块提供了另一种方法。我们可以从在这个矩形内排序计算过程开始，然后将其转换回一个构建块。

A2 方法细节

3 内存高效的构建块

现有流水线的内存效率问题。通过上述框架，我们可以方便地分析现有流水线的内存消耗模式。据我们所知，所有现有的流水线由于两个主要原因而内存效率低下：冗余的依赖链和不均衡的内存使用。在Zero Bubble【出处：Qi et al., Zero bubble pipeline parallelism, 2023】之前，反向传播通常被视为一个单一的过程，导致不必要的更长生命周期，从而产生更多的内存占用。在本文中，我们利用反向传播拆分策略（B和W）来消除这些冗余的生命周期。

不均衡内存的根源及V形构建块的提出。不均衡的内存源于各阶段生命周期的内在异构性。从图1a中，我们可以轻易看出各阶段的生命周期差异很大，第一个阶段的生命周期最长。因此，这在第一个设备上造成了内存瓶颈，并导致所有其他设备上的内存利用不足。为解决此问题，我们引入了一系列新颖的构建块，我们称之为V形（V-Shape）构建块。核心洞见来自公式1，该公式表明峰值内存取决于生命周期的总和。因此，当我们将多个阶段放置在同一设备上时，我们应该始终将长生命周期的阶段与短生命周期的阶段并置。当生命周期总和固定时，均衡的放置总是意味着更高的内存效率。这可以通过图2来证明，并行调度（用于interleaved 1F1B）是不均衡的，其内存瓶颈与l1 + l4成正比，而在V形调度中，它与l1 + l6成正比。

V形构建块的设计。V形调度要求将模型划分为设备数量两倍的阶段，并且后半部分阶段的设备放置顺序与前半部分相反。由于偏移量直接决定了每个阶段的生命周期，并因此通过公式1决定了峰值内存，我们可以进一步控制计算过程之间的偏移量，以生成具有不同内存特性的构建块。


图2: V形构建块确保所有设备上的峰值内存均衡，而并行构建块在第一个设备上存在内存瓶颈。

3.1 可控的均衡内存

统一偏移量简化设计。我们假设模型被均匀划分，即每个阶段的计算和内存都是相同的。对于单个微批次，我们用m表示每个阶段的激活内存，用M表示整个模型的总激活内存。注意$M=2dm$，其中d是设备数量。为了使其简单易处理，我们在F和B计算的前后半部分各自使用统一的偏移量来控制峰值内存。具体来说，我们在前d个阶段内的两个相邻F计算之间应用相同的偏移量$δ_0^F$（例如，图3b中$δ_0^F = 2$，图3c中$δ_0^F = 1$，图3d中$δ_0^F = 4$）。类似的约束也应用于F计算的另一半以及B计算的两半，分别表示为$δ_1^F$，$δ_0^B$，$δ_1^B$。为了保证设备间的峰值内存均衡，我们增加了两个约束条件，$δ_0^F = δ_1^B = δ_0$ 和 $δ_1^F = δ_0^B = δ_1$，为简化起见，我们使用符号$δ_0$和$δ_1$。例如，在图3d中，我们设置$δ_0 = 4$和$δ_1 = 2$。注意我们只控制不同设备间的偏移量。对于同一设备内的相邻计算（例如，最后一个阶段的F和B，最后一个设备中的两个F和两个B），我们使用暴力搜索来寻找最优解，确保它们的偏移量很小（小于重复间隔）。注意，在确定所有F和B计算后，W可以总是贪婪地放置，所以我们在暴力搜索期间不需要搜索它们的偏移量。


图3: 4个设备（d=4）下的V形构建块，其中白色文本颜色代表模型的前半部分阶段，黑色文本颜色代表后半部分。F、B和W分别代表前向传播、反向传播（为激活梯度）和为权重梯度的反向传播。

图4: V形调度与1F1B的比较，设置_为4个设备和8个微批次。稳定阶段遵循其构建块的模式。

通过偏移量控制峰值内存。根据公式1，我们可以分析关于d的渐近峰值内存。

忽略小常数后的直接控制。通过忽略小的常数，我们可以直接通过$δ_0$和$δ_1$的值来控制峰值内存。
表1: V-Min和V-Half的气泡和峰值内存忽略了小常数值。对于1F1B，δ0/δ1被重新定义为相邻前向/反向传播之间的偏移量。M代表整个模型的总激活内存，d是设备数量。

3.2 V形流水线调度

三种新颖的V形调度。通过改变$δ_0$和$δ_1$的值，我们提出了3种新颖的V形构建块（图3），并在图4中展示了它们基于我们框架的最终调度。V-Min的构建块（图3c）具有最小的偏移量，即$δ_0 = δ_1 = 1$，因此内存消耗最小。V-ZB通过设置$δ_0 = 4$和$δ_1 = 2$（如图3d所示），将气泡消除到几乎为0（图4d），从而实现了极致的吞吐量。V-Half的构建块（图3b）使用$δ_0 = 2$和$δ_1 = 1$，介于两个极端之间，消耗大约是1F1B所需激活内存的一半。尽管V-Min和V-Half的内存占用都低于1F1B，但假设F、B、W的运行时间相等，V-Min的气泡大约是1F1B的2/3，V-Half大约是1/2。我们在表1中展示了我们提出的V形调度与1F1B的比较。值得注意的是，V-Min的精确峰值内存是⌈ d+2 / 6 ⌉ M，V-Half是⌈ d+1 / 4 ⌉ M。为了避免V-Min和V-Half构建块中的冲突，（同一设备内的）偏移量对于不同的d值略有不同。详细信息在附录F中。

3.3 V-Min中的重复气泡

V-Min在真实场景中的问题。在F、B和W运行时间不同的真实场景中，V-Min会遭受重复气泡的影响。如图5所示，每个重复间隔T都存在气泡。因此，随着微批次数量的增加，气泡会增长。尽管V-Half也可能遇到同样的问题（当F、B和W的时间差异很大时），但由于其依赖关系较为松散，它在大多数经验案例中能生成很好地拼接的模式。如图5b所示，V-Half的吞吐量对运行时间的变化具有鲁棒性。此外，V-ZB的气泡在增加微批次数量时永远不会增长。我们将相关讨论留在附录E中。


图5: 我们从V-Min和V-Half调度中取一个重复的d × T网格，并为F/B/W分配不同的值。结果显示V-Min在每个重复网格中都有气泡，而V-Half则没有。

3.4 其他构建块

框架的泛化能力。除了V形构建块，我们还在附录H中提出了一些其他有趣的构建块，以展示我们框架的泛化能力。一些有用的例子包括：a) 1F1B-V，它在不做B-W拆分的情况下，实现了1F1B激活内存的2/3；b) 一个调度，其内存消耗比interleaved 1F1B少，但气泡率相同（图17c）。此外，我们在附录A中设计了一个自适应调度器，以更细的粒度控制内存。

A4 实验环境

模型架构：实验使用了一系列类似GPT-3【出处：Brown et al., Language models are few-shot learners, 2020】的模型，具体参数见表2。
表2: 实验中使用的模型。

硬件配置：
* GPU：最多40个NVIDIA A100 SXM 80G GPU。
* 节点：分布在5个节点上。
* 网络：通过RoCE RDMA网络互连。

软件配置：
* 代码实现：基于开源的Megatron-LM项目【出处：Narayanan et al., Efficient large-scale language model training on gpu clusters using megatron-lm, 2021】。
* 评估方法：在几次预热迭代后记录每次迭代的运行时间。

实验设置细节：
* 为了避免语言模型中的嵌入层和输出层成为流水线瓶颈并干扰效率评估，我们从初始和最终流水线阶段各减少了一个transformer层进行补偿，这与【出处：Qi et al., Zero bubble pipeline parallelism, 2023】中的设置类似。
* 对比方法：
* 本文方法：V-Min, V-Half, V-ZB。
* 基线方法：1F1B, Interleaved 1F1B (均来自Megatron-LM)。
* 其他对比方法：1F1B-R (带完全激活重计算的1F1B) 【出处：Chen et al., Training deep nets with sublinear memory cost, 2016】；ZB-1P, ZB-2P (自适应零气泡方法，内存限制分别为1F1B的1倍和2倍)【出处：Qi et al., Zero bubble pipeline parallelism, 2023】。

A4 实验结果

4.2 流水线调度对比

实验内容：在不同设置下，比较了各种流水线调度的吞吐量（以FLOPS利用率MFU衡量）和激活内存消耗。所有方法使用相同的微批次大小。

实验结果 (图6)：
* 吞吐量：V-ZB在所有方法中吞吐量最高，这与理论分析（图4）一致。V-Min的吞吐量与1F1B相当，但在微批次数量较大时落后于1F1B，这是由于3.3节讨论的重复气泡问题。
* 内存消耗：V-Min和V-Half显著减少了激活内存，分别降至约1/3和1/2。其他方法的内存消耗相似，除了1F1B-R（使用了完全重计算）。
* 结论：V-Min虽然存在重复气泡问题，但仍优于1F1B-R，为节省内存提供了有力的替代方案。


图6: 使用相同微批次大小的吞吐量和激活内存。

帕累托前沿分析 (图7)：
* 实验内容：在略有不同的设置下（为避免OOM，减小了9.6B和21B模型的微批次大小），绘制了各种方法的内存和MFU。
* 实验结果：V-Shape系列流水线调度位于帕累托前沿，表明它们在吞吐量和内存之间取得了最优的权衡。


图7: 不同设置下MFU和内存的帕累托前沿。

4.3 何时节省内存

实验内容1：与重计算对比
* V-Half和V-Min主要用于内存预算紧张的情况。传统方法是使用重计算（rematerialization），但这会降低吞吐量。
* 实验结果 (表7)：V-Half和V-Min的性能显著优于重计算方法(1F1B-R)。

实验内容2：利用节省的内存提升批大小
* 节省的内存可用于增加微批次大小，从而提高算术强度。
* 实验结果 (图8)：
* 在内存压力较大、微批次较小的大模型上，V-Half因其算术强度的增益，吞吐量可以超过V-ZB和其他基线。
* V-Min的算术强度增益不足以弥补其增加的气泡。
* 为V-Half/V-Min加倍/三倍微批次大小后，激活内存略高于其他方法，反映了理论分析中忽略的常数项。随着设备数量增加，这种增长不那么显著。


图8: 在相似内存限制下的吞吐量和激活内存。

4.4 与现有技术结合

实验内容：将本文方法与多种现有技术结合进行大语言模型训练，包括：Flash Attention, 张量并行(TP), 序列并行(SP), 以及Megatron-LM中的分布式优化器。在40个GPU上对PP、TP、DP大小和微批次大小进行网格搜索，并报告每种方法的最佳结果。

实验结果 (表3)：
* 低内存压力场景（序列长度较短）：V-ZB因其消除了气泡而表现最佳。
* 高内存压力场景（序列长度较长）：V-Half因其内存效率更高而表现最佳。
* 结论：在实际应用中，我们的方法与最先进技术结合依然表现出色，最佳策略的选择取决于具体的内存压力。
表3: V-Shape调度与其他内存节省方法结合。

A5 结论

工作总结：本文提出了一个通过重复构建块来构造流水线调度的框架。该框架能够直接从构建块的生命周期计算峰值内存。基于这一能力，我们设计了一系列内存高效的构建块。我们讨论了该系列中的三种代表性方法，即V-Min、V-Half和V-ZB，并通过实验证明我们的方法推进了大规模模型训练中吞吐量和内存的帕累托前沿。此外，我们通过构建块设计流水线调度的方法论可能会启发研究社区探索更多新颖的流水线调度。值得注意的是，重复构建块并非构建流水线的唯一方式，像贪心搜索等其他方法可能生成没有重复模式的流水线。

未来工作：未来，我们计划基于我们的框架进一步探索更多内存高效的流水线调度。V-Min的一个主要限制是，在增加微批次数量时会遭受不断增长的气泡。尽管V-Half缓解了这个问题，但仍有进一步降低内存消耗的空间。使用连续偏移量或更细粒度的离散化是解决这个问题的一种可能途径。我们将其留作未来的工作。

A6 附录

A 基于搜索的自适应调度器

更一般化的问题。现在我们考虑更一般化的场景，即在给定的激活内存限制下最小化气泡。一个直接的方法是简单地搜索所有可能的偏移量，并选择气泡最小的那个。然而，由于可能的偏移量数量呈指数级增长，这种朴素的方法不可行。在本节中，我们提出了一种更实际的搜索方法来解决这个普遍问题。

符号定义。我们使用上标$c \in \{0, 1\}$表示设备中的哪个阶段，使用下标$i \in \{1, 2, ..., d\}$表示设备索引。例如，$F_i^c$表示第i个设备中阶段$cd+i$的前向传播。我们定义从u到v的偏移量为$δ(u, v) = t(v) − t(u)$，其中$t(v)$表示计算v沿时间轴的单元格索引。为简化符号，我们定义$δ_{Fi}^0 = δ(F_i^0 , F_{i+1}^0)$, $δ_{Fi}^1 = δ(F_i^1 , F_{i-1}^1)$, $δ_{Bi}^1 = δ(B_i^1 , B_{i+1}^1)$和$δ_{Bi}^0 = δ(B_i^0 , B_{i-1}^0)$，来表示一个计算到其下一个计算的偏移量。

搜索空间与约束。我们不搜索所有可能的偏移量，而是将搜索空间限制在跨设备统一的偏移量上。我们还试图确保每个设备具有均衡的峰值内存。请注意，对于第3.1节中介绍的统一偏移量，峰值内存只落入一个小的离散集合中（{$\frac{k}{6}M$}，其中k为整数）。为了使其适用于更细粒度的内存控制，我们将重复模块分为两部分，分别包含前K行和后d-K行。更正式地，我们使用以下约束。

搜索过程与复杂度。请注意，上述约束具有良好的性质，即峰值内存在设备间是均衡的。由于我们总可以贪婪地填充$W_i^0$和$W_i^1$，我们只需要搜索第一个设备的排列、 $δ_{0K}$和K的值。如果我们将重复间隔视为常数，则计算复杂度为$O(d)$。

最终选择。对于搜索到的每个构建块，我们如2.1节所述重复构建块，检查冲突，进行