v5:Pipeline Parallel
v4 的 TP 每层要做 4 次同步 all-reduce,只能留在一台机器的 8 张卡里。8 张卡放不下一份模型时,就把 36 层分成几段放到不同的机器上,机器之间只传段边界的 activation(PP)。代价是流水线会空转,叫 bubble,后半课用 schedule 把它缩小。
按层切:一条装配线
TP 切的是每一层的“宽”,PP 切的是模型的“深”。36 层按顺序分成 p 段(stage),每段放一台机器,像装配线上的工位,装完交给下一个。
- 你的 4 台机器 → p = 4,每台 9 层。forward 把输出一段段往后交,backward 把 gradient 一段段往回传
- 每卡只放 1/p 的层,parameter、gradient、optimizer state 都是 1/p:128 GB ÷ 4 = 32 GB
- 和 TP 一样是几张卡合放一份模型,不占 DP 的名额
跨机只传一个 tensor
stage 之间只做点对点的 send / recv(p2p)。forward 把段边界的 activation 发给下一台,backward 把它的 gradient 发回上一台。
- 多大:一层的输出 [b, s, h],1 × 8192 × 4096 × 2 bytes = 64 MiB
- 多快:走 IB 约 1.3 ms,来回共约 2.7 ms。机内 TP 8、跨机 PP 4 时,一个 stage 算一个 micro-batch 约 31 ms(按 6ND、400 TFLOPS 粗估),p2p 不到计算的 1/10
- 对比 TP 跨机:一个 micro-batch 每卡要发 36 × 4 × 112 MiB = 15.75 GiB,走 IB 约 338 ms。PP 每卡最多发 2 × 64 = 128 MiB,少约 126 倍
- 还能藏起来:send / recv 能和其他 micro-batch 的计算同时进行,TP 的 all-reduce 只能干等
代价:bubble
整个 batch 一次送进去,同一时刻只有一个 stage 在算,p = 4 时 3/4 的时间在等。所以把 batch 切成 m 个 micro-batch,一个接一个送进去。下面是 p = 4、m = 8 的时间线,F3 是第 3 个 micro-batch 的 forward,B3 是它的 backward,.. 是空转,F 和 B 各画 1 格(实际 B 约是 F 的 2 倍,结论不变):
GPipe:先做完全部 forward,再做全部 backward
S1 F1 F2 F3 F4 F5 F6 F7 F8 .. .. .. .. .. .. B1 B2 B3 B4 B5 B6 B7 B8
S2 .. F1 F2 F3 F4 F5 F6 F7 F8 .. .. .. .. B1 B2 B3 B4 B5 B6 B7 B8 ..
S3 .. .. F1 F2 F3 F4 F5 F6 F7 F8 .. .. B1 B2 B3 B4 B5 B6 B7 B8 .. ..
S4 .. .. .. F1 F2 F3 F4 F5 F6 F7 F8 B1 B2 B3 B4 B5 B6 B7 B8 .. .. ..
1F1B:预热之后,forward 一个、backward 一个
S1 F1 F2 F3 F4 .. .. .. B1 F5 B2 F6 B3 F7 B4 F8 B5 .. B6 .. B7 .. B8
S2 .. F1 F2 F3 .. .. B1 F4 B2 F5 B3 F6 B4 F7 B5 F8 B6 .. B7 .. B8 ..
S3 .. .. F1 F2 .. B1 F3 B2 F4 B3 F5 B4 F6 B5 F7 B6 F8 B7 .. B8 .. ..
S4 .. .. .. F1 B1 F2 B2 F3 B3 F4 B4 F5 B5 F6 B6 F7 B7 F8 B8 .. .. ..
两种 schedule 都是 22 格,每个 stage 空转 6 格。开头 S4 等 3 格才拿到第一个 micro-batch,结尾 S1 等 3 格才拿到最后一个 gradient。
bubble = (p − 1) / m
\[ \text{bubble} = \frac{\text{空转时间}}{\text{理想计算时间}} = \frac{(p-1)(t_f + t_b)}{m\,(t_f + t_b)} = \frac{p-1}{m} \]- 上图 p = 4、m = 8:3/8 = 37.5%
- p = 4、m = 32:3/32 ≈ 9.4%
- p = 8、m = 32:7/32 ≈ 22%
所以要 m 远大于 p。但 global batch 固定时,m 越大 micro-batch 越小,GEMM 越跑不满。(占总时间是 \((p-1)/(m+p-1)\),上图是 6/22 ≈ 27%。)
GPipe 和 1F1B:bubble 一样,memory 不一样
一个 micro-batch 的 activation 要从 forward 存到它的 backward 用完。GPipe 先做完 8 个 forward,S1 同时存 8 份。1F1B 的 S1 先做 4 个 forward 灌满流水线(warm-up),之后 forward 一个、backward 一个,做完就释放,最多存 4 份(= p),像洗碗洗一个放一个。一个 micro-batch 在 9 层上存 9 × 1.14 ≈ 10.3 GB(只看 PP,每个 stage 一张卡,8K tokens、开 FlashAttention),所以 GPipe 的 8 份超过 80 GB,1F1B 只要一半。
| schedule(p = 4、m = 8) | bubble | S1 的 activation 峰值 | p2p |
|---|---|---|---|
| GPipe | 3/8 = 37.5% | 8 份(= m),约 82 GB | 1× |
| 1F1B | 37.5% | 4 份(= p),约 41 GB | 1× |
| interleaved 1F1B,v = 3 | 37.5% ÷ 3 = 12.5% | 约 41 × 1.25 ≈ 51 GB | 约 3× |
Interleaved:每卡拿 v 段不相邻的层
bubble 的长度取决于一个 micro-batch 穿过整条流水线要多久,每段切短就填得快、排空得也快。36 层、p = 4、v = 3 时每段只有 3 层,S1 拿第 1–3、13–15、25–27 层。
- bubble 除以 v:\(\frac{1}{v}\cdot\frac{p-1}{m}\)。上例 37.5% → 12.5%;p = 8、m = 32、v = 4 时 22% → 5.5%
- 代价是 p2p 约多 v 倍(forward 要跨 4 × 3 − 1 = 11 个段边界,原来只跨 3 个),activation 峰值也多一点(× 1.25)
- Megatron 要求 m 是 p 的整数倍
Zero Bubble 把 backward 拆成 B(对输入的 gradient,上一个 stage 在等)和 W(对 weight 的 gradient,没人等),用 W 填空转的格子,bubble 接近 0。
各个 stage 的 memory 峰值不一样
1F1B 下 S1 到 S4 依次最多存 4、3、2、1 份 activation,约 41、31、21、10 GB。S1 的 4 份 × 9 层 = 36 层,正好是不用 PP 时的 41 GB,所以 PP 把 model state 切成了 1/p,却没有让第一个 stage 的 activation 变小。S1 还放着 embedding,S4 放着 LM head 和 loss,所以 Megatron 允许给首尾两个 stage 少放几层。
你的经验你的 32 卡任务正好是 4 台 × 8 卡。换成 Megatron,最自然的切法是机内 TP 8、跨机 PP 4,每台机器放 9 层。8B 其实用不着这样切,这里只是对一下形状。
去查一下OpenManus-RL 基于 verl,verl 的 Megatron 训练端 config(ppo_megatron_trainer.yaml)里有 tensor_model_parallel_size、pipeline_model_parallel_size 和 virtual_pipeline_model_parallel_size,最后一个就是 interleaved 的 v。看看你们的 run 用的是 FSDP 还是 Megatron,是 Megatron 的话这几个值各是多少。
预告:PP + TP + DP 拼起来
| 维度 | 切什么 | 通信 | 多频繁 | 放在哪 |
|---|---|---|---|---|
| TP | 每层的矩阵 | all-reduce | 每层 4 次,只能干等 | 机内 NVLink |
| PP | 层 | p2p send / recv | 每个 micro-batch 在每个段边界来回各 1 次 | 跨机 |
| DP | batch | gradient all-reduce | 每个 step 1 次,可以 overlap | 最外层,跨机 |
总卡数 = TP × PP × DP。你的 32 卡可以是 TP 8 × PP 4(一份模型铺满 32 张卡),也可以是 TP 8 × PP 2 × DP 2。怎么选是第 8 课的内容。
30 秒英文回答
Pipeline parallelism splits the model by layers into p stages, usually one per node, because stages only pass the boundary activation point-to-point: about 64 MiB per micro-batch, versus four all-reduces per layer for tensor parallelism. The cost is the bubble: later stages first wait for activations, and earlier stages finally wait for gradients. With m micro-batches the bubble is (p−1)/m of compute time, so m should far exceed p. GPipe and 1F1B have the same bubble, but 1F1B keeps only p micro-batches of activations instead of m. Interleaving v chunks per GPU divides the bubble by v, at v times the p2p traffic.
自测
用不用 PP,memory 峰值一样吗?
不一样。model state 每卡变成 1/p(128 GB → 32 GB),activation 却不变小。1F1B 下第一个 stage 仍存 L 层的量(41 GB),最后一个只有约 10 GB。
怎么减小 bubble?
加大 m。用 interleaved 1F1B(bubble 除以 v)或 Zero Bubble。p 能小就小,多用 TP 或 DP。首尾 stage 少放几层,把计算量调均匀。
为什么 PP 跨机、TP 留在机内?
TP 每层 4 次 all-reduce 都要等结果,跨机一个 micro-batch 约 338 ms。PP 在段边界只传 64 MiB,流量少约 126 倍,还能和别的计算重叠。
1F1B 的 bubble 比 GPipe 小吗?
不小,两者都是 \((p-1)/m\)。1F1B 省的是 memory,第一个 stage 最多存 p 份 activation,GPipe 要存 m 份。
想看细一点 6.5 Pipeline Parallel(p = 4、m = 6 的时间线,DualPipe、Zero Bubble)· 6.2 TP communication(为什么 PP 可以跨机)· 6.9 怎么选 parallelism 配置 · 6.14 activation recompute