Penghui Qi∗, Xinyi Wan∗, Guangxing Huang & Min Lin Sea AI Lab {qiph,wanxy,huanggx,linmin}@sea.com

主要贡献

本文针对大规模分布式训练中的管道并行（PP）效率问题，焦点在于消除管道气泡，这些气泡被视为同步训练语义下不可避免的。研究目标是开发一种调度策略，实现零管道气泡，同时保持同步训练语义。创新点包括：将反向计算拆分为输入梯度和参数梯度两部分，以减少顺序依赖；基于此，手工设计新型管道调度，显著优于基线；进一步开发算法，根据特定模型配置和内存限制自动找到最优调度；引入技术绕过优化器步骤中的同步，实现真正零气泡。实验显示，该方法在相似内存限制下，吞吐量比1F1B调度高出23%，放松内存时可达31%。

背景知识/关键观察/设计原则

分布式模型训练背景：分布式模型训练已成为深度学习社区的焦点，尤其是大型复杂模型的出现。这些模型训练需要大量GPU，通过各种拓扑互联。过去几年提出了多种并行技术。数据并行（DP）【[Goyal et al., 2017, Accurate, large minibatch SGD: Training ImageNet in 1 hour, arXiv:1706.02677]；[Li et al., 2020, PyTorch Distributed: Experiences on Accelerating Data Parallel Training, arXiv:2006.15704]】是中小型模型的默认策略，因其简单性。但模型大小超出单个GPU时，模型并行成为解决方案【[Harlap et al., 2018, PipeDream: Fast and Efficient Pipeline Parallel DNN Training, arXiv:1806.03377]；[Huang et al., 2019, GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, NeurIPS 2019]；[Fan et al., 2021, DAPPLE: A Pipelined Data Parallel Approach for Training Large Models, PPoPP 2021]；[Zheng et al., 2022, Alpa: Automating Inter- and {Intra-Operator} Parallelism for Distributed Deep Learning, OSDI 2022]】。模型并行有两种主要方案：张量并行（TP）和管道并行（PP）。TP将一层中的矩阵乘法拆分到多个设备，而PP将整个模型分割成不同阶段，跨设备处理。ZeRO【[Rajbhandari et al., 2020, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, SC20]】通过跨设备分片参数，提供模型并行的强替代，同时保持DP的简单性。

大型训练中的并行策略交互：最近研究表明，在大规模训练中实现最优性能需要DP、TP和PP的非平凡交互。在丰富互连资源（如节点内NVLink）下，DP、TP和ZeRO的混合高效。而在跨服务器连接（尤其数千GPU规模）下，PP特别有利【[Fan et al., 2021, DAPPLE: A Pipelined Data Parallel Approach for Training Large Models, PPoPP 2021]；[Zheng et al., 2022, Alpa: Automating Inter- and {Intra-Operator} Parallelism for Distributed Deep Learning, OSDI 2022]；[Narayanan et al., 2021, Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC21]】。这突出了本文工作重点：提升PP效率。

PP效率的关键问题：PP实现的效率高度依赖设备空闲时间，即管道气泡。由于层间依赖，气泡似乎不可避免。早期工作GPipe【[Huang et al., 2019, GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, NeurIPS 2019]】通过增加管道中并发批次减少气泡比率，但导致峰值内存需求增加。为缓解此，GPipe丢弃部分中间激活，在反向传播时重计算，但引入约20%计算开销【[Fan et al., 2021, DAPPLE: A Pipelined Data Parallel Approach for Training Large Models, PPoPP 2021]】。异步PP改进包括PipeDream【[Harlap et al., 2018, PipeDream: Fast and Efficient Pipeline Parallel DNN Training, arXiv:1806.03377]】和PipeMare【[Yang et al., 2021, PipeMare: Asynchronous Pipeline Parallel DNN Training, MLSys 2021]】，理论上无气泡，提高效率，但牺牲精确优化语义。同步设置下改进包括1F1B调度，首先在PipeDream中提出用于异步【[Harlap et al., 2018, PipeDream: Fast and Efficient Pipeline Parallel DNN Training, arXiv:1806.03377]】，后引入同步【[Fan et al., 2021, DAPPLE: A Pipelined Data Parallel Approach for Training Large Models, PPoPP 2021]；[Narayanan et al., 2021, Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC21]】。1F1B通过早调度反向传播提供更快内存清除，在相同微批次下气泡比率类似，但峰值内存优势明显。基于1F1B，Narayanan et al.引入1F1B交织策略【[Narayanan et al., 2021, Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC21]】，通过将多个阶段分配到同一设备，进一步减少气泡，但增加通信和峰值内存。

当前挑战和机会：尽管努力，同步训练语义下剩余气泡仍是PP最大问题。本文发现，通过更细粒度表示和调度计算图，PP可进一步优化。经典深度学习框架以层粒度设计，而现代编译器使用不同中间表示进行多级优化【[Chen et al., 2018, TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, OSDI 2018]；[Roesch et al., 2018, Relay: A New IR for Machine Learning Frameworks, MAPL 2018]；[Sabne, 2020, XLA: Compiling Machine Learning for Peak Performance]；[Tillet et al., 2019, Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, MAPL 2019]；[Lattner et al., 2020, MLIR: A Compiler Infrastructure for the End of Moore’s Law, arXiv:2002.11054]】。更细粒度意味着更大搜索空间，但缺乏优化工具。选择合适粒度至关重要。

MLP的计算图
Figure 1: MLP的计算图。

传统粒度和关键观察：传统上，神经网络粒度化为堆叠层。每层关联两个函数：前向和反向。前向传递中，输入x通过参数化映射f(x, W)转换为输出y。反向传递涉及两个计算：∇xf (x, W )⊤ dℓdy 和 ∇W f (x, W )⊤ dℓdy，分别计算输入x和层参数W的梯度。为方便，使用B和W表示这些计算，F表示前向传递（Figure 1）。传统上，B和W分组为单个反向函数。这对用户友好，并在DP中有效，因为层i权重梯度通信可与层i-1反向计算重叠。但在PP中，此设计不必要增加顺序依赖计算，即层i-1的B依赖层i的W，通常有害于管道效率。

基于拆分的改进：基于拆分B和W，我们提出新管道调度，大幅提升效率。本文余下组织：第2节介绍基于F、B、W执行时间相同理想假设的手工调度。第3节移除假设，提出自动调度算法适用于现实条件。第4节详述绕过优化器步骤同步的方法，保持同步语义，实现零气泡。第5节在多样设置下评估方法与基线。第6节提出新型调度机制，进一步减少实现零气泡的内存需求，并评估性能。

工作范围：我们不探索大型分布式训练的通用混合策略，而是针对提升管道调度效率，提供与基线苹果对苹果比较。我们的方法正交于DP、TP和ZeRO，可作为大型训练中PP部分的并行替换。

方法细节

手工管道调度

关键观察基础：基于拆分B和W可减少顺序依赖从而提升效率的关键观察，我们从常用1F1B调度开始重新设计。如Figure 2所示，1F1B以热身阶段开始。在此阶段，工作者执行不同数量的前向传递，每个阶段通常比其后阶段多一个前向传递。热身后，每个工作者过渡到稳态，交替执行一个前向和一个反向，确保阶段间工作负载均匀分布。在最终阶段，每个工作者处理剩余在飞微批次的反向传递，完成批次。

1F1B管道调度
Figure 2: 1F1B管道调度。

改进版本：在我们改进版本中，将反向传递拆分为B和W传递，必须保持同一微批次的F和B在管道阶段间的顺序依赖。然而，W可在同一阶段相应B后灵活调度。这允许战略放置W填充管道气泡。有许多可能调度优于1F1B，在气泡大小和内存占用间权衡。我们在本节介绍两个特别有趣的手工调度，展示细粒度减少管道气泡的巨大潜力（见Figure 3）。为清晰起见，在初始设计中假设F、B、W时间成本相同，此假设为早期研究共享【[Narayanan et al., 2021, Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC21]；[Huang et al., 2019, GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, NeurIPS 2019]】。但在第3节，我们重新评估此假设，以优化现实场景调度效率。

手工管道调度，上：ZB-H1；下：ZB-H2
Figure 3: 手工管道调度，上：ZB-H1；下：ZB-H2。

ZB-H1调度：我们的第一个手工调度ZB-H1，确保所有工作者最大峰值内存使用不超过1F1B。ZB-H1大体遵循1F1B调度，但根据热身微批次数调整W起始点，确保所有工作者保持相同在飞微批次数。因此，如Figure 3（上）所示，气泡大小减少到1F1B的三分之一。此减少是因为所有工作者B比1F1B早启动，尾端气泡由晚启动W填充。由于W通常比B使用更少内存（Table 1），第一个工作者有最大峰值内存使用，与1F1B一致。

零气泡调度：当允许比1F1B更大内存占用且有足够微批次时，可能实现零气泡调度，我们标记为ZB-H2。如Figure 3（下）所示，我们在热身阶段引入更多F传递，填充初始B前的气泡。我们还重新排序尾部W传递，将布局从梯形变为平行四边形，消除所有气泡。重要的是，这里移除优化器步骤间的同步，我们在第4节讨论如何安全实现。

定量分析：我们用p表示阶段数，b表示每个微批大小。对于Transformer架构，用a表示注意力头数，s表示序列长度，h表示隐藏维度大小。我们用MB/MW表示存储一个B/W传递激活所需的内存，TF/TB/TW表示一个F/B/W传递的运行时间。为简单起见，仅对Transformer架构进行定量分析【[Vaswani et al., 2017, Attention Is All You Need, NeurIPS 2017]】，使用类似GPT-3的典型设置【[Brown et al., 2020, Language Models Are Few-Shot Learners, NeurIPS 2020]】，其中前馈隐藏维度为4h，每个注意力头维度为h/a。

FLOPs计算：如Narayanan et al.【[Narayanan et al., 2021, Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC21]】，仅考虑matmul操作计算FLOPs，因为它们贡献Transformer层大部分计算。对于前向传递中每个matmul，反向中有两个相同FLOPs的matmul，每个属于B或W。Transformer层FLOPs近似公式在Table 1。我们可见TW < TF < TB且TB + TW = 2TF。我们使用Korthikanti et al.方法估计B所需激活内存【[Korthikanti et al., 2023, Reducing Activation Recomputation in Large Transformer Models, MLSys 2023]】。B完成后，释放一些不再使用的激活，但保留一些额外梯度（Figure 1中的∇zL）用于W。W所需总内存，如Table 1所示，小于B。

Table 1: 每个Transformer层每个传递的FLOPs和激活内存

ZB-H1和ZB-H2量化：不假设TF = TB = TW时，ZB-H1和ZB-H2的峰值激活内存和气泡大小量化在Table 2。值得注意，工作者i的激活内存为(p − i + 1)MB + (i − 1)MW用于ZB-H1，(2p − 2i + 1)MB + (2i − 2)MW用于ZB-H2。如Table 1，W所需激活内存小于B。因此，峰值激活内存为pMB用于ZB-H1，(2p − 1)MB用于ZB-H2。

Table 2: 1F1B与我们手工调度比较

自动管道调度

手工调度的实际问题：虽然手工调度提供简单性和更好可理解性，但实际应用中面临几个问题。首先，在TF = TB = TW假设下调度引入不想要的气泡，尤其当这些值显著不同时。其次，阶段间传输激活/梯度的通信时间（记为Tcomm）常被忽略，导致管道流中明显延迟。最后，在可用内存不足以容纳足够微批次实现无气泡调度时，平衡最小化气泡大小和遵守内存限制特别具挑战性。

自动搜索算法：为解决这些挑战并确保泛化到实际场景，我们提出算法，给定管道阶段数p、微批次数m、激活内存限制Mlimit和运行时间估计TF、TB、TW、Tcomm，自动搜索最优调度。我们设计启发式策略，总产生最优或近最优解，尤其当m足够大时。我们还系统制定问题为整数线性规划（详见附录G），可用现成ILP求解器解决【[Forrest & Lougee-Heimer, 2005, CBC User Guide, INFORMS 2005]】当问题在一定规模下。这两种方法可结合：先用启发式解初始化，然后用ILP进一步优化。

启发式算法步骤：
- 在热身阶段，在内存限制内，调度尽可能多的F传递，以最小化第一个B前的气泡。结果调度可能仍有小于TF的小气泡，如果未达内存限制，调度另一个F可能延迟后续B。我们用二元超参数控制是否执行。
- 热身后，坚持一个F和一个B迭代调度的模式。当有大于TW的间隙时插入W填充气泡。当气泡发生但大小小于TW时，如果当前气泡使所有阶段最大累计气泡大小增大，仍插入W。当命中内存限制时，也插入W回收内存。通常，我们启发式策略进入遵循1F-1B-1W模式的稳态。
- 在整个过程中，管道阶段i总保证在F用尽前调度至少比阶段i+1多一个F。当此差异超过一，我们用另一个二元超参数决定是否在阶段i跳过一个F，如果不引起更多气泡。我们执行网格搜索找到最佳超参数组合。
- 在每个阶段，当F和B传递用尽时，我们逐一调度所有剩余W传递。

绕过优化器同步

当前实践中的同步：在大多数PP实践中，优化器步骤中通常跨管道阶段执行同步，以数值鲁棒性。例如，计算全局梯度范数用于梯度裁剪【[Pascanu et al., 2013, On the Difficulty of Training Recurrent Neural Networks, ICML 2013]】；在混合精度设置中执行NAN和INF全局检查【[Micikevicius et al., 2017, Mixed Precision Training, arXiv:1710.03740]】；两者都需要跨所有阶段的全归约通信。然而，优化器步骤的同步破坏平行四边形（Figure 3），使零气泡不可能。本节提出替代机制绕过这些同步，同时保持同步优化语义。

现有实现和观察：现有实现中，先启动全归约通信收集全局状态，然后是条件于全局状态的优化器步骤。然而，我们注意到大多数时候全局状态无影响，例如，NAN和INF全局检查很少触发，因为在鲁棒设置中大多数迭代无数值问题；梯度裁剪率经验上较低，不需每迭代同步全局梯度范数。

后更新验证机制：基于这些观察，我们提出用后更新验证替换事先同步。想法如Figure 4所示，在每个阶段优化器步骤前，从前一阶段接收部分归约全局状态，与当前阶段本地状态结合，传递到下一阶段。每个阶段优化器步骤由部分归约状态控制，例如，当发现NAN或部分归约梯度范数超过裁剪阈值时跳过更新。在下迭代热身阶段，全归约全局状态从最后阶段传播回第一阶段。接收全局状态后，每个阶段执行验证决定先前优化器步骤是否合法。如果需要修正梯度，将发出回滚（详见附录C），然后基于全归约全局状态重做优化器步骤。

替换优化器同步的后验证策略
Figure 4: 替换优化器同步的后验证策略。

内存高效零气泡调度

ZB-2p的内存问题：虽然ZB-2p有效实现近零气泡，但成本是比1F1B双倍内存消耗。此增加内存需求限制其实用性。为解决此，我们设计ZB-V，一种在与1F1B相同内存约束下实现最小空闲时间的调度方法。受Narayanan et al.交织1F1B策略启发【[Narayanan et al., 2021, Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC21]】，我们的方法将整个模型均匀分为正好2p块，每工作者分配两个块。与交织方案不同，我们的方法从第一工作者开始顺序分配模型块，到最后，然后从最后逆序回第一，创建独特的“V”形状（见Figure 8中第一个微批次的前向传递）。例如，在4阶段管道中分割16层Transformer模型，我们分配层1-2和层15-16给工作者1，层3-4和层13-14给工作者2，依此类推。

ZB-V调度。每个设备分配正好2块，白文字颜色代表第一块，黑文字颜色代表第二块。模型块间依赖序列遵循“V”形状模式，用于前向和反向传递
Figure 8: ZB-V调度。每个设备分配正好2块，白文字颜色代表第一块，黑文字颜色代表第二块。模型块间依赖序列遵循“V”形状模式，用于前向和反向传递。

优势：此方法确保每个微批次的前向和反向传递从同一工作者起始，不同于先前方法如1F1B和交织1F1B，其中前向从第一工作者开始，反向从最后开始。此区别提供两个显著优势：首先，第一工作者可立即启动反向传递，无需等待最后工作者反向返回，导致更快内存清除和减少实现最小空闲时间的内存需求。在TF = TB = TW条件下，ZB-V以pMB峰值激活内存实现零气泡，与1F1B最大峰值内存一致。值得注意，这几乎是ZB-H2所需(2p − 1)MB的一半。其次，峰值内存使用在所有工作者间固有平衡。此平衡由于所有模型块的统一计算负载和一致内存消耗。

ZB-V调度策略：在Figure 8中，ZB-V调度策略分三个阶段展开。在初始热身阶段，每个工作者（记为i）执行总共2p − 1个前向传递，包括2p − i个第一块传递和i − 1个第二块传递。热身后，所有工作者过渡到稳态，特征为重复1F-1B-1W模式。在稳态中，工作者执行计算组F-B-W，每个组对应特定块。对于给定工作者i，过程以执行p − i个第二块组开始。随后，工作者在第二块和第一块间交替处理一个组。此模式继续直到所有前向传递处理。在最终阶段，每个工作者专注于处理剩余B和W计算，优先B，W填充气泡。

自动搜索：我们采用与第3.1节类似启发式算法，自动搜索最优调度，考虑参数如管道阶段数p、微批次数m、激活内存限制Mlimit和剖析运行时间TF、TB、TW、Tcomm。由于热身和稳态中内存分布在所有工作者间固有平衡，我们可简单地将所有W右移，在内存约束内。此修改启用有效利用额外W填充调度尾部气泡，主要由于W比F和B持续时间短（详见附录D）。

实验环境

数据集：使用类似GPT-3的模型进行评估，未指定具体数据集，但实验基于语言模型训练设置，序列长度为2048，全球批次大小为512（见Table 3）。

模型架构关键参数：评估模型类似GPT-3，包括1.5B、6.2B、14.6B和28.3B参数模型。每个模型有特定层数、隐藏大小、注意力头数（见Table 3）。初始和最终管道阶段比中间阶段少一层Transformer层，以补偿额外嵌入查找和损失计算，避免成为瓶颈。

硬件配置：使用最多32个NVIDIA A100 SXM 80G GPU，分布在4个节点，通过RoCE RDMA网络互联。GPU数量包括8、16、32，跨节点设置。

软件配置：基于开源Megatron-LM项目实现【[Narayanan et al., 2021, Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC21]】。使用PyTorch分布式。固定随机种子初始化模型，验证ZB-1p、ZB-2p与1F1B的比特级相同损失。

实验结果

主要结果实验：比较ZB-1p（内存限pMB，与1F1B峰值内存相同）、ZB-2p（内存限2pMB，实现近零气泡）、1F1B和1F1B-I的吞吐量（Figure 5，Table 4）。ZB-2p在各种设置下一致优于其他方法，吞吐量与微批次数无关，因为气泡率近零（Table 5），接近上界（上界粗略估计为1F1B吞吐量乘以1/(1-1F1B气泡率)）。ZB-2p在相同内存下比1F1B高（附录F）。ZB-1p在8 GPU设置下与1F1B-I相当，在多节点下优于1F1B-I，因减少气泡无额外通信成本。m≤p情况见附录H，改进20%-30%。

不同管道调度吞吐量比较
Figure 5: 不同管道调度吞吐量比较。

Table 4: 实验结果细节

自动调度效率实验：使用气泡率=(cost − m(TF + TB + TW ))/cost评估（cost为所有阶段最大执行时间）。Table 5显示ZB-2p气泡率<1%，优于ZB-H2，证明自动算法更好适应现实TF、TB、TW、Tcomm估计。ZB-1p未显著优于ZB-H1，可能因内存限主导。Figure 6显示ZB-2p调度及剖析执行，几乎无气泡。

Table 5: 不同设置下1F1B、1F1B-I、ZB-H1、ZB-H2、ZB-1p、ZB-2p的气泡率

ZB-2p产生的调度（上）和其剖析执行过程（下）
Figure 6: ZB-2p产生的调度（上）和其剖析执行过程（下）。

内存限制实验：Figure 7显示内存限与气泡率关系，初始线性下降，理论上在(p−1)(TB+2Tcomm)+pTF/TF MB附近平台。经验上，2pMB是实现近零气泡的好阈值。超出拐点，零气泡成本高于收益（详见附录B）。

使用启发式算法的内存限与气泡率关系
Figure 7: 使用启发式算法的内存限与气泡率关系。Figure 7证实，在RoCE下，2pMB内存限实现近零气泡。

ZB-V评估实验：Table 6比较1F1B、ZB-1p、ZB-2p（半微批大小，双微批次数，相同内存）和ZB-V。ZB-V一致优于1F1B和ZB-1p，与ZB-2p相当。Table 7消融显示增大微批大小改进8%（14.6B和28.3B模型），但6.2B模型<3%，解释ZB-2p*在某些场景优于ZB-V。存在大微批大小与低气泡率的权衡。

Table 6: 相同内存下1F1B、ZB-1p、ZB-2p和ZB-V比较

Table 7: 双倍每个微批大小时的改进

ZB-V调度效率实验：Table 8显示ZB-V气泡率显著小于1F1B、1F1B-I、ZB-H1，与ZB-H2相当但内存仅一半。Figure 9显示内存限与气泡率关系，低于2pMB时ZB-V优于非ZB-V启发式。

Table 8: 不同设置下1F1B、1F1B-I、ZB-H1、ZB-H2和ZB-V的气泡率

ZB-V的内存限与气泡率关系，与第3.1节启发式方法比较
Figure 9: ZB-V的内存限与气泡率关系，与第3.1节启发式方法比较。

结论

本文引入新型策略，通过拆分反向计算中的激活梯度和参数梯度，提升管道并行效率，并设计自动管道调度算法，在不同内存预算下最小化气泡率。该算法产生的调度一致优于1F1B，甚至实现近零气泡率。为进一步减少内存，我们提出名为ZB-V的新调度机制，在TF = TB = TW时实现零气泡，同时遵守与1F1B相同内存限。我们的方法另一优势是可用较小微批次数（通常3p足够）实现最优效率，意味着更多微批可分区到数据并行维度，提升大型模型训练的可扩展性。未来工作可探索与DP、TP和ZeRO的更深集成。

附录

附录A：考虑数据并行时，在优化器步骤前启动全归约通信收集梯度。通常，此通信与计算传递重叠差，导致延迟，尤其通信带宽有限时。如Figure 3，通常迭代尾部调度多个W传递。每个W传递包括几个独立计算，为不同参数计算梯度。如Figure 10，我们可重新排序这些计算，聚类计算相同参数梯度的那些，从而实现计算与通信的最优重叠。

按W分组的原始调度与差重叠（上）和按参数分组的重排序调度与最优重叠（下）比较。数字i表示计算属于第i W，不同颜色代表不同参数的计算
Figure 10: 按W分组的原始调度与差重叠（上）和按参数分组的重排序调度与最优重叠（下）比较。

附录B：内存限与气泡率关系高度受初始阶段第一个B前气泡影响。对于第一个微批，前向从初始阶段到最终，反向反转直到回初始。第一个微批从开始到完成总时间至少p(TF + TB) + 2(p − 1)Tcomm，无法压缩因依赖链。记第一个B前F传递数为k(≥1)，气泡大小为β(≥0)，则有公式(见原文图像)。内存限下界与k成比例(公式1)，β与k反比例(公式2)。增加k且k < ⌊ (p−1)(TB+2Tcomm)+pTF/T ⌋时，β线性减少，内存限下界线性增加。当k = ⌊ (p−1)(TB+2Tcomm)+pTF/T ⌋时，β达最小而不延迟B，其值< TF，峰值激活内存至少⌊ (p−1)(TB+2Tcomm)+pTF/T ⌋MB。超出此，进一步减至零气泡不易。因每个阶段有< TF小气泡，调度另一个F将延迟B起始，导致前阶段更多F需求。理论上，初始阶段需额外p-1 F传递完全消除所有阶段第一个B前气泡(见Figure 11)，意味着总激活内存至少⌊ (p−1)(TB+2Tcomm)+(2p−1)TF/T ⌋MB。

32微批1.5B模型的零气泡调度
Figure 11: 32微批1.5B模型的零气泡调度。

附录C：回滚优化器步骤时，典型方法存储参数和优化器状态历史版本，需要时回退。但此内存低效，需要多复制操作，伤害训练性能。对于大多数优化器，我们注意到步骤函数算术可逆。在此观察下，我们提出新型技术执行原地优化器回滚，避免分配额外内存，仅在回滚时需额外计算。如算法1，我们展示如何回滚AdamW优化器步骤【[Loshchilov & Hutter, 2017, Decoupled Weight Decay Regularization, arXiv:1711.05101]】。

算法1: AdamW的原地回滚

附录D：对于第5节实验，我们记录ZB-2p在不同设置下的剖析时间TF、TB、TW和Tcomm。这些值用于计算第5.3和5.4节所有方法的气泡率。见Table 9。

Table 9: TF、TB、TW和Tcomm的剖析时间

附录E：本节提供优化器后验证策略有效性的消融研究。比较ZB-2p在后验证和全归约同步下的吞吐量。根据Table 10实验结果，同步版本比后验证ZB-2p性能下降约8%。

Table 10: ZB-2p与同步ZB-2p的吞吐量（样本/GPU/秒）比较

附录F：在相同内存消耗下，我们为1F1B和ZB-1p双倍每个微批大小，并与ZB-2p比较吞吐量(Table 11)。结果显示ZB-2p即使半微批大小也优于1F1B。经验上，大批大小提高GPU利用率和效率。但对于大模型，隐藏维度足够饱和设备利用率。基于此和结果，我们认为ZB-2p优于为1F1B增大批大小。在某些设备利用不饱和且m/p较大的实验中，双倍微批ZB-1p可能略优于ZB-2p。

Table 11: 相同内存下1F1B、ZB-1p和ZB-2p比较

附录G：管道中任何传递可唯一索引为(i, j, c)，i∈{1,...,p}为阶段，j∈{1,...,m}为微批，c∈{F,B,W}为传递类型。定义T(i,j,c)为时间成本，E(i,j,c)为结束时间。引入∆M(i,j,c)为传递(i,j,c)引起的内存增量。例如，∆M(·,·,F)=MB因前向导致净增MB激活存储用于反向。∆M(·,·,B)=MW - MB移除B存储内存但添加W所需，∆M(·,·,W)=-MW。我们搜索的变量是调度中传递顺序，引入O(i,j,c)→(i,j′,c′)∈{0,1}指示(i,j,c)是否在(i,j′,c′)前调度。优化目标(3)是最小化最长阶段时间。约束(4)(5)添加同一微批相邻阶段F和B的顺序依赖。(6)添加调度顺序约束。(7)限制峰值激活内存低于Mlimit。(见原文图像公式)。

附录H：PP性质上当微批次数m < 阶段数p时，有大气泡率。但零气泡方法仍可在这些罕见设置下提升20%-30%。忽略通信假设m<=p且TW<TB，1F1B迭代需(m + p − 1) ∗ (TF + TB + TW)完成，而ZB迭代需(m + p − 1) ∗ (TF + TB) + TW。实验结果见Table 12。注意当m<=p时ZB-1p和ZB-2p本质相同，内存与1F1B类似。

Table 12: 小微批次数下1F1B与ZB-2p比较

方法细节中的引用汇总

  • [Goyal et al., 2017]: Accurate, large minibatch SGD: Training ImageNet in 1 hour, arXiv:1706.02677. 在Introduction中描述为数据并行的默认策略。
  • [Li et al., 2020]: PyTorch Distributed: Experiences on Accelerating Data Parallel Training, arXiv:2006.15704. 在Introduction中描述为数据并行策略。
  • [Harlap et al., 2018]: PipeDream: Fast and Efficient Pipeline Parallel DNN Training, arXiv:1806.03377. 在Introduction中描述为模型并行和异步PP工作；在Handcrafted Schedules中描述为提出1F1B；在Experiments中作为比较方法。
  • [Huang et al., 2019]: GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, NeurIPS 2019. 在Introduction中描述为早期PP工作；在Handcrafted Schedules中共享时间假设。
  • [Fan et al., 2021]: DAPPLE: A Pipelined Data Parallel Approach for Training Large Models, PPoPP 2021. 在Introduction中描述为模型并行和PP优势证据；在GPipe开销中引用。
  • [Zheng et al., 2022]: Alpa: Automating Inter- and {Intra-Operator} Parallelism for Distributed Deep Learning, OSDI 2022. 在Introduction中描述为模型并行和PP优势证据。
  • [Rajbhandari et al., 2020]: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, SC20. 在Introduction中描述为模型并行的替代。
  • [Yang et al., 2021]: PipeMare: Asynchronous Pipeline Parallel DNN Training, MLSys 2021. 在Introduction中描述为异步PP改进。
  • [Narayanan et al., 2021]: Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC21. 在Introduction中描述为PP优势证据、1F1B交织；在Handcrafted Schedules中共享时间假设；在Quantitative Analyses中引用FLOPs计算；在Experiments中作为实现基础和比较方法；在ZB-V中作为交织策略启发。
  • [Chen et al., 2018]: TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, OSDI 2018. 在Introduction中描述为深度学习编译器优化。
  • [Roesch et al., 2018]: Relay: A New IR for Machine Learning Frameworks, MAPL 2018. 在Introduction中描述为编译器IR。
  • [Sabne, 2020]: XLA: Compiling Machine Learning for Peak Performance. 在Introduction中描述为编译器优化。
  • [Tillet et al., 2019]: Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, MAPL 2019. 在Introduction中描述为编译器优化。
  • [Lattner et al., 2020]: MLIR: A Compiler Infrastructure for the End of Moore’s Law, arXiv:2002.11054. 在Introduction中描述为编译器基础设施。
  • [Vaswani et al., 2017]: Attention Is All You Need, NeurIPS 2017. 在Quantitative Analyses中描述为Transformer架构。
  • [Brown et al., 2020]: Language Models Are Few-Shot Learners, NeurIPS 2020. 在Quantitative Analyses中描述为GPT-3设置；在Experiments中作为模型基础。
  • [Korthikanti et al., 2023]: Reducing Activation Recomputation in Large Transformer Models, MLSys 2023. 在Quantitative Analyses中描述为激活内存估计方法。
  • [Forrest & Lougee-Heimer, 2005]: CBC User Guide, INFORMS 2005. 在Automatic Scheduling中描述为ILP求解器。
  • [Pascanu et al., 2013]: On the Difficulty of Training Recurrent Neural Networks, ICML 2013. 在Bypassing Optimizer中描述为梯度裁剪。
  • [Micikevicius et al., 2017]: Mixed Precision Training, arXiv:1710.03740. 在Bypassing Optimizer中描述为混合精度设置。
  • [Loshchilov & Hutter, 2017]: Decoupled Weight Decay Regularization, arXiv:1711.05101. 在附录C中描述为AdamW优化器。