流水线并行

基本思想

• SP 和 CP 存在的问题

  • SP 和 CP 在处理长序列时很有帮助，但如果内存问题的根本原因并不是序列长度，而是模型本身的规模，这些方法的帮助就非常有限
  • 对于大模型（70B + 参数），仅仅是权重本身的大小，就已经可能超出单个节点上 4–8 张 GPU 所能承载的上限
  • 为了解决这个问题，需要引入另一种并行维度：流水线并行（Pipeline Parallelism，PP）

• 基本思想

  • 流水线并行是一种概念上非常简单但威力强大的技术 —— 将模型的层切分并分布到多张 GPU 上
  • 例如，如果有 8 张 GPU，可以将第 1–4 层放在 GPU 1 上，第 5–8 层放在 GPU 2 上，依此类推，这样一来，每张 GPU 只需要存储和处理模型的一部分层，大幅降低了单卡的参数显存需求
  • 虽然模型参数被很好地切分到了不同 GPU 上，但每张 GPU 上的激活值显存占用却保持不变 —— 这意味着，流水线并行并不能节省激活值显存，这是因为在开始第一次反向传播之前，每张 GPU 都需要先执行 $PP$ 次前向传播。虽然每张 GPU 只负责 $1/PP$ 的层数，但在第一次反向传播开始之前，它需要处理 $PP$ 个 micro-batch，因此最终需要存储的激活值数量为
  $$PP×(activs/PP)≈activs$$
  • 也就是说，激活值显存需求与不使用流水线并行时大致相同

• 这也引入了一种新的通信模式：不再像 ZeRO-3 的数据并行那样以通信的形式传递参数，而是以 “流水线” 的形式在 GPU 之间顺序传递激活张量

朴素实现

• 基本思想

  • 把模型的层简单地分布到多个设备上 —— 例如，第一个 GPU 负责前几层，第二个 GPU 负责模型的中间部分，依此类推
  • 此时，模型的前向传播就变成了一个 batch 的数据沿着模型深度被顺序地传递，依次使用每一个 GPU
  • 这种方式对互连带宽的需求非常低，因为只在模型被划分的中间点处以通信的形式传输中等大小的激活值。这和 TP 形成了鲜明对比，在 TP 中，通信会在每一层内部发生多次

• 流水线的效率

  • 上图中的灰色部分通常被称为 “Bubble” 气泡，可以通过计算气泡导致的时间损失来量化流水线的效率

  • 设 $t_f$ 和 $t_b$ 分别表示单个 micro-batch 在单个流水线 stage 上的前向和反向计算时间（一个常见的简化假设是 $t_b \approx 2 \times t_f$）

  • 如果能够做到完美并行，那么理想的 M 个 micro-batch 的总时间为

    $$M\cdot t_{fb}$$

  • 而在上面的朴素实现中，总的时间为

    $$K\cdot M\cdot t_{fb}$$

  • 气泡时间为

    $$(K-1)\cdot M\cdot t_{fb}$$

  • 气泡占比为

    $$\frac {K-1}{K}$$

  • 显然，当 K 越大，即 GPU 的数量越多时，空置的比例接近 1，即 GPU 的资源都被浪费掉了

• 内存占用问题

  • 在流水线并行的反向传播（backward）阶段，我们需要逐层计算梯度，而这一过程依赖于每一层在前向传播时产生的中间激活张量 $z$

  • 假设模型共有 $L$ 层，每层的维度为 $d$，在流水线并行的设置下，模型会被切分为多个阶段分布在不同 GPU 上

  • 对于每一块 GPU 而言，若不考虑模型参数本身的存储开销，仅为了保存这些中间激活张量，其额外的空间复杂度为：

    $$O\left (M \times \frac {L}{K} \times d\right)$$

  • 其中 $M$ 是 micro-batch 的数量，$K$ 是参与流水线并行的 GPU 总数，$\frac {L}{K}$ 则代表了每块 GPU 上分配到的模型层数

  • 从这个复杂度公式可以看出，随着模型规模的扩大，当 $N$（微批次数）、$L$（模型总层数）或 $d$（层维度）增大时，每块 GPU 上需要存储的中间激活张量总量会显著增加，从而导致 GPU 内存占用急剧上升，甚至成为训练更大模型的瓶颈

  • 因此，如何优化这部分内存开销，是流水线并行技术中需要重点解决的问题之一

• 朴素实现是顺序、依次的，而流水线并行的核心挑战，正是如何有效绕开这种顺序性，让所有 GPU 始终保持忙碌，而不会出现某一张卡在计算、其他卡在等待的情况

All forward, All backward（AFAB）

• 流水线并行的核心逻辑，是在模型并行的基础上结合数据并行的思路，通过对训练数据的精细化拆分来提升并行效率

  • 首先将完整的模型按层 / 模块切分，分配到不同 GPU 上（比如 GPU 1 负责前 10 层，GPU 2 负责后 10 层），解决单卡放不下大模型的问题
  • 为了让各 GPU 上的模型片段能 “流水线式” 持续工作（而非等待整块数据处理完），将原本的训练批次（称为 mini-batch，即最小批次）进一步拆分成更小的子批次，这些子批次被称为 micro-batch（微批次）
  • 将多个 micro-batch 依次送入由多块 GPU 组成的 “流水线” 中 —— 当第一个 micro-batch 完成 GPU 1 的计算并传入 GPU 2 时，GPU 1 可立即处理第二个 micro-batch，而非闲置等待，从而让所有 GPU 始终处于高效计算状态，充分利用硬件资源

• 气泡计算

  • 前向阶段：第 1 个微批次到达设备 $K$ 需要经过 $K-1$ 个设备的传递，耗时 $(K-1) t_f$，设备 $K$ 完成所有 $M$ 个微批次的前向传播还需要 $N \cdot t_f$，因此前向总耗时为 $(K-1+M) t_f$

  • 反向阶段：同理，反向传播从设备 $K$ 传回设备 1，最后一个微批次完成反向传播的总耗时为 $(K-1+M) t_b$

  • 总时间为

    $$M\cdot t_{fb}+(K-1)\cdot t_{fb}=(K+M-1)\cdot t_{fb}$$

  • 气泡总时间为

    $$(K-1)\cdot t_{fb}$$

  • 气泡占比为

    $$\frac {K-1}{K+M-1}$$

  • 从公式看出，要降低气泡占比，必须增大 $N$，但 AFAB 的致命缺陷是显存占用极大，因为必须等所有 $N$ 个前向做完才做反向，设备必须在显存中缓存 $N$ 份激活值（Activations）

  • 在千亿大模型训练中，显存是绝对的瓶颈，导致 $N$ 无法开得很大，从而限制了执行效率的上限

  • Gpipe 的实验证明，当 $M\ge 4K$ 时，气泡产生的空转时间占比对最终训练时长影响是微小的，可以忽略不计

• 这种调度方式被称为 all forward, all backward（AFAB）：先完成所有前向传播，再完成所有反向传播

  • 它的优点在于，前向和反向仍然各自保持顺序结构，整体训练代码的组织方式几乎不需要很大改动，因此实现起来非常简单
  • 然而，需要处理的还有激活值带来的显存开销，在 AFAB 这种实现中，必须一直把所有激活值保存在显存中，直到进入反向传播阶段，这会很快导致显存爆炸

• Gpipe 的做法是时间换空间，论文中称为 re-materialization，也即 active checkpoint

  • 具体来说，就是几乎不存中间结果，等到 backward 的时候，再重新计算一次 forward
  • 每块 GPU 上，只保存来自上一块的最后一层输入 z，其余的中间结果计算完后即丢弃，在反向传播时，再根据保存的输入 z 重新前向传播计算得到
  • 每块 GPU 峰值时刻存储大小 = 每块 GPU 上的输入数据大小 + 每块 GPU 在 forward 过程中的中间结果大小
  • 每块 GPU 上固定需要保存起始输入，记起始输入为 $N$（即 mini-batch 的大小）
  • 每个 micro-batch 是流水线形式进来的，算完一个 micro-batch 才计算下一个
  • 在计算一个 micro-batch 的过程中，会产生中间变量，它的大小为 $\frac {N}{M} * \frac {L}{K} * d$（其中 $M$ 为 micro-batch 个数）
  • 因此，每块 GPU 峰值时刻的空间复杂度为 $O\left (N + \frac {N}{M} * \frac {L}{K} * d\right)$
  • 将其与朴素模型并行中的 GPU 空间复杂度 $O\left (N * \frac {L}{K} * d\right)$ 比较，可以发现，由于采用了 micro-batch 的方法，当 $L$​ 变大时，流水线并行相比于朴素模型并行，对 GPU 内存的压力显著减小

• micro-batch 对 BN 的影响

  • BN 是对一个 batch 内所有样本在某个特征维度上计算均值和方差，再对该 batch 的数据做归一化

  • 在流水线并行的 micro-batch 划分中：

    • 把原本大小为 $N$ 的 mini-batch 切分成了 $M$ 个 micro-batch，每个 micro-batch 只有 $\frac {N}{M}$ 个样本
    • 前向计算时，GPU 只能拿到当前的 micro-batch 数据，无法直接获取整个 mini-batch 的数据
    • 如果直接用 micro-batch 的 $\frac {N}{M}$​ 个样本计算 BN 的均值 / 方差，会因为 batch 尺寸过小导致归一化效果失真（比如统计量偏差、模型不稳定）

  • Gpipe 兼顾了流水线并行的效率和 BN 的正确性，采用 “训练时用微批量，追踪全局统计量” 的策略

    • 训练阶段（前向计算）：对当前正在计算的 micro-batch（$\frac {N}{M}$ 个样本）计算临时均值 / 方差，并用这个临时值对当前 micro-batch 数据做归一化（保证流水线能正常推进）

    • 统计量更新（后台）：同时将每个 micro-batch 的临时均值 / 方差汇总，计算整个 mini-batch（$N$​ 个样本）的移动平均均值 / 方差（Moving Average），并持续更新这个全局统计量

      $$\text{running\_mean} = \gamma \times \text{running\_mean} + (1-\gamma) \times \mu_{\text{micro}}$$

    • $\gamma$ 为衰减系数，通常取 0.99

    • 测试 / 推理阶段：不再计算 batch 均值 / 方差，直接使用训练阶段追踪到的全局移动平均均值 / 方差做归一化，保证测试时的统计量是基于完整 mini-batch 的，结果更稳定

  • LN 和 BN 的核心区别在于归一化的维度，这也是 LN 不受 micro-batch 影响的关键：

    • BN：跨样本、同特征维度归一化（依赖 batch 内的多个样本）
    • LN：同样本、跨特征维度归一化（只依赖单个样本自身）

  • LN 只需要单个样本的所有特征维度数据就能计算均值 / 方差，完全不依赖 batch 内的其他样本

  • 即使是 micro-batch，只要能拿到单个样本的完整特征，就能正确计算 LN，因此不受 micro-batch 划分

One forward, one backward（1F1B）/ Llama 3.1 流水线

• 为了打破 AFAB 的 “显存墙”，1F1B 采取了交替调度的策略

  • 一旦最后一个设备 $K$ 完成了第 1 个微批次的前向传播，它立刻进行第 1 个微批次的反向传播，并将梯度传回。后续所有设备都进入 “做一个前向，紧接着做一个反向”（1F1B）的稳定状态（Steady State）
  • 1F1B 的核心思想是：在中间的稳定阶段中，计算过程是在一次前向传播和一次反向传播之间交替进行的，即尽可能早地开始反向传播
  • 1F1B 方法中 bubble 的大小和 FAFB 保持一致，因此从纯计算效率角度看，训练效率并没有得到显著提升
  • 但和 AFAB 不同的是，只需要为 $K$ 个 micro-batch 保存激活值，而不是为 $M$ 个 micro-batch 保存激活值，这显著缓解了 AFAB 中出现的激活显存爆炸问题
  • 正因为激活显存占用降低，就可以使用更多的 micro-batch，而这反过来又能够减小 bubble 的影响

• 气泡计算

  • 预热（Warm-up）与排空（Cooldown）：设备 1 需要先做 $K$ 个前向（注：严格来说是取决于设备位置，设备 $i$ 会缓存 $K-i+1$ 个激活值），才能等来第一个反向

  • 实际总耗时：尽管调度顺序变了，但关键路径的总长度并没有改变，第一个微批次走到底并传回设备 1 需要 $K \cdot t_{fb}$ 的时间，此后设备 1 还需要做剩下的 $M-1$ 个 1F1B

    $$T_{1F1B} = K \cdot t_{fb} + (M - 1) \cdot t_{fb} = (K - 1 + M) t_{fb}$$

  • 气泡总时间为

    $$(K-1)\cdot t_{fb}$$

  • 气泡占比为

    $$\frac {K-1}{K+M-1}$$

  • 可以看出，1F1B 的气泡时间与 AFAB 完全一样，都是 $(K - 1) t_{fb}$，但 1F1B 实现了显存的降维打击

  • 通过尽早执行反向传播，微批次的激活值被尽早释放，峰值显存占用从 $O (M)$ 骤降到了 $O (K)$

    • 优化逻辑：因为显存占用与 $M$ 解绑（前提是 $M \ge K$），现在可以尽可能地增大 $M$（比如把 $M$ 设为 $4K$ 甚至 $8K$），而不必担心 OOM（内存溢出）
    • 根据气泡占比公式 $\frac {K-1}{K+M-1}$，由于 $M$ 可以开得很大，1F1B 间接地大幅度压缩了流水线气泡占比，极大地提升了系统整体的 GPU 计算效率

• 1F1B 中 bubble 对性能的影响

  

  • 在左图中，当 micro-batch 数量小于或等于流水线并行度减一（$m = p - 1$）时，可以清楚地看到 pipeline bubble 的破坏性影响：吞吐率很低，而且随着 PP 度数增加进一步下降
  • 右图展示了当 micro-batch 数量远大于流水线并行度（例如 $m = 32 \gg p - 1$​）时，低 PP 度数下的性能显著改善，但在 PP 度数非常大时，仍然会受到限制
  • 在实际训练中，并不能无限制地增加 micro-batch 的数量来维持 $m \gg p - 1$，因为最终还要受到目标全局 batch size 的约束
  • 当 PP 度数不断增加而 micro-batch 数达到上限时，bubble 的大小最终仍会按照线性关系增长：
  $$r_{bubble}= \frac {p-1}{m}$$
  • 一个有趣的现象是：在 micro-batch 数较小的情况下，从一个节点（$p = 8$）扩展到两个节点（$p = 16$）时，性能只下降了 14%
  • 相比之下，张量并行在类似的跨节点场景中出现性能下降大约 43%，可见流水线并行在跨节点网络带宽较低的场景中仍然保持较好的扩展性

• 通信 - 计算掩盖

  • 在流水线阶段切换时，设备 $i$ 必须把前向的激活值（Activations）或反向的梯度（Gradients）通过 P2P（Point-to-Point，通常经由 NVLink 或网络）发送给设备 $i+1$ 或 $i-1$

  • 将单次微批次的前 / 反向总耗时拆解：

    • $t_{comp}$：纯计算耗时（包括前向和反向）
    • $t_{comm}$：纯 P2P 通信耗时（发送 / 接收激活值和梯度）

  • 没有通信掩盖时（串行执行）：设备必须等计算完成，再发送数据；等接收完数据，再开始下一步计算

    $$t_{fb\_串行} = t_{comp} + t_{comm}$$

  • 有通信掩盖时（异步执行）：利用 GPU 的独立 DMA 引擎（如 CUDA 中不同的 Stream），在进行第 $i$ 个微批次的矩阵乘法计算时，同时（异步）接收第 $i+1$ 个微批次的输入，并发送第 $i-1$ 个微批次的输出

    $$t_{fb\_掩盖} = \max (t_{comp}, t_{comm})$$

  • 由于在大模型训练中，计算量（如 $O (B \cdot S \cdot H^2)$）通常远大于 P2P 传输量（如 $O (B \cdot S \cdot H)$），所以通常 $t_{comp} > t_{comm}$

  • 完美掩盖后，$t_{comm}$ 被完全隐藏：

    $$t_{fb\_掩盖} \approx t_{comp}$$

  • 这意味着，流水线中的数据传输成本在数学期望上被抹零了

  • 在 1F1B 调度的稳定状态（Steady State）下，每个 GPU 都在循环执行 “1 个前向 + 1 个反向”

  • 通过通信掩盖，1F1B 的核心循环变得极其紧凑

    • 当 GPU 计算微批次 $i$ 的前向（$F_i$）时，后台正在发送 $F_{i-1}$ 的激活值给下一个 GPU，同时接收 $B_{i}$ 需要的上一层反向梯度
    • 在上文得出 1F1B 的气泡时间为 $(K-1) t_{fb}$，加入掩盖后，这个时间从 $(K-1)(t_{comp} + t_{comm})$ 降至 $(K-1) t_{comp}$，不仅总吞吐量提升，流水线两端填满和排空的绝对等待时间也缩短了

Interleaving stages（交错流水线）

• 基本思想

  • 虽然 1F1B 解决了显存问题并允许增大 $M$，但当模型切得极细（$K$ 很大）时，气泡的绝对时间 $(K - 1) t_{fb}$ 依然很可观
  • Interleaved 1F1B 提出虚拟流水线阶段（Virtual Stages），原本每个 GPU 负责模型连贯的 1 块（比如 GPU 1 负责层 1-10），现在将模型切分为 $v \cdot K$ 块（$v$ 为交错因子，通常设为 2 等），每个 GPU 负责不相邻的 $v$ 块（比如 GPU 1 负责层 1-5 和 51-55）
  • 这种方案需要额外的通信，因为同一个 micro-batch 在一次前向或反向计算中会多次经过同一张 GPU，而在之前的方案中只需要经过一次
  • 不过，每一次前向和反向计算的时间都会被缩短为原来的 $\frac {1}{v}$，其中 $v$ 是每张 GPU 上的 stage 数量（模型 chunk 数）

• 气泡计算

  • 耗时切分：因为每个虚拟块的层数变为原来的 $1/v$，所以处理每个微批次的前 / 反向时间也变成了 $\frac {t_f}{v}$ 和 $\frac {t_b}{v}$，单个虚拟块的完整耗时为 $\frac {t_{fb}}{v}$

  • 气泡产生机制：流水线气泡只产生在流水线 “填满前” 和 “排空时”

  • 在 Interleaved 模式下，第一个微批次只需要经过 $K$ 个虚拟块（而不是 $K$ 个完整 GPU 阶段）就能到达最后一个 GPU，耗时从 $(K-1) t_f$ 缩短为 $(K-1)\frac {t_f}{v}$

  • 实际总耗时：总耗时等于处理所有 $M$ 个微批次的理想时间，加上流水线两端的绝对气泡时间：

    $$T_{Interleaved} = M \cdot t_{fb} + (K - 1) \frac {t_{fb}}{v}$$

  • 气泡时间：

    $$Bubble_{Interleaved} = \frac {K - 1}{v} t_{fb}$$

  • 气泡占比：

    $$Fraction_{Interleaved} = \frac {\frac {K - 1}{v} t_{fb}}{(M + \frac {K - 1}{v}) t_{fb}} \approx \frac {K - 1}{vM + K - 1}$$

  • 可以看到，直接缩减气泡：Interleaved 1F1B 在物理层面直接将气泡体积缩小为原来的 $1/v$

  • 在 $M$ 相同的情况下，原本气泡占用了 $(K-1) t_{fb}$ 的算力，现在只占用 $\frac {K-1}{v} t_{fb}$，例如 $v=2$ 时，气泡直接减半

  • 需要注意的是，这种优化是以增加 GPU 间的通信次数（通信量翻了 $v$ 倍）为代价的，因此 $v$ 不能无限大，通常 $v \in \{2, 4\}$ 是计算与通信 tradeoff 后的最优解

• 由此，可以同时通过增加 micro-batch 数 $m$，以及增加交错阶段数 $v$，来进一步减小 bubble；同时从量化角度看，通信量也会随 $v$ 成比例增加，因此这本质上是一个 trade-off

  • $m = 1, v = 1$ 对应的是最朴素的流水线并行
  • $v = 1$ 对应 AFAB 或 1F1B
  • $v \neq 1$ 对应交错流水线

• 通信 - 计算掩盖

  • Interleaved 1F1B 把每个 GPU 负责的层数拆成了 $v$ 个虚拟块（交错因子），气泡公式变为 $\frac {K-1}{v} t_{fb}$

  • 切分得越细，网络边界就越多，原本 1 个微批次在 GPU 间只需传递 1 次，现在被切成 $v$ 块后，GPU 之间的 P2P 通信次数直接翻了 $v$ 倍

  • 而且，每次 P2P 传输的边界激活值大小（Batch $\times$ SeqLen $\times$ HiddenSize）并没有变小

  • 如果无法掩盖通信，Interleaved 模式的单个微批次真实耗时将变成：

    $$t_{fb\_interleave} = t_{comp} + v \cdot t_{comm}$$

  • 此时，整个流水线的总耗时变为：

    $$T = N \cdot (t_{comp} + v \cdot t_{comm}) + \frac {p - 1}{v} (t_{comp} + v \cdot t_{comm})$$

  • 只要 $N$ 足够大，新增的通信惩罚 $N \cdot v \cdot t_{comm}$ 将远远超过缩减的气泡收益，导致负优化

  • 如果完美掩盖通信，由于我们把计算切分成了 $v$ 块，单次虚拟块的计算时间变为 $\frac {t_{comp}}{v}$

  • 只要切分后，单次计算时间仍然大于单次通信时间，即：

    $$\frac {t_{comp}}{v} \ge t_{comm}$$

  • 那么每一次额外的 P2P 通信就都可以被对应虚拟块的计算完美隐藏在后台，此时单位时间：

    $$t_{fb\_interleave} \approx \max \left (t_{comp}, v \cdot t_{comm} \right) = t_{comp}$$

  • 这时候，总耗时公式就变成了真正的理想状态：

    $$T = N \cdot t_{comp} + \frac {p - 1}{v} t_{comp}$$

  • 通信掩盖将 Interleaved 1F1B 增加的 $O (v)$ 倍通信开销强行降维打击为 $0$，它是 Interleaved 1F1B 能够实现在真实集群中性能落地的绝对前提

  • 这也是为什么在实践中（如 Megatron-LM），$v$ 的大小（如 $v=2$ 或 $4$）取决于 $\frac {t_{comp}}{v}$ 与 $t_{comm}$ 的临界点，不能无限增大

• 交错流水线的调度策略

  

• 在交错流水线方案中，调度本身也会变得更加复杂，因为在任意一张 GPU、任意一个时刻，都需要决定

  • 优先让较早的 micro-batch 通过较后的层，从而尽快闭合前向和反向回路（depth-first，目标是尽快把 batch 跑完整个模型）
  • 还是优先让较晚的 micro-batch 通过较早的层，从而尽可能填满整个流水线（breadth-first，目标是最大化流水线利用率）
  • Breadth-First Pipeline Parallelism