跳转到主要内容

流水线并行:GPipe、1F1B 与气泡账

流水线并行的调度结构:GPipe 与 1F1B 的气泡公式、峰值激活差异、interleaved 变体与 microbatch 选型

· 约 5 分钟阅读

本文目录

训练基础 的执行图是单卡视图;当模型大到单卡装不下,ZeRO 切静态状态,而流水线并行(PP)把层切到不同卡上。PP 特有的代价是气泡:下游 stage 必须等上游算完才能开工。这页把气泡的结构、公式和调度变体讲清楚。

30 秒复习
  • 一句话:PP 把模型按层分段到 p 张卡,batch 拆成 m 个 microbatch 在段间流水;GPipe 先全量前向再全量反向,1F1B 尽早反向把峰值激活从 m 份降到 ≤p 份,但两者的气泡占比都是约 (p−1)/m。
  • 三个判断:气泡大小由 p 和 m 决定,与调度方式无关(1F1B 不缩气泡);1F1B 的真正收益是峰值激活(内存),不是时间;interleaved 1F1B 让每卡管 v 段,气泡缩 v 倍,代价是通信次数 ×v。
  • 核心模型:bubble ≈ (p−1)/m;工程经验 m ≥ 4p 把气泡压到可接受;峰值激活 GPipe = m 份 stage 激活,1F1B ≤ p 份。
  • 边界:公式假设各 stage 均衡且忽略 p2p 通信;embedding/输出头让首尾 stage 天然不均,跨节点慢链路会放大 p2p 延迟。

1. 为什么会有气泡

PP 把 L 层切成 p 段,每张卡(或每组卡)一段。一个 batch 的完整前向要求数据沿 stage 0→1→…→p−1 顺序流动,反向再流回来——依赖链是串行的。如果整批一次过完,任意时刻只有一个 stage 在干活,其余全闲。解决办法是把 batch 拆成 m 个 microbatch 错峰流水:stage 0 算完 microbatch 1 的前向立刻发给 stage 1,自己接着算 microbatch 2。

但流水的两端——启动时的填充(fill)和结束时的排空(drain)——填不满,这段 GPU 空转就是气泡。

2. GPipe:全量前向 + 全量反向

流水线调度:GPipe vs 1F1B

GPipe 的调度是 AFAB(all-forward-all-backward):先把 m 个 microbatch 的前向全部做完,再统一做反向。对 stage 0 来说,气泡卡在中间——前向全做完后,要等最后一个 microbatch 从下游逐级反向回来才能开始自己的反向。

气泡时间 = (p−1)·(t_f + t_b)(等首段填充 + 尾段排空),与有用功 m·(t_f + t_b) 相比:

bubble 占比=(p−1) (tf+tb)m (tf+tb)+(p−1)(tf+tb)≈p−1m(m≫p)\text{bubble 占比} = \frac{(p-1)\,(t_f + t_b)}{m\,(t_f+t_b) + (p-1)(t_f+t_b)} \approx \frac{p-1}{m} \quad (m \gg p)

GPipe 的痛点不只是气泡:AFAB 要求所有 m 个 microbatch 的激活同时存活(反向开始前不能释放),峰值激活 = m × 单 stage 单层激活。m 越大气泡越小,激活却越多——两者直接打架。

3. 1F1B:不缩气泡,缩激活

1F1B(one-forward-one-backward)的调度分三段:

  • Warmup:先做若干次前向(stage i 做 min(p−i−1, m) 次),让流水线填满;
  • Steady:一次前向 + 一次反向交替——每个 microbatch 的反向尽早开始,它的激活用完即释放;
  • Cooldown:排空剩余反向。

关键事实:1F1B 的气泡占比与 GPipe 完全相同,仍是 (p−1)/m——填充和排空的总量没变,只是从中间挪到了末尾(对 stage 0 而言)。它真正的收益是峰值激活:steady state 下每个 stage 最多只有 ≤p 个 microbatch 的激活在飞(in-flight),而不是 m 个。于是 m 可以开得更大(压气泡)而不用担心显存——两个目标解耦了。

4. Interleaved 1F1B:用段数换气泡

Megatron 的 interleaved(virtual pipeline)变体让每张卡负责 v 个不相邻的 stage 段(例如卡 0 管第 1、9 层而不是连续八层)。等效流水线深度变成 p·v,每次流动的激活片段变小:

bubbleinterleaved≈1v⋅p−1m\text{bubble}_{\text{interleaved}} \approx \frac{1}{v} \cdot \frac{p-1}{m}

代价:p2p send/recv 次数 ×v(每个 microbatch 要多走 v−1 次跨卡交接),在慢链路上可能吃掉收益。

5. 实践口径

  • m 的选择:经验 m ≥ 4p 把气泡压到 25% 以下;同时 microbatch 不能小到让单卡 GEMM 失去效率(回到 Roofline 的算术强度约束)。
  • PP 放最慢链路:p2p 只传边界激活(B·S·H 量级),比 TP 的层内 all-reduce 轻得多,所以 PP 通常跨节点、TP 留节点内——见 Megatron 并行 与 GPU Communication。
  • 负载不均衡:首尾 stage 多背 embedding 和输出头,切层时通常让它们少背一层;MoE 的 EP 维度还要与 PP 正交组合。
  • 与 ZeRO 的关系:PP 切层(模型维度),ZeRO/FSDP 切状态(数据维度),两者正交可叠加;激活显存在 PP 下按 stage 数分摊,配 activation checkpointing 进一步压缩。

相关页面

修改历史1 次提交