流水线并行

标签

AI/infra/并行策略

字数

1711 字

阅读时间

8 分钟

TP 切的是单层内部，出不了机（见 08-张量并行与序列并行）。流水线并行（PP）切的是「哪些层放在哪台设备」 —— 通信只在 Stage 边界发生，量小到可以跨机，所以它才是跨机扩展的主要手段。

它要解决的核心问题只有一个：Pipeline Bubble（流水线气泡）。

问题：气泡从哪来

把模型的 $L$ 层均匀分到 $P$ 个 Stage，每个 Stage 负责 $L/P$ 层。像工厂流水线 —— 产品依次经过各工位。

但朴素实现里，一个 micro-batch 必须串行经过所有 Stage：Stage 0 在算的时候，Stage 1 到 $P-1$ 全在等。

$$\text{朴素 PP 的气泡比例}=\frac{P-1}{P}$$

$P$ 个 Stage 里有 $P-1$ 个时间步是浪费的。 4 个 Stage 就是 75% 的时间在空转 —— 这显然不可接受。

朴素 PP 的时间轴（一个 micro-batch 串行经过所有 Stage）：

  时间 ──────────────────────────────────────────▶
  Stage 0   [ 算 ][ 等 ][ 等 ][ 等 ]
  Stage 1   [ 等 ][ 算 ][ 等 ][ 等 ]
  Stage 2   [ 等 ][ 等 ][ 算 ][ 等 ]
  Stage 3   [ 等 ][ 等 ][ 等 ][ 算 ]
                    └── P−1 个时间步里同时只有 1 个 Stage 在干活 ──┘

  气泡比例 = (P−1)/P —— 4 个 Stage 就是 75% 的时间在空转

三代调度策略

GPipe：微批次

把 mini-batch 切成 $M$ 个 micro-batch，依次注入流水线。

F F F F . . .        ← 所有 micro-batch 先做完前向
. . . . B B B B      ← 再统一做反向

$$\text{气泡比例}=\frac{P-1}{M+P-1}$$

$M$ 越大，气泡越小。缺点有两个：

缺点	原因
峰值激活显存是 $M$ 份	前向全部做完才能开始反向 → 所有 micro-batch 的激活都得留着
$M$ 与显存直接冲突	$M$ 增大气泡减小，但显存爆炸

1F1B（One Forward One Backward）

交替执行前向和反向，尽早释放激活。 三个阶段：

阶段	做什么
Warmup	填充流水线（只做前向，不做反向）
Steady	每做一次 F 紧接一次 B
Cooldown	排空流水线（只做反向）

$$\text{气泡比例}=\frac{P-1}{M+P-1}(\mathbf{\text{与 GPipe 相同}})$$

但峰值激活显存从 $M$ 份降到 $P$ 份。

这是 1F1B 的全部价值所在：气泡比例一模一样，但显存降下来了。原因是 Steady 阶段「F 完立刻 B」，同一时刻最多只有 $P$ 个 micro-batch 的激活值存活（每个 Stage 一份）。

GPipe 是「先攒满再统一反向」，1F1B 是「边走边清」—— 同样的气泡，一半的代价。这也是 Megatron-LM 的默认调度。

GPipe 与 1F1B 的调度对比（F = 前向，B = 反向）：

  GPipe：所有 micro-batch 先做完前向，再统一做反向

    时间 ────────────────────────────────────────▶
    F F F F . . . .
    . . . . B B B B
    ⇒ 峰值激活 = M 份（所有 micro-batch 的激活都得留着）

  1F1B：稳态阶段「F 完立刻 B」，边走边清

    时间 ────────────────────────────────────────────────────▶
    预热 F F F │ 稳态 F B F B F B F B │ 收尾 B B B
               └── 同一时刻最多 P 个 micro-batch 的激活存活 ──┘

  ⇒ 气泡比例与 GPipe 完全相同，峰值激活从 M 份降到 P 份
  ⇒ 这就是 1F1B 的全部价值：同样的气泡，一半的代价

Interleaved 1F1B（虚拟 Stage）

每个设备不只负责连续的 $L/P$ 层，而是负责多段非连续的层。

若每个设备承担 $v$ 个虚拟 Stage（每个虚拟 Stage 有 $L/(P\cdot v)$ 层）：

$$\text{气泡比例}=\frac{P-1}{M\cdot v+P-1}$$

等效于把 $M$ 扩大了 $v$ 倍 —— 不改 micro-batch 数就拿到更小的气泡。

代价是通信频率增加 $v$ 倍 —— 每个虚拟 Stage 之间都要跨设备传激活。这是一次「用额外通信换气泡」的交易，只在机间带宽充足时才划算。

四代对比

调度	气泡比例	峰值激活显存	通信频率
朴素 PP	$(P-1)/P$	1 份	低
GPipe	$(P-1)/(M+P-1)$	$M$ 份	低
1F1B	$(P-1)/(M+P-1)$	$P$ 份	低
Interleaved 1F1B	$(P-1)/(Mv+P-1)$	$P$ 份	高（$v$ 倍）

读这张表的方式：气泡公式决定「浪费多少时间」，峰值激活决定「能不能装下」，通信频率决定「对带宽的要求」。三者不能同时最优 —— 1F1B 是默认平衡点，Interleaved 是「带宽换时间」。

工程挑战

挑战	说明
负载均衡	不同层的计算量不同（Embedding 层 vs Transformer 层），Stage 不能简单按层数均分
Embedding 与 Loss 的放置	Embedding 在最前端、Loss 在最后端，通常放在第一个和最后一个 Stage
跨机激活传输	每 Stage 边界传 $b\times s\times h$（batch × seq × hidden）—— 这是 PP 唯一的通信，量小但要 IB 支撑
与 Activation Checkpointing 的配合	PP 天然与重计算配合：反向时重算而非存储激活

PP 的通信量是四个并行维度里最小的 —— 每个 micro-batch 边界只传一次激活值，不是每层。这就是它能跨机的原因，也是 03-多卡互联与集群网络 里那张「哪个维度放哪条链路」表的第三行。

而「Embedding 和 Loss 放首尾」不只是习惯：它们的参数量与计算特征和 Transformer 层不同，混在中间会同时破坏负载均衡和显存分布。

什么时候该用 PP

PP 不是首选 —— 优先级在 DP/ZeRO 之后（先省冗余、再切层）。它的适用条件是：

• 模型层数多、单层不大（TP 切不动，因为单层装得下；DP 又装不下整个模型）
• 跨机带宽有限（TP 出不去，而 PP 的通信量小到 IB 够用）
• 能接受气泡（或愿意用 Interleaved 拿通信换）

反过来说：如果单层就装不下，该上 TP；如果只是状态冗余，该上 ZeRO。PP 解决的是第三种问题 —— 层数太多。

相关

• 08-张量并行与序列并行 —— 切单层的方案，与本篇互补
• 07-ZeRO 显存优化系列 —— 切状态冗余，优先级在 PP 之前
• 03-集合通信与 NCCL —— PP 用的点对点 Send/Recv
• 03-多卡互联与集群网络 —— 「PP 可以跨机」的带宽依据
• 10-混合精度与显存优化 —— 激活重计算与 PP 的配合

参考

• https://caomaolufei.github.io/AIInfraGuide/guides/模块三-分布式训练/第7章-流水线并行
• https://arxiv.org/abs/1811.06965
• https://arxiv.org/abs/1806.03377
• https://arxiv.org/abs/2104.04473

贡献者

YJ

文件历史

最后编辑于 大约 18 小时前查看完整历史

aa2e3-docs：补充

ae1ee-docs：内容完善

04687-docs：笔记