Synergistic Tensor and Pipeline Parallelism

作者/机构: Mengshi Qi∗1, Jiaxuan Peng1, Jie Zhang, Juan Zhu, Yong Li, Huadado Ma1 (1State Key Laboratory of Networking and Switching Technology, Beijing University of Posts and Telecommunications, China)

A1 主要贡献

本文旨在解决混合模型并行（张量并行TP与流水线并行PP结合）中存在的两个核心问题：TP引入的显著集体通信开销（TP气泡）和PP的同步低效问题（PP气泡）。现有工作通常孤立地解决这些问题，或专注于重叠TP通信，或专注于灵活的PP调度以减少PP气泡。本文提出了一种协同的张量和流水线并行调度方案，旨在同时减少这两种气泡。

核心问题与研究目标
* TP气泡：张量并行（TP）虽然能减少单个设备的内存需求，但会引入额外的集体通信操作（如All-Reduce）。随着TP规模的增大，这些通信开销（TP气泡）占比显著增加，严重影响训练性能。例如，在TP=8、序列长度为6144的配置下，TP气泡占总时间的27.5%。
* PP气泡：流水线并行（PP）虽然能有效利用计算资源，但各阶段之间需要仔细同步，以最小化设备空闲时间（PP气泡）。
* 研究目标：提出一种新的调度方法，从协同优化的角度出发，利用PP中丰富的计算操作来重叠（隐藏）TP气泡，并通过精确的调度来减少PP气泡，从而提升混合并行训练的整体吞吐量。

创新点与主要贡献
本文的主要贡献总结如下：
1. 提出协同并行的新视角和编织执行块：本文从协同优化的角度重新思考混合模型并行。将PP中的前向和后向传递解耦为细粒度的计算单元，然后将这些单元交错编织在一起，形成“编织执行块”（braided execution blocks）。这种结构可以在调度层面有效地消除TP通信气泡。
2. 设计新颖的“V”形PP调度：基于编织执行块，设计了一种新颖的PP调度。该调度采用“V”形数据流，实现了跨阶段的内存占用均衡，并显著降低了PP气泡率。
3. 广泛的实验验证：在多种规模的主流大语言模型（LLMs）和多模态大语言模型（MLLMs）上进行了广泛实验，验证了该方法的普适性和有效性。实验结果表明，与现有方法相比，该调度方案将LLM的训练吞吐量提升了高达12%，MLLM的训练吞吐量提升了高达16%。

此外，本文还提供了一个包含激活卸载（activation offloading）的增强版调度方案，作为在内存受限场景下的初步尝试。


图 1: 在Qwen2 Transformer层的前向传递中，PP内计算与TP通信重叠的加速效果。TP通信的比例随着TP规模的增加而显著增长，与朴素实现相比，我们的调度能有效重叠这部分开销。

A3 背景知识/关键Observation/设计原则

编织执行块

TP通信开销的来源与挑战。尽管张量并行（TP）能减少内存消耗，但它在Attention和MLP计算之后的前向和后向传递中引入了额外的通信开销，如图2中的橙色块所示。如图3a中的蓝色块所示，后向传递中的通信（All-Reduce）可以自然地与权重梯度计算重叠，实现并行执行。然而，前向传递中的通信（实心块）由于数据依赖性，在软件层面难以与其他操作重叠，导致大量的TP气泡。为了解决这个问题，我们设计了两种类型的执行块，通过编织（braiding）前向和后向计算单元来重叠TP通信与PP计算，如图3所示。这些执行块是我们调度方案的基础。


图 2: 单个TP Transformer层中的主要计算和通信操作示意图。f在前向中是恒等操作，在后向中是All-Reduce。g与f相反。

执行块的构建方式。具体来说，我们将整个Transformer层分解为几个细粒度的计算单元：Pre-Attn单元、Attn单元、Pre-MLP单元和MLP单元。然后，将来自前向和后向传递的单元按顺序交错，形成一个编织执行块，如图3a所示。此外，遵循Zero Bubble【【索引30，Zero bubble (almost) pipeline parallelism，2024，The Twelfth International Conference on Learning Representations】】的做法，后向传递中的Attn和MLP单元都被进一步拆分为两个部分：激活梯度计算和权重梯度计算。在必要时，我们会激活权重梯度计算的分离，以平衡工作负载并减少流水线气泡。这种分离不会破坏执行块的结构，因为后续的前向计算可以填充分离所引入的气泡，如图3b所示。Pre-Attn和Pre-MLP单元根据其计算依赖关系插入到计算流中。所有这些操作仅涉及模型级别的修改，使得我们的方法相比于内核级别的实现更加用户友好。


(b) 包含前向和激活后向的执行块。
图 3: 两种类型的执行块，它们通过编织前向（F）和后向（B）计算单元来重叠TP通信（All-Reduce），表示为AR。

计算融合修改。此外，考虑到Attn和MLP单元后的残差计算会引入额外的数据依赖，并增加后向传递的工程复杂性，我们修改了前向和后向计算，将残差计算融合到Attn和MLP单元中，然后再应用操作g。这种修改保持了计算的等价性，并且不影响模型收敛。对于单个TP rank，Attn单元修改后的前向计算可以表示为：

$$X_{ln} = LayerNorm(X), \quad X_{attn} = AR(Attention(X_{ln}) + \frac{detach(X)}{t})$$

其中，$t$表示张量并行大小，detach(·)是停止对$X$进行梯度计算的操作，Attention(·)是在没有任何通信的情况下执行的注意力计算，AR(·)代表All-Reduce通信操作。隐藏状态$X$的梯度计算$g(X)$则表示为：

$$g(X)=AR\left(\frac{\partial X_{attn}}{\partial X_{ln}}\right) \cdot \frac{\partial X_{ln}}{\partial X}+1,$$

其中，+1项考虑了来自残差连接的梯度贡献。类似地，MLP单元也采用了类似的计算融合策略。

构建原则

为了确保流水线调度的有效性，我们概述了几个旨在优化内存效率和最小化气泡的指导原则。

均衡的内存占用。正如【【索引29，Pipeline parallelism with controllable memory，2024，The Thirty-eighth Annual Conference on Neural Information Processing Systems】】和【【索引18，Hanayo: Harnessing wave-like pipeline parallelism for enhanced large model training efficiency，2023，Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】】中所讨论的，1F1B-I【【索引23，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，Proceedings of the international conference for high performance computing, networking, storage and analysis】】由于虚拟阶段的并行数据流，常常导致第一个设备出现内存瓶颈，而后续阶段的内存资源却未被充分利用。这种不平衡导致资源使用效率低下。为了解决这个问题，当前的调度方法【【索引18，Hanayo: Harnessing wave-like pipeline parallelism for enhanced large model training efficiency，2023，Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】、
【索引29，Pipeline parallelism with controllable memory，2024，The Thirty-eighth Annual Conference on Neural Information Processing Systems】、
【索引17，Deepseek-v3 technical report，2024，arXiv preprint arXiv:2412.19437】】利用“V”形调度策略，以实现所有设备上峰值内存的更均衡分布（见图4）。我们遵循这一原则，为我们的虚拟阶段实现了“V”形数据流，以优化内存效率，这也支持了编织执行块的调度，从而提高整体性能。


图 4: 微批次1的并行流水线流与“V”形流水线流的比较，后者显示出跨阶段内存平衡的改善。这归功于设备0上更早的后向传递。

气泡最小化。如第3节详述的编织式计算-通信执行块，展示了计算单元的排列方式。这些块的实际构建遵循两种主要模式：（1）将一个虚拟阶段的前向单元与另一个虚拟阶段的后向单元配对，或（2）将同一模型块的前向和后向单元在不同微批次间进行编织。构建模式的选择与PP中虚拟阶段的特性密切相关。具体来说，当将模型划分为虚拟阶段时，必须保持这些阶段的执行时间大致相等，以减少流水线气泡。对于多模态大语言模型（MLLM），视觉编码器（如ViT）和语言模型（LM）之间由于隐藏层大小的差异，通常存在显著的执行时间差距。这是多模态模型的一个关键挑战，其中较早的虚拟阶段通常包含更多的ViT层，而较晚的阶段则包含较少的LM层，以实现执行时间的平衡。然而，虚拟阶段间计算单元数量的不平衡限制了模式（1）的有效性，因为不匹配的单元数量无法在跨虚拟阶段配对时有效隐藏TP通信。相比之下，模式（2）涉及重叠同一模型块的前向和后向单元，能更有效地最小化TP气泡。因此，我们采用模式（2）来构建基础块，以保证通用性，特别是考虑到多模态模型的情况。

A2 方法细节

我们首先在4.1节介绍构建整体调度的几个关键原则，然后在4.2节介绍我们提出的调度方案。接着，在4.3节对PP和TP气泡及峰值内存进行理论分析。最后，在4.4节提出一个增强版的调度方案，以适应内存受限的场景。

协同张量和流水线并行调度

调度方案的整体结构。基于上述构建原则，我们构建了一个协同的张量和流水线调度方案，该方案为单个微批次实现了跨阶段的“V”形数据流，如图5所示。与现有的调度方案【【索引22，Pipedream: Generalized pipeline parallelism for dnn training，2019，Proceedings of the 27th ACM symposium on operating systems principles】、
【索引23，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，Proceedings of the international conference for high performance computing, networking, storage and analysis】、
【索引30，Zero bubble (almost) pipeline parallelism，2024，The Twelfth International Conference on Learning Representations】】类似，我们提出的方法包括三个主要阶段：预热（warm-up）阶段、稳定（steady）阶段和冷却（cool-down）阶段。首先，预热阶段在第一次后向传递开始前，以最大可能数量的在途（in-flight）微批次启动。这有助于保持较低的PP气泡率，并促进编织块的构建，因为较早微批次的后向传递与后续微批次的前向传递共同构成了执行块。参与预热阶段的微批次数量由阶段数决定。在此阶段，重叠的前向和后向（F&B）计算从第一个和第二个微批次的编织计算开始，并激活权重梯度分离，以便快速将梯度传播到下一阶段（最后一个阶段除外）。


图 5: 在4个设备和12个微批次的设置下，协同张量和流水线并行调度的示意图。深色和浅色块分别表示模型块0和1的计算。F、B和W分别代表前向、后向和权重梯度后向计算。

稳定阶段与冷却阶段的调度逻辑。经过几轮重叠的前向和权重（F&W）以及F&B计算后，调度进入稳定阶段。在此阶段，只要有新的微批次持续到达，权重分离就会被停用。在这种情况下，系统对模型块1执行一次F&B，然后对模型块0执行一次F&B。当微批次的供应耗尽时，维持完整的F&B变得不可行，过程会退化为一次完整的后向传递，然后是一次分离的F&B。在这个退化阶段，权重分离被重新激活，以使F&B的时间与完整后向传递的时间对齐。在冷却阶段，所有剩余在途微批次的后向传递都将完成。流水线气泡由之前步骤中存储的权重梯度计算来填充。此外，在稳定阶段激活权重分离会导致内存累积，因为相应的权重计算会一直挂起直到最后几个步骤。因此，权重梯度计算所需的已保存激活值会与计算流并行地卸载到CPU，并在需要时重新加载。

理论分析

参数定义与理论比较。我们定义PP阶段数和微批次数分别为$p$和$m$，其中$p \ll m$。为简化起见，此处不考虑多模态模型。对于LLM，我们假设模型沿层维度被均匀划分为多个块，并且每个微批次在所有设备上对每个模型块需要相同的激活内存$M_a$。一个块的前向传递计算时间表示为$T_F$，激活梯度和权重梯度的计算时间分别为$T_B$和$T_W$。一个模型块的TP通信时间表示为$T_{AR}$，该时间在前向和后向传递中保持一致。

表 1: 1F1B-I、Zero Bubble V 和我们提出的调度方案在理论PP气泡大小、未重叠的TP通信气泡和峰值激活内存方面的比较。

理论分析结果。如表1所示，具有2个虚拟阶段的1F1B-I【【索引23，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，Proceedings of the international conference for high performance computing, networking, storage and analysis】】在第一个设备上需要$(3*p-2)*M_a$的峰值激活内存，该值随着阶段索引的增加而递减。在后向传递中，TP通信被自然重叠，但前向传递中的TP通信是未重叠的，导致总TP通信时间为$2 * m * T_{AR}$。相比之下，Zero Bubble V (ZB-V)【【索引29，Pipeline parallelism with controllable memory，2024，The Thirty-eighth Annual Conference on Neural Information Processing Systems】】实现了更小的理论气泡大小$(p - 1) * (T_F + 2T_{AR} + T_B - 2T_W)$，并保持了比1F1B-I更低的峰值激活内存。然而，由于将后向传递解耦为激活和权重梯度计算，前向和激活后向传递中的TP通信都变得未重叠，导致总TP通信时间为$4 * m * T_{AR}$，并且实际气泡更大。我们的调度方案与ZB-V调度相比，表现出近似的气泡大小$(p - 1) * (T_F + T_{AR} + T_B - T_W)$，这远小于1F1B-I。并且，得益于计算与通信的重叠，它实现了大幅减少的TP开销，尽管这在峰值激活内存方面有所牺牲。

增强型流水线调度

内存优化动机与策略。考虑到实际应用中CUDA内存有限的需求，我们尝试降低我们调度策略的峰值内存占用。峰值内存使用率的升高主要源于两个阶段：预热阶段，此时在途微批次的数量最大化以维持低气泡率；以及稳定阶段，此时模型块0的激活值会持续存在较长时间才被释放，我们称之为生命周期。现有的方法如激活检查点（activation checkpoint）和卸载（offloading）【【索引40，Accelerating the training of large language models using efficient activation rematerialization and optimal hybrid parallelism，2024，2024 USENIX Annual Technical Conference (USENIX ATC 24)】、
【索引34，Adapipe: Optimizing pipeline parallelism with adaptive recomputation and partitioning，2024，Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】】已被用于减少内存。然而，检查点技术会在后向传递中引入额外的前向计算，增加了执行时间。因此，当主机和设备之间有高速带宽可用时，激活卸载在内存减少和端到端性能之间提供了更好的权衡。


图 6: 我们提出的带有卸载和重载的调度方案在预热和稳定阶段对单个微批次（PP size=2）的示意图。

分阶段卸载策略。在图6所示的预热阶段，内存使用量迅速累积，模型块1的激活值生命周期较短，因此不优先进行卸载。卸载时间$T_o$由卸载比例$\alpha$决定，并被限制为小于前向时间$T_F$，以避免干扰后续操作。在稳定阶段，模型块0的激活值比模型块1的激活值持续时间长得多，使其成为卸载的主要目标。由于编织执行块的存在，此阶段的卸载比例$\alpha$可以设置得比预热阶段更高。为了最小化频繁的双向CPU-GPU数据传输引起的PCIe带宽争用，我们避免卸载模型块1的激活值，确保它们可以及时重新加载而不会中断计算。比例$\alpha$根据硬件特性（如PCIe带宽和FLOPs）进行调整。

A4 实验环境与结果

实验环境

  • 模型：实验评估了Qwen2（LLM）【【索引39，Qwen2 technical report，2024，arXiv preprint arXiv:2407.10671】】和Qwen2-VL（MLLM）【【索引36，Qwen2-vl: Enhancing vision-language model’s perception of the world at any resolution，2024，arXiv preprint arXiv:2409.12191】】系列模型。所有模型都利用了Flash Attention 2【【索引3，Flashattention-2: Faster attention with better parallelism and work partitioning，2023，arXiv preprint arXiv:2307.08691】】来提升效率。模型具体配置见表2。
  • 硬件配置：实验在最多4个节点、共32个NVIDIA A800 SXM4 80G GPU上进行。部分实验在16个配备PCIe Gen 5互连的NVIDIA H20 GPU和16个NVIDIA H20 96GB GPU上进行。
  • 软件配置：实现基于开源的Megatron-Core项目【【索引33，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，arXiv preprint arXiv:1909.08053】】。
  • 实验设置：
    • LLM场景：遵循【【索引30，Zero bubble (almost) pipeline parallelism，2024，The Twelfth International Conference on Learning Representations】】中的方法，均匀切分模型，但考虑到Qwen模型152,064的大词汇表，最后一个阶段比其他阶段少两层。
    • MLLM场景：ViT编码器分配给设备0的第一个虚拟阶段，LM模型均匀分布在其余虚拟阶段，最后一个虚拟阶段同样比其他阶段少两层。
    • 对比方案：主要与以下流水线并行调度进行比较：a) Megatron-LM【【索引23，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，Proceedings of the international conference for high performance computing, networking, storage and analysis】】中实现的交错式1F1B（1F1B-I）；b) 【【索引29，Pipeline parallelism with controllable memory，2024，The Thirty-eighth Annual Conference on Neural Information Processing Systems】】中引入的ZB-V，它通过“V”形数据流实现均衡的峰值内存分布；c) 本文提出的调度。为保持一致性，所有调度方案均配置为每个设备两个虚拟阶段。

表 2: 实验中评估的模型配置。

实验结果

在LLM上与现有调度的比较

  • 实验内容：在2节点（16 GPU）和4节点（32 GPU）设置下，评估了不同流水线策略、序列长度和模型规模配置下Qwen2模型的训练吞吐量。

  • 实验结果与分析：如图7和图8所示，本文提出的调度方案在所有设置下的吞吐量均优于其他策略。

    • 在TP=8、序列长度为6144、PP=2的配置下，相比1F1B-I，吞吐量提升最高达到12.2%（图7右侧）。这得益于该配置下TP相关气泡占比较大，而本文方法能有效缓解这些气泡。
    • 在PP=8的配置下（图8），性能提升不太明显。这是因为本文调度中的显式流水线通信在计算后立即执行且无法与计算重叠，对性能产生负面影响。尽管如此，本文方法仍优于其他两种方法。
    • 在TP=8时，12B和26B LLM均获得最高吞吐量提升，表明该方法在扩展到更大模型（如TP=16）时具有巨大潜力。
    • ZB-V的性能与1F1B-I相当甚至更差，因为它解耦了后向传递，暴露了原本在1F1B-I中被重叠的TP气泡，引入了额外的气泡开销。

  • 内存占用：如图9所示，本文方案的峰值激活内存占用略高，但在内存不受限的场景下能实现最佳吞吐量。对于内存受限的场景，其增强版变体能显著降低内存占用，详见5.4节。


图 8: 在32个GPU上对26.3B LLM的实验结果（吞吐量）。

图 9: 在2个节点上对12.1B LLM测试的峰值激活内存占用。

图 7: 在16个GPU上对12.1B LLM的实验结果（吞吐量，单位：样本/秒）。

在MLLM上与现有调度的比较

  • 实验内容：在2节点和4节点设置下，评估了不同ViT/LM序列长度和模型规模配置下Qwen2-VL模型的吞吐量和峰值激活内存。
  • 实验结果与分析：如表3所示，本文的调度方法在所有计算均衡的场景中均表现出优越性能。
    • 在TP=8、PP=4的配置下，吞吐量提升最高达到11.7%，显著高于TP=4、PP=4时的2%提升，这与LLM实验的观察一致，即更大的TP规模通常带来更大的性能增益。
    • 一个特别显著的结果是在PP=2设置下（ViT计算强度相对较低）实现了16.7%的性能提升。
    • 在PP=8场景下（ViT的FLOPs超过LM块），当微批次数量增加时，本文方法展现出更优的扩展性和鲁棒性，相比ZB-V和1F1B-I获得了更大的提升。

表 3: 14.9B、28.8B和30.3B MLLM的实验结果，显示了吞吐量（样本/秒）和峰值激活内存占用（GB）。mbs表示微批次数。

增强型调度的效率

  • 实验内容：由于A800的PCIe带宽有限，在16个配备PCIe Gen 5互连的NVIDIA H20 GPU上对12.1B LLM（PP=4）进行实验，评估增强版调度（带卸载）的效率。
  • 实验结果与分析：如图10所示，
    • 吞吐量：带卸载的变体由于额外的数据传输CPU开销，导致吞吐量有微不足道的下降。
    • 内存占用：卸载策略实现了10%到19.2%的峰值内存降低，使得最终峰值内存占用接近1F1B-I（40G）。同时，所有阶段的内存占用近似相等，表明调度有效利用了内存资源来提高吞吐量。


图 10: 在4个PP阶段上对12.1B LLM的实验结果，包括吞吐量和峰值激活内存占用。*表示带卸载的增强版变体。

在最大化资源利用率下的结果

  • 实验内容：在16个NVIDIA H20 96GB GPU上，通过调整微批次大小来最大化内存利用率，并评估不同并行配置下的最优吞吐量和模型FLOPs利用率（MFU）。
  • 实验结果与分析：如表4所示，在最大化内存利用率下，本文方法表现出优越性能。
    • 在TP=2, PP=8配置下，尽管H20 GPU上TP气泡占比较低，本文的增强版调度仍实现了比1F1B-I略高的吞吐量（2.74样本/秒）和MFU（92.86%），同时保持了更低且更均衡的内存占用（68 GB）。

表 4: 在最大化内存利用率下，对12.1B LLM在16个H20 GPU上评估的实验结果。*表示我们带激活卸载的增强版调度。

A5 结论

本文提出了一种新颖的混合并行策略，该策略协同地集成了张量并行（TP）和流水线并行（PP），用于LLM和MLLM的分布式训练。通过将Transformer层分解为细粒度的计算单元，并在每个PP阶段内交错前向和后向计算，我们构建了能有效重叠TP通信气泡的执行块。基于这些执行块，我们设计了一个“V”形PP调度，该调度实现了跨阶段的内存均衡，并显著减少了TP和PP气泡，尽管代价是增加了峰值内存。实验结果表明，我们的方法在LLM和MLLM上均持续优于现有的调度方法。我们还提出了一个结合卸载功能的增强版变体，以更好地适应内存受限的场景。

未来工作：计划在未来的工作中，探索在大规模机器学习系统中更高效（内存和吞吐量方面）的混合调度方案。

A6 附录

A 流水线调度的预热阶段

预热阶段的构建原则。构建预热阶段的指导原则是确保重叠的前向和后向传递能够正确执行。如图11顶部所示，一个关键要求是前向和后向传递之间的重叠必须跨越不同的微批次。具体来说，前向传递中的微批次索引应大于后向传递中的微批次索引。为了满足这个条件，在重叠的F&B执行开始之前需要一个额外的前向传递，如图11(b)所示。然而，这种方法需要将后向传递解耦为激活和权重梯度计算，这会暴露TP通信并引入额外的PP通信开销。关于内存高效的预热阶段的更多细节，请参阅B节。


图 11: 预热阶段的构建：(a) 错误的预热阶段；(b) 内存高效的预热阶段；(c) 吞吐量高效的预热阶段。

B 与其他调度的比较

不同调度的气泡分析。我们在相同的配置下比较了1F1B-I【【索引23，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，Proceedings of the international conference for high performance computing, networking, storage and analysis】】、ZB-V【【索引29，Pipeline parallelism with controllable memory，2024，The Thirty-eighth Annual Conference on Neural Information Processing Systems】】和我们提出的调度，配置为4个流水线阶段，每个阶段2个虚拟阶段，以及12个微批次，如图12所示。直观上，ZB-V和我们的调度具有相似数量的PP气泡，都显著少于1F1B-I。1F1B-I中的PP气泡主要来自预热和冷却阶段。在预热阶段，跨虚拟阶段的并行数据流导致大量气泡。例如，微批次1的虚拟阶段2的生命周期很长，如果不引入更多微批次的额外前向传递，就无法有效填充，但这会增加阶段1的峰值内存。因此，1F1B-I在预热阶段会产生更大的PP气泡开销。同时，冷却阶段的气泡由于固有的数据依赖性是不可避免的。

我们调度方案的气泡特性。与ZB-V相比，我们的调度在预热阶段结束时会产生少量PP气泡，这归因于重叠F&B块相对较长的执行时间和特定的PP数据流模式。通常，激活梯度计算时间$T_B$满足$T_B > T_W$，其中$T_W$是权重梯度计算时间。因此，在冷却阶段结束时也会出现小气泡，其大小与ZB-V和我们提出的调度大致相等。


图 12: 我们工作中比较的同步流水线并行调度，设置了4个流水线阶段和12个微批次在一个训练迭代内。深色表示第一个虚拟阶段，浅色表示第二个虚拟阶段。ˆ表示以内存高效的预热阶段开始的张量和流水线调度。

不同变体的适用场景。此外，我们还展示了调度Oursˆ(d)，它以内存高效的预热阶段开始，如图12所示。尽管调度(d)的峰值内存占用低于我们的标准调度(c)，但它在结束时引入了额外的PP气泡，对吞吐量性能产生负面影响。因此，对于涉及大规模模型训练、具有大TP规模且内存充足的场景，我们更倾向于使用调度(c)以获得更高性能。我们的增强版调度变体，包括激活卸载，在假定有足够带宽的情况下，更适合用于降低峰值内存使用。调度(d)仅适用于内存极其有限且TP开销较大的场景。

C LLM的详细实验结果

详细数据展示。我们基于12.1B和26.3B的Qwen2模型提供了详细的实验结果，总结在表5、6和7中。值得注意的是，表5中报告的GPU内存占用略高于理论估计。我们将这种差异主要归因于我们工程实现的实际问题，例如细粒度计算单元的实现和权重梯度计算的解耦。例如，在TP=8下，分析每个虚拟阶段单个微批次的激活内存显示，ZB-V消耗3.6 GB，而我们的方法使用4.3 GB，增加了近20%。虽然这个问题对吞吐量没有影响，但我们计划在未来的工作中进一步完善工程实现以减少这部分开销。

表 5: 12.1B和26.3B LLM在峰值内存（GB）方面的实验结果，其中序列长度为3072和2048的样本的微批次大小设置为2。

表 6: 12.1B和26.3B LLM在吞吐量（样本/秒）方面的实验结果。mbs代表微批次的数量。

D 在H20 GPU上与其他调度方法的比较

H20 GPU上的性能表现。我们还在16个NVIDIA H20 96G GPU上进行了实验，以评估我们提出的调度与其他基线的性能，结果总结在表8中。我们提出的调度方案优于其他方法，但性能提升远不及在A800 GPU上实现的那么显著。这种差异主要归因于H20 GPU的特性，即其BF16 FLOPs较低但通信带宽较高。此外，我们分析了12.1B Qwen2模型中单个Transformer层的Attention和MLP模块的计算和通信时间，如图13所示。很明显，TP通信的比例显著低于在A800 GPU上测得的比例。因此，通过优化TP气泡获得的好处略有减弱。

表 7: 12.1B和26.3B LLM在模型FLOPs利用率（MFU, %）方面的实验结果。

表 8: 在16个H20 GPU上对12.1B LLM评估的实验结果。mbs和seq_len分别表示微批次数和序列长度。


图 13: 在H20 GPU上分析的Attention和MLP模块的计算和通信时间比例。

E 与现有训练策略的评估

E.1 激活检查点（Activation Checkpointing）

与AC的兼容性。我们提出的调度与激活检查点策略是兼容的。我们在Qwen2-12.1B上使用激活检查点（AC）进行了实验，批处理大小为128，序列长度为6k，结果报告在下表中：

表 9: 不同激活检查点配置下的性能比较。

AC实验结果分析。如表9所示，禁用AC可获得最高的吞吐量4.79样本/秒，但需要56.0 GB的峰值内存。在MLP模块上选择性启用AC，可将峰值内存降低20.2%至44.7 GB，相应地吞吐量下降12.5%至4.19样本/秒。将AC扩展到Attention和MLP模块，峰值内存进一步降低至41.5 GB（减少25.9%），而吞吐量降至3.94样本/秒。最激进的配置，即在所有模块上应用AC，实现了最大的内存节省，将峰值内存降低35.2%至36.3 GB，代价是吞吐量减少21.7%，最终为3.75样本/秒。这些结果证实了我们的调度框架与激活检查点策略完全兼容。

E.2 数据并行与上下文并行

与DP和CP的兼容性。为了证明我们的调度与数据并行（DP）和上下文并行（CP）的兼容性，我们在一个12.1B的Qwen2模型上进行了实验。如表10所示，在相同的张量并行（TP=2）和流水线并行（PP=4）配置下，我们的方法在两种设置中都达到了最高的吞吐量：使用CP时为2.71样本/秒（CP=2, mbs=128, 序列长度12k），使用DP时为9.40样本/秒（DP=2, mbs=256, 序列长度4k）。我们的方法在两种并行范式下都持续优于基线调度（1F1B-I和ZB-V），证实了其对不同分布式训练策略的鲁棒性和适应性。

F 通信重叠对GEMM执行效率的影响

背景。在大型语言模型的分布式训练中，将通信（如AllReduce）与计算（如GEMM）重叠，被广泛用于提高硬件利用率和减少训练时间。

表 10: 兼容数据并行和上下文并行的12.1B LLM实验结果。

资源竞争问题。然而，这种重叠可能导致资源竞争，特别是对于流式多处理器（SMs），这可能会降低GEMM的性能。

微基准测试。为了评估这种影响，我们在两种场景下进行了微基准测试：（1）GEMM主导AllReduce，可以实现完全重叠；（2）GEMM提前完成，导致部分通信暴露。在这两种情况下，我们将重叠执行与顺序基线（先GEMM后AllReduce）以及单独的GEMM时间进行了比较。

表 11: GEMM、AllReduce及其组合在顺序和重叠执行下的执行时间（毫秒）。

结果分析。结果总结在表11中，显示在第一种场景中，GEMM耗时8.605毫秒，AllReduce耗时3.364毫秒，重叠执行在9.251毫秒内完成，仅比单独的GEMM慢7.5%，比顺序执行（11.969毫秒）快22.6%。这表明通信被有效隐藏，干扰最小。

第二种场景分析。在第二种场景中，GEMM非常短（0.334毫秒），而AllReduce占主导（1.643毫秒），重叠执行在1.685毫秒内完成，这比AllReduce时间仅有2.6%的开销，并且仍然比顺序基线（1.977毫秒）快14.8%。值得注意的是，GEMM本身并没有被明显减慢；相反，总时间的轻微增加源于无法隐藏的通信尾部。