文章标题：Mario: 流水线并行中近乎零开销的激活检查点
作者/机构：Weijian Liu, Mingzhen Li, Guangming Tan, Weile Jia (中国科学院计算技术研究所，SKLP；中国科学院大学)

A1 主要贡献

本文提出了 Mario，一个流水线优化器，旨在将激活检查点（activation checkpointing）技术以近乎零开销的方式集成到现有的流水线并行方案中，以解决大规模语言模型（LLM）训练中的内存瓶颈和效率问题。

  • 核心问题：现有的流水线并行方案虽然提高了设备利用率，但通常会导致额外且不均衡的激活内存占用。直接应用激活检查点技术虽然能解决内存问题，但会引入显著的重计算开销（GPipe 报道约为 23%），从而降低训练效率。

  • 研究目标：设计一种优化器，能够将激活检查点的重计算开销隐藏在流水线气泡（pipeline bubbles）中，从而实现内存占用减少和平衡，同时不牺牲甚至提升训练吞吐量。

  • 主要创新点：

    1. 设计了一个流水线优化器 Mario：它通过将激活检查点技术“镶嵌”到流水线中，有效减少并平衡了各设备的内存占用。具体而言，Mario 通过提前执行重计算指令，将其开销隐藏在流水线固有的以及特意重塑（reshaped）的气泡中。
    2. 提出了一个精确的基于模拟器的性能模型：该模型能即时提供性能反馈，结合了动态规划和预定义依赖关系来精确估计执行行为，避免了昂贵的真实硬件运行。
    3. 利用检查点节省的内存扩大了流水线调度空间：通过激活检查点释放的内存，Mario 能够支持更大的微批次大小（micro-batch size）以提高设备利用率，或支持更长的序列长度（sequence length）以提升模型处理上下文的能力。此外，本文还引入了一个自动调度器，用于寻找更优的参数配置。
    4. 在 Megatron-DeepSpeed 框架中实现了 Mario：并在 GPT3 和 LLaMA2 等 Transformer 模型上进行了评估。实验结果表明，Mario 能显著提升现有先进流水线方案（如 1F1B、Chimera、Interleave）的训练速度。

A3 背景知识与动机

2.1 流水线并行

并行化方案简介：为了训练大型模型和处理大规模数据集，采用了多种并行化方案。数据并行（DP）在每个设备上保存完整的模型权重副本，内存效率不高。ZeRO 【35, ZeRO: memory optimizations toward training trillion parameter models, 2020, SC; 36, ZeRO-infinity: breaking the GPU memory wall for extreme scale deep learning, 2021, SC】通过分区参数、梯度和优化器状态改进了 DP。张量并行（TP），包括序列并行（SP），每个设备计算模型算子的一部分，但在跨节点训练时通信开销大。流水线并行（PP）将模型层划分为多个阶段，每个设备负责一个或多个阶段，以流水线方式处理微批次。PP 与 DP 和 TP 是正交的，可以组合使用以获得更好的资源利用率【3, AutoDDL: Automatic Distributed Deep Learning With NearOptimal Bandwidth Cost, 2024, IEEE Transactions on Parallel and Distributed Systems; 16, Whale: Efficient Giant Model Training over Heterogeneous GPUs, 2022, USENIX ATC; 49, Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning, 2022, OSDI】。

Table 1. 各流水线方案的内存占用。

Table 2. 符号表。

Table 3. 指令表。

同步流水线并行：本文专注于同步 PP 方法。这些方法会触发梯度同步和流水线刷新，以确保训练迭代中的所有微批次使用相同的权重，但这会引入流水线气泡【23, Chimera: efficiently training large-scale neural networks with bidirectional pipelines, 2021, SC】。异步方法则存在权重版本不匹配【45, PipeMare: Asynchronous Pipeline Parallel DNN Training, 2021, MLSys】或权重陈旧【29, PipeDream: generalized pipeline parallelism for DNN training, 2019, SOSP; 30, Memory-Efficient Pipeline-Parallel DNN Training, 2021, ICML】的问题，两者都会损害收敛性【2, Demystifying Parallel and Distributed Deep Learning: An In-depth Concurrency Analysis, 2019, ACM Comput. Surv.; 5, Why Dataset Properties Bound the Scalability of Parallel Machine Learning Training Algorithms, 2021, IEEE Transactions on Parallel and Distributed Systems】。


Figure 1. 流水线并行方案的发展。

现有 PP 方案及其权衡：图 1 总结了典型的 PP 方案及其对训练吞吐量的改进。现有的 PP 工作致力于减少流水线气泡和提高计算效率，但通常以牺牲内存效率为代价。表 1 总结了各种流水线方案的权重内存和激活内存。在权重内存方面，Chimera 维护两个模型权重副本，而其他方案只维护一个。峰值激活内存由正在处理的微批次数量决定，因此初始流水线阶段通常比最终阶段消耗更多内存。峰值激活内存与参与流水线的设备索引呈线性相关。如表 1 最后一列所示，使用 Mario 后，峰值激活内存可以减少到 $A$ 或 $A/2$，使其与设备数量 $D$ 无关，且成本接近于零。

2.2 激活检查点与重计算

激活检查点的基本原理：激活检查点通过增加额外的计算来换取更低的内存占用【4, Training deep nets with sublinear memory cost, 2016, arXiv; 11, MegTaiChi: dynamic tensor-based memory management optimization for DNN training, 2022, ICS; 21, Dynamic Tensor Rematerialization, 2021, ICLR; 26, Exploiting Input Tensor Dynamics in Activation Checkpointing for Efficient Training on GPU, 2023, IPDPS】。它允许在前向计算中丢弃激活张量，并在后向计算中通过重放前向计算（即重计算）来重新生成它们。由于重计算的开销通常低于 GPU 和 CPU 之间的数据传输开销，激活检查点在主流深度学习框架中被广泛采用。

激活检查点在流水线中的应用现状：一些工作已将检查点应用于训练大型模型，但仅将其视为弥补流水线方案引入的内存占用的补救措施。GPipe 【13, GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, 2019, NeurIPS】使用检查点来减少多个并发微批次产生的多版本激活张量，并考虑了提前调度重计算以与气泡重叠。Megatron-LM 也以朴素【31, Efficient large-scale language model training on GPU clusters using megatron-LM, 2021, SC】和选择性【22, Reducing activation recomputation in large transformer models, 2023, MLSys】的方式使用检查点来减少峰值内存。AdaPipe 【39, AdaPipe: Optimizing Pipeline Parallelism with Adaptive Recomputation and Partitioning, 2024, ASPLOS】为特定算子选择检查点，然后重新平衡流水线阶段。PipeDream-2bw 【30, Memory-Efficient Pipeline-Parallel DNN Training, 2021, ICML】保存了两个版本的模型权重，因此需要检查点来弥补额外的内存占用。

未被充分挖掘的潜力：然而，将检查点与流水线并行相结合的潜力仍未得到充分探索。通过消除激活内存，流水线执行可以以更灵活的方式进行调整，但先前的工作忽略了在这种情况下存在的加速机会。

3.1 流水线中高效检查点的需求

激活检查点的开销问题：为了消除激活内存占用与设备数量之间的线性关系，当前的流水线方案依赖于激活检查点，但代价是重计算开销，GPipe【13, GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, 2019, NeurIPS】报告该开销占单次训练迭代的 23%。理论上，一个块的重计算时间等于其前向时间，导致约 33% 的开销（假设后向时间是前向时间的两倍，记为 $T$）。图 2 展示了一个 GPT3-125M 模型在 4 阶段 1F1B 方案下的执行情况，比较了无检查点的基线与带检查点的四个步骤。在步骤 1 中，当应用检查点并将阶段 $i$ 的重计算置于其后向计算之前时，总计算时间从 21$T$（基线）增加到 28$T$。


Figure 2. 一个 4 阶段的 1F1B 流水线可以实现近乎零开销的激活检查点。X 轴表示时间线，y 轴表示流水线阶段。红色/绿色框突出显示了不同步骤中的优化。

实现近乎零开销检查点的可能性：然而，如果将重计算精心“镶嵌”到流水线中，激活检查点可以做到近乎零开销。具体步骤如下：
- 步骤 2：提前开始重计算，将其填充到流水线气泡中，使总时间减少到 25$T$。
- 步骤 3：移除冗余的重计算，特别是那些紧邻其相应后向计算的重计算，将时间进一步减少到 23$T$。
- 步骤 4：采用细粒度的协同调度，例如提前执行额外的微批次来重塑流水线气泡，将时间减少到 22$T$，接近基线的 21$T$。
在 4 个 A100-40G GPU 节点上复现上述步骤，证明了近乎零开销的激活检查点是可行的。这可以实现：1) 减少内存占用；2) 缓解设备间激活内存不平衡的问题；3) 扩大流水线并行的可行调度空间。

3.2 流水线性能建模的复杂性

现有建模方法的局限性：为流水线调度（特别是与激活检查点结合时）建立精确的性能模型非常复杂。简单地计算前向和后向计算网格数量的方法过于简化，无法捕捉真实场景。原因包括：
1. 阶段划分不均衡：流水线不仅包含相同的 Transformer 层，还包含其他各种层（如第一个/最后一个阶段的嵌入层和层归一化层）。
2. 前向与后向计算时间比非整数：前向和后向计算的执行时间并非简单的整数比（如通常假设的 1:2）。例如，根据参考文献【22, Reducing activation recomputation in large transformer models, 2023, MLSys】，一个 Transformer 层的比例约为 1:1.6。基于这些假设的偏差性能数据会导致次优的流水线建模。
3. 框架实现紧耦合：当前的流水线方案与定制框架紧密集成（例如，DeepSpeed 基于 num_in_flight 公式调度指令，Chimera 使用队列，Hanayo 使用未开源的动作列表），这使得在不进行实际运行的情况下跨多种方案评估检查点策略变得困难。

真实运行的成本高昂：然而，通过真实运行来评估这些流水线调度既昂贵又耗时。例如，仅在 2048 个 GPU 上初始化 Megatron-LM 框架就需要 1047 秒【18, MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs, 2024, NSDI】。因此，一个可靠且可移植的性能建模工具（如流水线模拟器）至关重要。这样的工具能够以轻量级和公平的方式估计流水线方案的效率，从而在无需大量实际测试的情况下促进进一步的参数调优。

4. 概述

目标：Mario 旨在通过集成激活检查点来优化流水线调度，从而以近乎零的成本减少和平衡设备间的内存占用。因此，Mario 能够使用更大的微批次大小来提高计算效率，以及更长的序列长度来提高模型质量。Mario 专注于 PP，并且只优化同步流水线调度，使其与 DP/TP/SP 方案正交。经 Mario 优化的流水线调度可以无缝集成到其他并行策略中。

用户接口：如代码清单 1 所示，用户需要指定 Mario 配置和模型配置。Mario 配置包括流水线方案、全局批量大小、设备数量和每个设备的内存。目前，Mario 支持典型的流水线方案，包括 "V" 形（1F1B）、"X" 形（Chimera）和 "W" 形（Interleave）。用户也可以选择“Auto”选项来自动确定最优的流水线方案。

指令列表作为中间表示：一旦指定，流水线方案会为每个设备扩展成一个训练迭代的指令列表，其功能类似于流水线的中间表示（IR）。如图 3(b) 所示，指令列表主要由基本的前向（$FW_i^p$）和后向（$BW_i^p$）指令组成。每个指令有两个关键属性：微批次 ID（下标 $i$）和分区 ID（上标 $p$）。模型配置被实例化为一个可训练模型，然后被划分为多个流水线阶段，并映射到可用设备上。

设计概览：如图 3 所示，Mario 包含三个主要组件：图调优器（graph tuner）、基于模拟器的性能模型和调度调优器（schedule tuner）。
- 图调优器：将激活检查点和重计算“镶嵌”到流水线中，识别模式以在流水线气泡内重叠重计算，并重塑这些气泡以实现进一步的重叠。
- 基于模拟器的性能模型：通过预测每个块的延迟和内存使用来评估调优后流水线的效率。它通过动态规划算法对所有可能的流水线方案的虚拟流水线进行建模，考虑了水平/垂直依赖关系。
- 调度调优器：通过迭代方法进一步优化图调优器生成的调度。它利用轻量级模拟器快速估计兼容调度的吞吐量，从而避免在 GPU 集群上进行昂贵的执行。


Figure 3. Mario 概览：(a) 图调优器，(b) 基于模拟器的性能模型，和 (c) 调度调优器。

# Listing 1. Mario 的编程接口。
import mario
mario_conf = {
    'pipeline_scheme': 'Auto|V|X|W|...',
    'global_batch_size': 128,
    'num_device': 32,
    'memory_per_device':'40G'
}
model_conf = {
    'type': 'GPT3|LLaMA2|...',
    'hidden_size': 4096, ...
}
schedule = mario.optimize(mario_conf, model_conf, global_bs)
mario.run(schedule)

A2 方法细节

5.1 图调优器

基于指令列表构建计算图：我们基于初始的流水线调度（即指令列表）构建流水线执行的计算图。指令列表类似于流水线的中间表示（IR），每个指令代表计算/通信类型及其依赖关系。指令的依赖关系表明指令 I 必须在指令集 S 完成后执行，这些依赖可分为水平和垂直两种。我们专注于前向（$FW$）和后向（$BW$）指令，并在此基础上插入辅助指令（如 SEND_act, RECV_act, SEND_grad, RECV_grad, RE）来完善流水线执行过程。表 3 总结了常用指令。为了简化，我们假设各阶段延迟均衡且后向延迟是前向的两倍，但模拟器使用真实延迟以保证高保真度。

图级别优化遍（Passes）：Mario 应用图级别的优化遍来修改和转换计算图，确保每次变换后计算图的所有依赖关系都得以维持。我们使用 1F1B 方案的指令列表来演示这些优化遍，这些遍同样适用于 Chimera 和 Interleave 等其他方案。这些优化遍是迭代应用的。

  • Pass 1: apply-checkpoint（应用检查点）：此遍为所有成对的前向和后向指令应用激活检查点。前向指令（$FW_i$）被替换为带检查点的前向指令（$FW_{ckpt,i}$），并在相应的后向指令（$BW_i$）之前插入一个重计算指令（$RE_i$），以确保后向指令所需的激活张量得到恢复。$RE_i$ 和 $BW_i$ 之间的距离应最小化，因为 $RE_i$ 生成的激活会一直保存在内存中直到被 $BW_i$ 消耗。较短的距离意味着峰值内存占用时间更短，从而允许调度其他高开销指令。因此，每个阶段只保留一份激活副本，从而平衡了设备间的内存占用。

  • Pass 2: overlap-recompute（重叠重计算）：此遍通过共同提前重计算和后向指令，将重计算指令重叠在流水线气泡中。重计算指令 $RE_i$ 仅依赖于同一设备上的前向指令 $FW_{ckpt,i}$，因此可以放置在 $FW_{ckpt,i}$ 和 $BW_i$ 之间的空间，该空间通常包含流水线气泡。后向指令 $BW_i$ 依赖于同一设备（设备 $d$）的 $RE_i$ 和下一设备 $d+1$ 的 $BW_{i+1}$。因此，设备 $d$ 的 $RE_i$ 可以与设备 $d+1$ 的 $BW_{i+1}$ 并发执行，有效地将其隐藏在流水线气泡中。为了保持流水线效率，$RE_i$ 必须放置在接收来自设备 $d+1$ 的 $BW_{i+1}$ 梯度的指令（RECV_grad）之前。如果 $RE_i$ 被错误地放置在 RECV_grad 之后，它必须等待 RECV_grad 完成，而 RECV_grad 又依赖于设备 $d+1$ 的 $BW_{i+1}$，这会导致设备 $d$ 上的 $RE_i$ 等待设备 $d+1$ 上的 $BW_{i+1}$，从而失去并发执行的机会。

  • Pass 3: remove-redundancy（移除冗余）：此遍移除冗余的激活检查点。当 $FW_{ckpt,i}$ 和 $RE_i$ 相邻时，对应的激活刚被丢弃就立即被恢复，这导致了与没有检查点时相似的更高内存占用。因此，这种情况下的检查点是冗余的，应被移除。

  • Pass 4: prepose-forward（前置前向计算）：此遍将带检查点的前向指令提前到最早的流水线气泡中。当带检查点的前向指令（例如 $FW_{ckpt,i}$）之前存在流水线气泡时，它可以被提前到气泡中，使其原始位置变为空闲。这使得 Pass 2 可以将其他重计算（例如 $RE_{i+1}$）移动到这个空闲位置，从而增加了未来 $RE_{i+j}$（$j>1$）指令被重叠的机会。在没有检查点的传统流水线中，虽然也可以提前前向指令，但这会因为增加在途（on-the-fly）微批次数量而增加内存占用，使得这种方法不可行。具体来说，与 $FW_{ckpt,i}$ 对应的 SEND_act 指令也需要被考虑以避免通信死锁。这可以分两种情况处理：1) 如果下一设备的 $FW_{ckpt,i}$ 也被提前（例如 $FW_{4|5}$），SEND_act 可以与 $FW_{ckpt,i}$ 一同提前。2) 如果没有（例如 $FW_{6|7}$），SEND_act 应保留在原位。$FW_{ckpt,i}$ 的输出被临时存储在缓冲区中，SEND_act 在发送给下一设备的 $FW_{ckpt,i}$ 之前从缓冲区读取数据。这是必要的，因为 Mario 支持阻塞式点对点通信方法，SEND_act 和 RECV_act 必须配对以避免死锁。

5.2 基于模拟器的性能模型

构建统一的性能模型：为了理解流水线的计算图，我们构建了一个基于模拟器的性能模型。与常见的分析性能模型【23, Chimera: efficiently training large-scale neural networks with bidirectional pipelines, 2021, SC; 49, Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning, 2022, OSDI】不同，我们将估算过程表述为动态规划算法，无需手动识别关键路径或对理想流水线进行近似。对于 GPT3-13B（64 个微批次，Chimera 方案）在 32 个 GPU 上的训练，模拟延迟约为 700 毫秒，显著加速了调度调优器。

虚拟流水线（Virtual Pipeline）：我们引入虚拟流水线的概念来统一表示各种流水线方案。尽管这些方案在指令执行顺序上差异很大，但它们都遵循基本原则：$FW$ 指令按序遍历所有阶段，然后 $BW$ 指令以相反顺序执行。虚拟流水线旨在封装这一原则。我们通过引入分区 ID（part_id，上标 $p$）来区分单个设备上的不同阶段。结合微批次 ID（micro_id，下标 $i$），虚拟流水线可以抽象地表示几乎所有流水线方案。例如，Chimera 中的双向流水线（向上和向下）分别被分配 part_id 0 和 1。

依赖关系查找：如算法 1 所示，虚拟流水线中指令间的依赖关系被简化为两个函数：find_prev_inst（定位前一阶段的指令）和 find_next_inst（定位下一阶段的指令）。find_prev_inst 函数对于 $FW$ 指令的逻辑方向与 $BW$ 指令相反。在 1F1B 方案中，目标指令的设备索引沿逻辑方向前进一步。对于 Chimera，上/下行流水线分别遵循逻辑/相反方向。对于 Interleave，设备索引使用模运算在逻辑方向上进行，如果跨越阶段则调整 part_id。此外，我们提供了一个灵活的接口供用户扩展以支持其他新兴流水线方案。

算法 1 在虚拟流水线中查找依赖关系。

轻量级性能剖析（Lightweight Profiling）：为了支持模拟器，我们通过离线剖析收集计算时间、内存占用和通信时间。剖析过程是轻量级的，只收集十个训练迭代的数据（例如，剖析 LLaMA2-13B 只需 142 秒）。我们遵循以下准则：1) 将 Transformer 块作为剖析的基本单元。2) 选择 1F1B 方案中的第 (D-1) 个设备进行剖析，因为它通常包含多个 Transformer 块且有更多可用内存。3) 将内存占用分为静态部分和动态部分。4) 关注预热、稳定和冷却阶段的整体模式，而非单个发送/接收指令。我们应用线性回归（$y = ax + b$）来预测执行时间和静态/动态内存，其中 $x$ 是 Transformer 块的数量，$b$ 代表框架开销。点对点通信时间的预测也采用类似方法，其中 $x$ 是微批次的数量。

动态规划公式化：为了推导流水线指令的时间线，我们开发了一个基于上述估计器的动态规划算法。关键是保留指令间的依赖关系，并逐步推断每个指令的最早开始时间。指令的开始时间由其前驱指令（跨设备的垂直依赖和设备内的水平依赖）的最晚完成时间决定。


Figure 4. V|X|W 形流水线的模拟过程。

模拟过程：图 4 展示了详细的时间线模拟过程，包含五个主要步骤：1) 使用预训练的估计器为每个指令分配延迟。2) 初始化没有前驱的 $FW$ 指令。3) 更新 $FW$ 指令的开始时间。4) 识别仅以其对应 $FW$ 为前驱的 $BW$ 指令并标记。5) 更新 $BW$ 指令的开始时间。总计算时间由所有设备上最后一个 $BW$ 指令的完成时间得出。内存模拟在设备级别进行，我们累积估计的静态内存并跟踪峰值动态内存。

可视化：Mario 支持计算图的可视化，允许用户直观地观察流水线执行状态、气泡分布等，而不是仅仅依赖于执行时间和吞吐量。这种可视化有助于对带检查点的流水线进行细粒度调整。图 5 展示了一些可视化示例。


Figure 5. 通过 Mario 模拟器进行的流水线可视化。

与 TP 和 DP 的集成：我们将 TP 视为流水线阶段内的内部块，其带来的内存、通信时间和延迟变化直接影响指令，并将在剖析中被捕获。对于 DP，我们只对流水线冷却阶段的 allreduce 时间进行建模，无需沿 DP 维度复制流水线模拟。

5.3 调度调优器

参数调优问题：调度调优器基于 Mario 配置和模型配置优化参数。它利用图调优器和基于模拟器的性能模型。由于 Mario 通过释放大部分激活张量显著减少了内存占用，我们可以增加微批次大小以更好地利用可用内存。我们将其公式化为一个优化问题，如下面的方程 1 所示。待调优的参数包括：1) $C$：是否启用 Mario 的激活检查点；2) $S$：选择哪种流水线方案；3) $PP_{dim}$：流水线并行维度；4) $DP_{dim}$：数据并行维度；5) $mbs$：微批次大小；6) $Mem_{avail}$：可用设备内存。主要目标是最大化由性能模型 $M$ 和估计器 $P$ 估算的训练吞吐量。

优化问题公式

max T = M(P, C, S, PP_dim, DP_dim, mbs)
s.t.
C ∈ {True, False}
S ∈ {1F1B, Chimera, Interleave}
4 ≤ PP_dim ≤ D
PP_dim × DP_dim = D
mbs ∈ {1, 2, 4, 8, ...}
Mem_avail: 用户指定的可用设备内存
D: 可用设备总数

求解方法：为了支持调度调优器，我们扩展了模拟器 $M$，当估计的峰值内存超过 $Mem_{avail}$ 时，$M$ 返回零作为惩罚。我们还扩展了 $P$ 以支持 $DP_{dim}$ 参数，通过乘以一个效率系数来模拟数据并行的可扩展性。我们采用网格搜索来解决这个优化问题，因为搜索空间是可管理的，并且得益于模拟器的效率和准确性，每次迭代都是轻量级的。

A4 实验

实验环境

  • 模型：GPT3（1.6B 和 13B 参数）和 LLaMA2（3B 和 13B 参数）。模型配置总结在表 4 中。
  • 硬件：一个包含 16 个节点的集群，每个节点配备 2 个 64 核 ARMv8 CPU（鲲鹏 920）和 4 个 NVIDIA A100-40G GPU。
  • 软件：基于 Megatron-DeepSpeed (v0.2.0) 和 DeepSpeed (v0.11.2) 实现。驱动版本 510.85.02, CUDA 11.8, cuDNN 8.4。
  • 基线方案：1F1B (V), Chimera (X), Interleave (W)。
  • 评估配置：
    1. base (baseline)：原始流水线方案，无额外优化。
    2. ckpt (checkpoint)：启用了激活检查点的流水线方案。
    3. ovlp (overlap)：ckpt 配置经过 Mario 的四遍优化。
    4. lmbs (larger micro-batch size)：在 ovlp 配置上进一步增大微批次大小，同时保持全局批量大小不变。

Table 4. 实验中使用的 LLM 模型。

实验结果

6.1 在 GPT3-1.6B 和 LLaMA2-3B 上的性能
- 实验设置：使用 8 个 A100-40G GPU 的流水线。全局批量大小为 128。
- 结果：如图 6 所示，lmbs 配置实现了最高的吞吐量。与 base 相比，lmbs 在 V/X/W 方案上平均取得了 1.25 倍的加速。通过激活检查点，峰值内存占用减少到 base 的 1/2 到 1/3。节省的内存允许我们将微批次大小加倍，提高了计算效率并减少了 p2p 通信。与 base 相比，ovlp 平均慢了 29.22%，这是因为在小模型下，重计算无法被完全隐藏，但随着模型或微批次大小增加，此开销可被消除（见表 5）。Interleave(W)方案由于内存限制，微批次大小设为 V/X 的一半，导致吞吐量较低。


Figure 6. 在 8 个 A100-40G GPU 上 GPT3-1.6B 和 LLaMA2-3B 的性能。

6.2 在 GPT3-13B 和 LLaMA2-13B 上的性能
- 实验设置：使用 32 个 A100-40G GPU 的流水线。
- 结果：如表 5 所示，由于激活内存不平衡，基线方案会出现内存不足（OOM）。与 ckpt 相比，ovlp 和 lmbs 分别实现了 1.13 倍和 1.36 倍的加速。在 LLaMA2-13B 上，V-ovlp 达到了 V-base 吞吐量的 94.7%，仅有约 5% 的开销，验证了 Mario 的“近乎零开销”特性。

Table 5. 在 32 个 A100-40G GPU 上 GPT3-13B 和 LLaMA2-13B 的性能。带下划线的值由 Mario 模拟器估计，原始情况为 OOM。

6.3 峰值内存占用
- 结果：如图 7 所示，激活检查点显著减少并平衡了设备间的内存使用。ovlp 相较于 ckpt 没有引入额外的内存占用。lmbs 虽然通过增加微批次大小使用了更多内存，但设备间的内存占用仍比 base 更加均衡。这证明 Mario 能有效解决 SOTA 流水线方案中高且不均衡的内存使用问题。


Figure 7. 各设备上的峰值内存占用。

6.4 模型参数扩展
- 实验设置：在 16 个 A100-40G GPU 上对 GPT3 模型进行参数扩展实验。
- 结果：如图 8 所示，Mario 显著增加了可训练参数的规模。V-base 可处理 3B 参数，而 V-lmbs 可处理 16B（提升 5.3 倍）。X-lmbs 可处理 7B（提升 2.3 倍）。W-lmbs 可处理的参数规模是 W-base 的约 20 倍。


Figure 8. 在 16 个 A100-40G GPU 上对 GPT3 模型进行模型参数扩展。

6.5 序列长度扩展
- 实验设置：在 16 个 A100-40G GPU 上使用 GPT3-1.6B 模型进行序列长度扩展实验。
- 结果：如图 9 所示，经 Mario 优化的 PP:8 TP:2 配置支持的序列长度比未优化的 PP:8 TP:2 平均长 1.49 倍，比 PP:8 TP:1 平均长 2.80 倍。这是因为 Mario 节省的大量激活内存可用于容纳更长的序列。


Figure 9. 在 16 个 A100-40G GPU 上对 GPT3-1.6B 进行序列长度扩展。

6.6 模拟器准确性
- 实验设置：在 8 个 A100-40G GPU 上使用 GPT3-1.6B 评估模拟器准确性。
- 结果：如图 10 所示，估计的峰值内存与实际测量值非常接近，平均绝对百分比误差（MAPE）仅为 5.1%。吞吐量估计的 MAPE 为 9.4%，虽然略有高估，但重要的是，估计吞吐量的偏序关系与实际运行一致，足以用于参数调优。


Figure 10. Mario 模拟器的准确性。

6.7 集群实验
- 实验设置：在 64 个 A100-40G GPU 上使用 GPT3-13B 模型，并引入数据并行。
- 结果：如图 11 所示，1F1B、Chimera 和 Interleave 的最佳配置分别为 V-64-16、X-64-16 和 W-64-32（格式为：方案-PP维度-微批次大小），且都启用了 Mario。实验表明，更大的 PP 维度配合更大的微批次大小能获得更高吞吐量，而 Mario 通过消除内存瓶颈使得这种配置成为可能。整个调优过程耗时 210 秒，远快于手动配置（一次约 10 分钟）。


Figure 11. 调优迭代过程中的吞吐量曲线。

A7 补充细节

7.1 讨论

流水线阶段划分：通过阶段划分可以实现计算平衡或内存平衡，但通常无法两者兼得。为了平衡内存，可以给后面的阶段分配更多层，但这会加剧计算不平衡，对流水线更为不利。实际上，均匀划分被 SOTA 方案（如 Chimera【23, Chimera: efficiently training large-scale neural networks with bidirectional pipelines, 2021, SC】）和主流框架（Megatron-LM）广泛采用。不均衡划分的收益有限（例如 AdaPipe【39, AdaPipe: Optimizing Pipeline Parallelism with Adaptive Recomputation and Partitioning, 2024, ASPLOS】报告的加速 < 4%），并且可能损害重计算的隐藏效果。

内存碎片：内存碎片问题在动态细粒度检查点方案中很常见。然而，Mario 采用静态、粗粒度（整个流水线阶段）的检查点，因此碎片问题要轻微得多。此外，与碎片相关的优化可以直接应用于 Mario。

应用于更大规模系统：尽管实验受限于硬件规模，但 8-GPU 和 32-GPU 的流水线评估已能代表大规模场景。根据 LLM 训练的最佳实践，PP 维度通常在 4 到 16 之间。例如，在 6144 个 H100 GPU 上训练一个 462B 模型时，PP 维度仅为 16。因此，文中的 32-GPU 流水线已经足够大。未来工作将探索 Mario 与 ZeRO-DP/Offload、TP 等技术的结合，以发现和利用新的“气泡”来隐藏重计算。

8. 相关工作

解决 LLM 训练中的内存墙问题：这是 LLM 训练中的一个关键问题。解决方案包括：通过并行训练框架（Megatron-LM, DeepSpeed, Colossal-AI）在设备间分配内存占用；通过 ZeRO【35, ZeRO: memory optimizations toward training trillion parameter models, 2020, SC; 36, ZeRO-infinity: breaking the GPU memory wall for extreme scale deep learning, 2021, SC】 和 FSDP【48, PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel, 2023, VLDB Endow.】划分优化器状态；通过 ZeRO-Offload【38, ZeRO-Offload: Democratizing Billion-Scale Model Training, 2021, USENIX ATC】等技术利用异构内存进行卸载；以及通过激活检查点【4, Training deep nets with sublinear memory cost, 2016, arXiv; 15, Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization, 2020, MLSys; 21, Dynamic Tensor Rematerialization, 2021, ICLR; 39, AdaPipe: Optimizing Pipeline Parallelism with Adaptive Recomputation and Partitioning, 2024, ASPLOS】用计算换内存。

LLM 训练中的性能模型：为了在巨大的搜索空间中协调并行方案，性能模型至关重要。Chimera【23, Chimera: efficiently training large-scale neural networks with bidirectional pipelines, 2021, SC】、MixPipe【47, MixPipe: Efficient Bidirectional Pipeline Parallelism for Training Large-Scale Models, 2023, DAC】、ZB-H1【34, Zero Bubble (Almost) Pipeline Parallelism, 2024, ICLR】 和 AdaPipe【39, AdaPipe: Optimizing Pipeline Parallelism with Adaptive Recomputation and Partitioning, 2024, ASPLOS】都为其流水线构建了性能模型。

利用流水线气泡：流水线气泡可以用来重叠其他任务。PipeFisher【32, PipeFisher: Efficient training of large language models using pipelining and Fisher information matrices, 2023, MLSys】重叠 K-FAC 优化器。Bamboo【40, Bamboo: Making Preemptible Instances Resilient for Affordable Training of Large DNNs, 2023, NSDI】重叠冗余计算以实现容错。BPipe【20, BPipe: Memory-Balanced Pipeline Parallelism for Training Large Language Models, 2023, ICML】重叠激活传输以平衡内存。Mario 的未来工作可以进一步采用 ZB-H1【34, Zero Bubble (Almost) Pipeline Parallelism, 2024, ICLR】的分割后向部分来重叠剩余的气泡。

A5 结论

通过 Mario，本文成功展示了将激活检查点与流水线并行方案中的气泡重叠以实现近乎零开销内存优化的可行性。通过利用消除的内存空间，Mario 能够在不改变全局批量大小的情况下增加微批次大小以增强训练吞吐量，并支持更长的序列长度以提高模型质量。未来，希望 Mario 能激发更多关注，探索在 LLM 训练的所有并行维度中实现近乎零开销激活检查点的潜力。