Obscura: Concealing Recomputation Overhead in Training of Large Language Models with Bubble-filling Pipeline Transformation

作者/机构: Yuzhou Huang and Yapeng Jiang, Sun Yat-sen University; Zicong Hong, Hong Kong University of Science and Technology; Wuhui Chen, Sun Yat-sen University; Bin Wang and Weixi Zhu, Huawei Technologies; Yue Yu, Peng Cheng Laboratory; Zibin Zheng, Sun Yat-sen University

A1 主要贡献

本文旨在解决流水线并行训练大语言模型（LLM）时，早期阶段面临的内存瓶颈问题。虽然重计算（recomputation）可以缓解此问题，但会引入额外的计算开销。

核心问题: 传统方法在流水线并行训练中为解决内存瓶颈而引入重计算，但这会带来显著的计算开销，延长训练时间。现有方法或者对所有阶段进行重计算（All-Stage Recomputation），效率低下；或者选择性重计算（selective recomputation），但在严格内存限制下仍有较高开销。

研究目标: 提出一个名为 Obscura 的计算高效的流水线训练系统，旨在给定内存约束下，优化并隐藏重计算带来的开销，从而提升训练吞吐量。

核心观察与创新点:
1. 关键观察: 作者发现，在流水线并行中，反向传播过程中的“气泡”（bubbles）可以被用来隐藏重计算的开销。具体来说，反向传播阶段之间的“后向气泡”（Backward Bubbles）可以被有效利用，而第一个前向传播和第一个反向传播之间的“前向气泡”（Forward Bubbles）则未被利用。

  1. 流水线转换（Pipeline Transformation）: 基于上述观察，本文提出了一种新颖的流水线转换方法，将未被利用的“前向气泡”转换为“后向气泡”，从而更好地隐藏重计算开销。这构成了“稻草人流水线”（strawman pipeline）的基础。

  2. Obscura系统: 为解决“稻草人流水线”面临的挑战（如何充分利用后向气泡、如何减少额外内存使用、如何解决计算不平衡），本文提出了 Obscura 系统，包含三个关键组件：

    • 依赖松弛（Dependency Relaxation）: 通过一种左移方法重新调度前向传播，策略性地将其插入反向传播之间，以打破紧密的数据依赖，从而充分利用后向气泡来隐藏重计算开销。
    • 感知交换的重计算（Swapping-Aware Recomputation）: 结合激活交换（swapping）技术来减少调整阶段在第一个反向传播前的激活内存占用。同时，将流水线执行时间建模为一个优化问题，以找到最小化通信和重计算开销的最优重计算策略。
    • 分区调整（Partition Adjustment）: 引入一种阶段分区调整算法，将计算层从调整阶段重新分配到非调整阶段，以平衡各阶段的工作负载，提高整体性能。

主要贡献总结:
* 发现并提出利用流水线气泡来隐藏重计算开销的新机会，并设计了一个“稻草人流水线”来实现这一目标。
* 提出了 Obscura 系统，通过依赖松弛、感知交换的重计算和分区调整三个关键组件，显著提升了流水线训练性能。
* 在 DeepSpeed 上实现了 Obscura，并通过对 Llama-2 和 GPT-3 等多种模型的评估，证明与广泛使用的重计算基线相比，吞吐量提升最高可达 1.32 倍。

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

2.1 流水线并行

模型分区的并行策略。流水线并行【索引5，Dapple: A pipelined data parallel approach for training large models，2021，PPoPP】、【索引8，Gpipe: Efficient training of giant neural networks using pipeline parallelism，2019，NeurIPS】、【索引15，Pipedream: Generalized pipeline parallelism for dnn training，2019，SOSP】将模型层划分为不同的阶段，分布在多个GPU上。输入批次被分割成更小的微批次（micro-batches），这些微批次以流水线方式顺序处理。在每个前向或反向传播前后，阶段之间需要通信来传输激活值和梯度。


Figure 2: 1F1B流水线的调度。

GPipe的内存问题。GPipe【索引8，Gpipe: Efficient training of giant neural networks using pipeline parallelism，2019，NeurIPS】在开始反向传播之前，会顺序处理所有微批次的前向传播。这种调度方法需要存储所有微批次的激活值，导致巨大的GPU内存使用。

1F1B流水线及其调度。相比之下，1F1B流水线【索引5，Dapple: A pipelped data parallel approach for training large models，2021，PPoPP】通过更早地启动反向传播来减少内存开销，有效避免了所有微批次激活值的累积。其调度如Figure 2所示，包括三个阶段：预热（warmup）、稳定（steady）和结束（ending）。预热阶段从第一个前向传播到第一个反向传播开始，内存使用呈线性增长。稳定阶段是1F1B的核心，交替执行前向和反向传播，重用反向传播释放的内存来存储前向传播的激活值，从而确保内存使用稳定。结束阶段与预热阶段相反。此外，本文定义了两种气泡：前向气泡（对角线块），出现在最后一个前向和第一个反向传播之间；以及后向气泡（交叉线块），出现在第一个和最后一个反向传播之间。

减少流水线气泡的策略。为了减少流水线气泡，一些基于1F1B调度的策略【索引1，Varuna: scalable, low-cost training of massive deep learning models，2022，EuroSys】、【索引14，Chimera: efficiently training large-scale neural networks with bidirectional pipelines，2021，SC】、【索引15，Pipedream: Generalized pipeline parallelism for dnn training，2019，SOSP】、【索引16，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC】、【索引26，Pipemare: Asynchronous pipeline parallel dnn training，2021，MLSys】被提出，但它们通常会增加内存压力或影响收敛性。例如，Chimera【索引14，Chimera: efficiently training large-scale neural networks with bidirectional pipelines，2021，SC】通过双向调度减少气泡，但代价是更高的内存消耗。PipeDream【索引15，Pipedream: Generalized pipeline parallelism for dnn training，2019，SOSP】通过将第二个微批次的前向传播与第一个微批次的反向传播重叠来消除气泡，但在收敛性方面面临挑战。

1F1B的内存瓶颈。尽管1F1B流水线比GPipe内存效率更高，但它在早期阶段面临内存瓶颈【索引10，Bpipe: memory-balanced pipeline parallelism for training large language models，2023，ICML】、【索引29，Mpress: Democratizing billion-scale model training on multigpu servers via memory-saving inter-operator parallelism，2023，HPCA】。具体来说，在p个阶段中，阶段s需要存储多达p-s个微批次的激活值。因此，越早的阶段（s → 0）内存使用越高。如Figure 3所示，对于13B模型，阶段0的内存使用明显高于阶段7，差距高达24GB。对于18B模型，阶段0的内存使用超出了GPU内存限制（OOM），导致训练失败。

2.2 重计算

重计算技术的基本原理。减少激活值内存使用的一个直接方法是重计算技术，该技术被集成到现代框架如DeepSpeed【索引18，Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters，2020，KDD】和Megatron【索引20，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，arXiv】的各种并行策略中。


Figure 3: 在1F1B流水线中训练13B和18B Llama-2模型时各阶段的内存使用情况。18B模型的结果是估算的。

重计算的实现与开销。形式上，一个流水线阶段由一组算子组成，每个算子在前向传播期间产生并存储激活值。重计算技术通过在前向传播期间丢弃一部分激活值，并在反向传播期间重新计算它们来减少内存使用。虽然这种方法有效降低了内存消耗，但它引入了额外的计算开 ઉ销，因为被丢弃的激活值必须通过在反向传播中重新执行相应的算子来生成。

最小化重计算开销的研究。多项研究【索引7，Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory，2019，arXiv】、【索引9，Checkmate: Breaking the memory wall with optimal tensor rematerialization，2020，MLSys】、【索引11，Reducing activation recomputation in large transformer models，2023，MLSys】、【索引22，Adapipe: Optimizing pipeline parallelism with adaptive recomputation and partitioning，2024，ASPLOS】、【索引27，Accelerating the training of large language models using efficient activation rematerialization and optimal hybrid parallelism，2024，USENIX ATC】探讨了在特定内存约束下最小化重计算开销的策略。这些工作利用算子不同的计算和空间复杂度来确定一个最优的重计算策略，该策略由内存使用高但重计算开销低的“成本效益高”的算子组成。然而，仅依靠重计算来节省内存高度依赖于模型，并且不可避免地会引入巨大开销，尤其是在严格的内存限制下。

结合卸载与重计算。为了解决这个限制，一些工作【索引2，Efficient combination of rematerialization and offloading for training dnns，2021，NeurIPS】、【索引17，Capuchin: Tensor-based gpu memory management for deep learning，2020，ASPLOS】、【索引24，Superneurons: Dynamic gpu memory management for training deep neural networks，2018，PPoPP】、【索引28，v pipe: A virtualized acceleration system for achieving efficient and scalable pipeline parallel dnn training，2021，IEEE TPDS】将卸载（offloading）与重计算相结合，利用张量活跃性分析和动态内存状态来决定是交换数据还是丢弃数据。虽然这些技术通过交换减少了重计算开销，但它们受到PCIe带宽的限制，通常会引入显著的通信开销。此外，它们可能会干扰流水线并行所需的通信【索引17，Capuchin: Tensor-based gpu memory management for deep learning，2020，ASPLOS】，在某些情况下还会产生额外的CPU计算延迟【索引19，{Zero-offload}: Democratizing {billion-scale} model training，2021，USENIX ATC】。

BPipe的特定流水线方法。BPipe【索引10，Bpipe: memory-balanced pipeline parallelism for training large language models，2023，ICML】引入了一种特定于流水线的方法，利用高带宽的GPU间连接将激活值从早期阶段传输到后期阶段的空闲内存中，以缓解内存瓶颈。此外，它还结合了选择性重计算【索引11，Reducing activation recomputation in large transformer models，2023，MLSys】以进一步减少内存使用。然而，后期阶段可用的空闲内存有限，限制了BPipe交换技术的有效性，使得在严格的内存约束下更需要依赖重计算。由于其缺乏对GPU间交换和重计算的优化集成，BPipe仍然会产生显著的重计算开销。

本文的独特性。总结来说，重计算开销仍然是一个重大挑战。我们的工作Obscura发现了一个在流水线气泡内隐藏重计算开销的机会，并引入了一种新颖的流水线调度来更好地隐藏这种开销。这种方法还使得重计算和交换能更有效地结合，实现高效的内存节省。据我们所知，我们是第一个探索通过策略性地隐藏重计算来减少开销的工作。


Figure 4: 按需重计算中的观察。

3.1 1F1B流水线调度中的观察

按需重计算优于全阶段重计算。现有工作通常在流水线遇到OOM错误时，对所有阶段实施重计算（称为全阶段重计算）。然而，内存瓶颈主要集中在早期阶段。因此，重计算应选择性地应用于超出内存限制的早期阶段，我们称之为按需重计算。

观察1：按需重计算性能更优。我们评估了四阶段按需重计算流水线的执行时间。Full重计算策略丢弃几乎所有激活值，并在反向传播期间执行一次额外的前向传播。如Figure 5(a)所示，Recomp-N表示对s ≤ N的阶段应用重计算，Recomp-All表示对所有阶段应用重计算，Recomp-Non表示不进行重计算。通过比较所有配置的阶段0，我们观察到Recomp-0和Recomp-1的执行时间分别比Recomp-Non增加了19%和25%，优于Recomp-All的33%增幅。

深入分析按需重计算。为了研究按需重计算的益处，我们深入研究了Figure 4中选择性地对阶段0和1应用重计算（Recomp-1）的1F1B流水线调度，从而得出了另一个关键观察。

观察2：后向气泡可隐藏重计算开销。如Figure 4所示，应用Recomp-1后，阶段1中的后向气泡大部分被微批次6的重计算所导致的时间延长所填充，从而防止其影响后续执行。相比之下，微批次0-5由于没有足够的后向气泡来吸收重计算开销，导致整体执行向右移动，延长了流水线时间。在阶段0中，微批次6和微批次0-5的重计算也观察到类似的模式。


Figure 5: 单次迭代的计算和气泡时间。Fwd+Bwd表示总的前向和后向时间；Fwd BB和Bwd BB分别表示前向和后向气泡时间；'Recomp'表示重计算时间。

数据验证后向气泡的作用。为了验证这一发现，我们分析了Figure 5(a)中计算和气泡时间的变化。比较Recomp-Non与Recomp-0在阶段0的情况，发现后向气泡被重计算任务所取代。在Recomp-1的阶段0和1也观察到类似现象。相比之下，Recomp-All中的后向气泡未受影响，因为最后阶段没有后向气泡，且所有阶段的重计算时间相同，使得最后阶段成为最慢的阶段。这种延迟通过阶段间依赖关系传播，阻止了早期阶段利用其后向气泡。

后向气泡的局限性。尽管后向气泡可以部分隐藏重计算开销，但其有限的可用性只允许它们隐藏一小部分开销。为了探索增加后向气泡的机会，我们提出了第三个观察。

观察3：前向气泡未被利用。如Figure 4所示，前向气泡在重计算阶段中仍未被使用。这是因为前向气泡在第一次反向传播之前产生，而重计算本质上发生在反向传播期间。此外，前向气泡的持续时间明显长于后向气泡。在一个p阶段的流水线中，阶段s包含2(p − s − 1)个前向气泡，这是后向气泡数量(p − s − 1)的两倍。

数据验证前向气泡未被利用。在Figure 5(a)中，与Recomp-Non相比，Recomp-0的阶段0中的前向气泡保持不变，其气泡时间比后向气泡时间长（0.50s vs. 0.37s，由于阶段间同步通信，不是两倍关系）。在Recomp-1和Recomp-All中也观察到类似趋势，前向气泡略有延长。这是因为阶段s的重计算延迟了阶段s-1中第一次反向传播的启动，从而增加了前向气泡的持续时间。


Figure 6: 流水线转换过程（上）和生成的稻草人流水线（下）。

3.2 稻草人流水线：隐藏重计算开销的优化机会

将前向气泡转换为后向气泡。基于我们的三个关键观察，我们提议将前向气泡转换为后向气泡，以更有效地隐藏重计算开销。为了实现这一概念，我们引入了一个稻草人流水线，它是通过流水线转换从1F1B流水线派生而来的，如Figure 6(a)所示。此转换包括两个关键步骤：

  1. 步骤1：识别和重计算。在1F1B流水线中，超出内存限制的阶段被识别为“调整阶段”，并对其进行重计算。其余的称为“非调整阶段”，保持不变。
  2. 步骤2：迁移。对于调整阶段，稳定阶段中的前向传播被尽可能地向左迁移到前向气泡中，导致前向气泡向右移动并转换为后向气泡。

稻草人流水线的性能。Figure 6(b)展示了提议的稻草人流水线，其性能在Figure 5(b)中进行了评估。术语Strawman-N与Recomp-N的含义相同。与Recomp-0和Recomp-1相比，Strawman-0和Strawman-1中的重计算开销分别减少到8%和17%。这是通过消除前向气泡并用重计算任务取而代之实现的。

稻草人流水线面临的挑战。在实践中，与其对调整阶段应用Full重计算策略，不如采用一种选择性策略，只丢弃部分激活值以精确满足内存约束，这样更有效。这种方法减少了重计算开销并改善了开销隐藏效果。然而，它引入了三个关键挑战：

  1. 挑战1：如何充分利用后向气泡。调整阶段和非调整阶段之间的紧密数据依赖性——微批次在前向传播中顺序执行，在反向传播中逆序执行——限制了后向气泡隐藏重计算的能力。后向气泡只能隐藏其前面的重计算，而不能隐藏后面的。例如，如Figure 6(b)所示，微批次1的重计算无法利用前面的气泡，因为它必须等待阶段2中的反向传播完成。此外，当采用选择性重计算策略时，每个微批次的重计算时间变得比气泡时间短，导致后向气泡未被充分利用。使用仅丢弃MLP激活值的MLP重计算策略，Figure 5(c)显示，与Figure 5(b)中的Strawman-0相比，之前被充分利用的后向气泡现在有很大一部分未被使用。

    Table 1: 不同流水线中各阶段的内存使用情况。

  2. 挑战2：如何减少额外内存使用。与1F1B流水线相比，稻草人流水线中的调整阶段在预热阶段存储了更多前向传播的激活值。因此，稻草人流水线表现出更高的内存使用，可能导致OOM错误。Table 1展示了在MLP重计算策略下各阶段的内存使用情况。结果清楚地表明，稻草人流水线中调整阶段的内存使用显著高于1F1B流水线。

  3. 挑战3：如何解决计算不平衡。按需重计算在调整阶段和非调整阶段之间引入了工作负载不平衡。虽然前向和后向气泡可以部分隐藏重计算，但它们有限的可用性以及随着微批次增大而增加的开销，使得一些重计算无法被隐藏。这导致在非调整阶段形成额外的后向气泡，如Figure 6(b)所示。具体来说，如Figures 5(b)和5(c)所示，由计算不平衡引起的后向气泡在非调整阶段是显而易见的。

A2 方法细节

4. Obscura 概览

Obscura系统架构。为了应对上述挑战，我们提出了Obscura，一个训练系统，其架构如Figure 7所示。Obscura由两个主要模块组成：Obscura规划器（Obscura Planner）和Obscura运行时（Obscura Runtime）。Obscura规划器将原始的1F1B流水线转换为Obscura流水线，并生成一个部署配置供Obscura运行时使用。Obscura运行时在执行期间管理分布式部署和模型训练。


Figure 7: Obscura的系统架构。

Obscura规划器的子模块。Obscura规划器包含四个子模块：流水线转换（Pipeline Transformation），它将前向气泡转换为后向气泡以更好地隐藏重计算开销；依赖松弛（Dependency Relaxation），它松弛阶段间的紧密数据依赖以充分利用后向气泡；感知交换的重计算（Swapping-Aware Recomputation），它结合交换和重计算以高效地满足内存约束；以及分区调整（Partition Adjustment），它优化阶段分区策略以消除非调整阶段的气泡。Obscura运行时包括分布式系统（Distributed System），用于实现分布式训练，以及Obscura流水线引擎（Obscura Pipeline Engine），它实现了Obscura流水线的核心运行时。

规划器与运行时的协作。值得注意的是，Obscura规划器离线运行，仅使用几次训练迭代即可生成Obscura配置。具体来说，流水线转换子模块运行几次迭代以确定调整阶段并收集运行时性能数据，包括算子的计算和内存成本。其余子模块使用这些性能数据来解决小规模优化问题并执行轻量级算法，引入的开销可以忽略不计。此外，Obscura运行时在线运行，并在训练期间管理系统行为。

5. 依赖松弛

核心目标与方法。为了充分利用后向气泡，我们首先对调整阶段的计算和气泡进行了深入分析。基于此分析，我们提出了依赖松弛，该方法将未使用的后向气泡重新用于其他计算。此外，我们优化了前向传播的迁移，以最小化不必要的前向传播迁移所导致的额外激活内存占用。

5.1 计算与气泡分析

紧密数据依赖的限制。我们以Figure 8展示了相邻的调整和非调整阶段的简化稳定阶段。在Figure 8(a)中，左侧显示了迁移后后向气泡的未充分利用情况，而右侧则描绘了剩余的前向传播。由于两个阶段之间存在紧密的数据依赖，反向传播无法向左移动以利用左侧的气泡。具体来说，阶段s中微批次x的反向传播必须等待同一微批次在阶段s+1的反向传播完成，这阻止了使用前面的气泡进行重计算。前向传播也遇到类似的问题：它们无法向右移动以利用右侧的气泡（如果有），因为这会延迟后续阶段同一微批次的前向传播，从而延长流水线。


Figure 8: 相邻调整和非调整阶段的简化稳定阶段调度。

利用气泡的可行移动。然而，相反方向的移动来利用气泡是允许的。如Figure 8(b)所示，阶段s中微批次y的前向传播可以向左移动以利用左侧的气泡，而不会违反其数据依赖。同样，微批次x的反向传播可以向右移动。这一观察提供了一个关键见解：与其仅依赖相邻的重计算任务来利用未使用的后向气泡，我们可以将剩余的前向传播向左移动来利用这些气泡。随后，反向传播可以向右移动，以利用调整后的前向传播所产生的气泡。

调整后的效果。Figure 8(c)展示了这种调整的结果。阶段s中微批次x反向传播之前的气泡被利用，从而减少了执行时间。这表明之前未使用的后向气泡被有效地用于隐藏重计算开销。

5.2 阶段间依赖松弛

具体调整方法。基于上述分析，我们提出了依赖松弛来提高后向气泡的利用率，如Figure 9所示。关键调整在于稳定阶段中，将剩余的前向传播向左移动，并将反向传播向右移动。为了防止过多的激活累积，左移的前向传播与反向传播交错进行，类似于1F1B方法。例如，微批次4-7的前向传播被左移，并与相应的反向传播（相应地右移）以交错方式进行编排。

依赖松弛的效果。这种调整放宽了阶段间的紧密数据依赖，确保一个阶段的执行不再直接影响另一个阶段。结果是，未被充分利用的后向气泡得到了有效利用，减少了总的迭代时间，如Figure 9中间部分所示。


Figure 9: 应用于稻草人流水线的两步优化过程，以实现Obscura流水线。

5.3 迁移优化

减少不必要的迁移。我们通过减少迁移到前向气泡中的前向传播数量来优化流水线，最终形成了Figure 9底部所示的Obscura流水线。这种优化减少了预热阶段的激活内存占用，而不影响重计算开销的隐藏效果。其动机是，如Figure 9中间部分所示，阶段0中后向气泡的持续存在。

隐藏能力的决定因素。在一个p路流水线中，阶段s最多可以有 min(m − 1, 3(p − s − 1)) 个后向气泡，其中m表示微批次的总数。这表明理论上较早的调整阶段可以隐藏更多的重计算开销。然而，由于调整阶段最终反向传播的数据依赖性，所有调整阶段的总时间延长不能超过最后一个调整阶段k的时间延长。由此分析，我们得出结论，隐藏重计算开销的能力由阶段k中的气泡数量决定。这也解释了在依赖松弛后，较早调整阶段中观察到的未使用气泡。因此，对于s < k的调整阶段，只需将其后向气泡的数量与阶段k对齐即可。

6. 感知交换的重计算

目标与方法。为了减少由前向传播迁移引起的额外内存使用，我们为Obscura流水线引入了一种激活交换方案，该方案在预热阶段减少激活值的存储。此外，我们分析了重计算和通信开销之间的权衡，并构建了一个优化问题，以在给定的交换配置下确定最优的重计算策略。

6.1 激活交换

引入交换的动机。仅仅通过增加相对于1F1B流水线的重计算来消除额外的内存消耗，会引入更多的重计算开销，这与我们利用气泡隐藏重计算开销的目标相冲突。为了解决这个问题，我们将交换技术集成到Obscura流水线中，以减轻重计算的压力。


Figure 10: 激活交换方案。

激活交换方案。我们提出了一种激活交换方案，以单个微批次为粒度将激活值传输到外部存储器，称为“激活块”，如Figure 10所示。该方案包括三个交换阶段：驱逐（eviction）、驱逐-加载（eviction-loading）和加载（loading）。驱逐阶段发生在预热阶段，此时内存使用线性上升，需要持续驱逐激活块以维持目标内存水平。驱逐-加载阶段贯穿整个稳定阶段，在此期间，被驱逐的激活值被重新加载用于反向传播。为避免超出内存限制，在重新加载前必须驱逐一个激活块。利用全双工通信，驱逐和加载可以并行执行。加载阶段发生在结束阶段，此时激活值被加载而无需驱逐，因为反向传播会消耗它们，从而确保内存使用稳定。

方案的普适性。各种交换实现【索引10，Bpipe: memory-balanced pipeline parallelism for training large language models，2023，ICML】、【索引27，Accelerating the training of large language models using efficient activation rematerialization and optimal hybrid parallelism，2024，USENIX ATC】都是该方案的具体实例。通过抽象其细节并关注核心方面，我们将该方案集成到Obscura流水线的优化中，确保与不同设计的兼容性。

内存分析。在激活交换方案中，内存使用由两个关键参数决定：1) $λ$，即驱逐阶段传输的激活块数量；2) $β$，即每个激活块传输的数据比例。设 $µ_{adj}$ 为调整阶段预热阶段的前向传播次数，$a_{adj}$ 表示调整阶段每个微批次存储的激活值大小。交换后调整阶段s的激活内存消耗可表示为：

6.2 感知交换的重计算

重计算与交换的权衡。Obscura中的重计算和激活交换在重计算成本和通信成本之间形成了一种权衡。1) 增加重计算，如Figure 11(a)所示，可以减少预热阶段为满足内存限制而驱逐的激活块数量，并允许较短的传输时间与计算重叠，但会产生更高的重计算开销。2) 减少重计算，如Figure 11(b)所示，降低了重计算开销（大部分可被气泡隐藏），但需要更多的交换来满足内存约束。这包括在预热阶段驱逐更多的激活块，以及每个激活块需要更长的传输时间且无法与计算重叠，从而导致显著的通信开销。为了平衡重计算和通信开销，我们对Obscura流水线的执行时间和内存使用进行建模，并构建我们的优化问题。


Figure 11: 重计算与交换之间的权衡。

执行时间建模。我们考虑一个p路1F1B流水线，有m个微批次，在LLM的均匀阶段分区下，前向时间f和后向时间2f在所有阶段都相同。重计算时间表示为r，最后一个调整阶段表示为阶段k。对于交换，如Figure 11(b)所示，驱逐或加载操作在三个交换阶段的最大传输时间分别为f、3f+r和2f+r。超过这些阈值会引入通信开销，其中驱逐-加载阶段的成本随微批次数量的增加而增加。为减轻显著的开销，我们强制约束 $a_{adj}β ≤ (3f + r) × B$，其中B是CPU-GPU通信带宽。因此，通信成本为：

通信成本公式说明。Eq.2适用于阶段k没有剩余气泡的情况，因为气泡的存在意味着Full重计算的开销可以被完全隐藏，使得交换变得不必要。第一项和第二项分别代表在预热和结束阶段产生的通信成本，前者有气泡来减轻成本。$ε$ 代表由于交换设计【索引10，Bpipe: memory-balanced pipeline parallelism for training large language models，2023，ICML】、【索引27，Accelerating the training of large language models using efficient activation rematerialization and optimal hybrid parallelism，2024，USENIX ATC】的变化而未被前两项覆盖的开销，其范围为[0, 3f + r]。

总执行时间公式。执行时间，包括计算时间和气泡时间，表示如下：

其中 $γ = 3f + r$。$T_s$ 是从阶段s到最终阶段的流水线执行时间，主要由计算、气泡延迟和通信开销决定。max函数中的第一项捕捉了阶段间的数据依赖关系。

内存使用建模。我们定义a为没有重计算时每个微批次的激活大小。阶段s的内存使用为 $M_s = W_s + A_s$，其中$W_s$是模型参数（包括优化器状态）的大小，$A_s$是激活的内存使用。对于调整阶段，$A_s = A_{adj_s}$；对于非调整阶段，$A_s = (p − s) × a$。

优化目标。基于上述分析和建模，我们构建优化问题如下：

优化变量与求解。优化变量包括 $λ_s$、$β$、$a_{adj}$ 和 $r$。$λ_s$ 在各阶段间遵循一个固定的关系：阶段s-1比阶段s多驱逐一个激活块。因此，$λ_s$ 可以表示为 $λ_k$，范围为 [2, p - k - 2]。$β$ 是 [0, 1] 范围内的一个百分比。$a_{adj}$ 和 $r$ 对应于重计算策略，即选择进行重计算的算子集合。我们将重计算策略量化如下：定义流水线阶段的算子为 $U_{op} = \{op_0, ..., op_{n−1}\}$，其中n表示算子总数。则有 $f ≈ ∑_{op_i∈U_{op}} f_{op_i}$ 和 $a ≈ ∑_{op_i∈U_{op}} a_{op_i}$，其中 $f_{op_i}$ 和 $a_{op_i}$ 表示算子 $op_i$ 的计算时间和激活大小。重计算策略表示为 $r = ∑_{i=0}^{n−1} R_{op}^{0−1} [i] × f_{op_i}$ 和 $a_{adj} = a − ∑_{i=0}^{n−1} R_{op}^{0−1} [i] × a_{op_i}$。如果 $R_{op}^{0−1} [i] = 1$，则 $op_i$ 被重计算；$R_{op}^{0−1} [i] = 0$ 表示不重计算。通过将 $a_{adj}$ 和 $r$ 的量化定义代入Eq.4，优化问题变成一个整数规划（IP）问题。我们通过枚举与交换相关的变量的可能值来求解，从而推导出相应的重计算策略，然后选择能实现最佳性能的交换配置和重计算策略的组合。例如，如果我们采用BPipe的【索引10，Bpipe: memory-balanced pipeline parallelism for training large language models，2023，ICML】交换设计，我们只需要迭代$λ_k$，因为$β$固定为1。

7. 分区调整

目标与方法。为了平衡调整阶段和非调整阶段之间的计算，我们首先优化了Obscura流水线的执行时间和内存使用模型，考虑了不均匀的阶段划分。基于此模型，我们提出了一个阶段分区调整算法，通过找到最优的划分策略来优化流水线性能。

7.1 不均匀阶段划分下的建模

模型更新。考虑一个不均匀阶段划分的流水线，我们定义 $f_s$、$r_s$ 和 $a_s$ 分别为阶段s的前向时间、重计算时间和每个微批次的激活大小。在流水线并行中，执行时间由最慢的阶段决定。我们令 $C_{adj} = \max_{s≤k}(3f_s + r_s)$ 和 $C_{non} = \max_{s>k}(3f_s)$ 分别表示最慢的调整阶段和非调整阶段每个微批次的总计算时间。因此，Obscura流水线中阶段s的执行时间可以表示如下：

执行时间与内存模型近似。Eq.5是一个近似，其中 $C_{non}$ 和 $C_{adj}$ 分别替代了 $T_s$ 中的 $3f$ 和 $3f + r$ 项。对于内存使用，我们将 $M_s$ 中的 $a$ 和 $a_{adj}$ 替换为 $a_s$，如下所示：

7.2 分区调整算法

算法设计。为了平衡调整阶段和非调整阶段的计算，我们提出了一个分区调整算法，将层从调整阶段转移到非调整阶段。为了进行更细粒度的调整，每个transformer层被分成一个注意力层和一个MLP层，总共产生2L+2层，其中L是transformer层的数量，常数2代表嵌入层和解码器头层。在均匀阶段划分中，第一和最后一个阶段包含2L/p + 1层，而中间阶段包含2L/p层。

算法的适用场景与简化。Obscura流水线中显著的计算不平衡仅在微批次数量很大时发生，此时轻微的分区调整可以消除不平衡。为简单起见，我们忽略了调整对开销隐藏和感知交换的重计算的影响。

算法伪代码。分区调整算法在Algorithm 1中给出，它接受代表调整和非调整阶段及其层数的字典 $D_{adj}$ 和 $D_{non}$。目标是使 $T̂_{k+1}$ 尽可能接近 $T̂_k - C_{adj}$，以平衡阶段间的计算。第5-8行将一个层从最慢的调整阶段a转移到最快的非调整阶段b。算法输出更新后的字典 $D_{adj}$ 和 $D_{non}$。

8. 设计强化：增强识别调整阶段时的可隐藏性

问题。如§5所述，一方面，一个阶段的后向气泡数量与阶段ID成反比，且最大可隐藏开销由最后一个调整阶段决定。因此，选择更多的阶段作为调整阶段会降低Obscura流水线隐藏重计算开销的能力，最终使其性能退化到与1F1B流水线相当。另一方面，在流水线转换过程中，识别调整阶段是根据内存使用和硬件内存容量来指导的。在高内存使用或有限内存容量的条件下，会有更多的阶段被指定为调整阶段，从而降低了Obscura流水线的性能。

Algorithm 1: 分区调整

机会。最近的研究【索引27，Accelerating the training of large language models using efficient activation rematerialization and optimal hybrid parallelism，2024，USENIX ATC】表明，transformer模型中成本效益高的算子，如Llama-2模型中的RMSNorm、SiLU和Mul，表现出高空间复杂度和低计算复杂度。将这些算子纳入重计算集（称为CMB重计算策略）可以减少约40%的内存占用，而开销仅增加2-3%。

优化。为了增强原始的识别步骤，我们引入了一种名为CMB-Identifying的替代方法。当遇到OOM问题时，我们首先应用CMB重计算策略，而不是立即将一个阶段指定为调整阶段。这种方法减少了调整阶段的数量，从而提高了Obscura流水线隐藏重计算开销的能力。

A4 实验环境

A4 实验结果

训练性能


Figure 12: 训练不同尺寸的Llama-2和GPT-3模型在不同全局批量大小下的端到端性能。

内存分析


Figure 13: Llama-2 28B模型中的内存使用情况。

计算与气泡分析


Figure 14: Llama-2 23B模型中所有阶段的计算、气泡和总执行时间。

附加性能验证


Figure 15: 4阶段流水线中的端到端性能。

敏感性分析


Figure 16: Llama-2 23B模型中的敏感性分析。

A5 结论

本文提出了Obscura，一个计算高效的流水线训练系统。Obscura的核心创新在于引入了一种新颖的流水线转换，它能够利用现有的流水线气泡来有效隐藏重计算的成本。此外，系统还集成了依赖松弛、感知交换的重计算以及计算平衡等先进技术，这些技术协同工作，显著提升了训练的整体性能和资源利用率。实验证明，Obscura在训练各种规模的大语言模型时，相比现有方法具有明显的吞吐量优势。