Megatron TP/PP 与气泡率

Megatron-LM 的 3D 并行(TP × PP × DP)实现细节。TP 列/行并行与双 AllReduce: 每个 transformer 层内, attention 的 QKV 投影用列并行(输出沿 hidden 维切分)、输出投影用行并行(输入沿 hidden 维切分), MLP 同理——列并行分支的输出是各卡局部和, 必须 AllReduce 才能得到完整结果; 每层前向 2 次 AllReduce(attention 输出投影 + MLP 第二个线性层), 反向再 2 次(列并行权重梯度需要 AllReduce), 共 4 次/层/迭代, 每次消息量 ~$b \times s \times h$。PP 1F1B 调度: 朴素 greedy 调度(先跑完所有 microbatch 的前向再统一反向)会同时驻留 m 个 microbatch 的激活、显存尖峰且反向启动晚; 1F1B(one-forward-one-backward)每个 stage 交替执行一个前向与一个反向, 同时 in-flight 的 microbatch 数被限制在 p 以内, 峰值显存更低。气泡率公式: 设 p 个 stage、每个 batch 拆 m 个 microbatch, 稳态下气泡占比 $\text{bubble} = \frac{p-1}{m+p-1}$——p 固定时 m 越大气泡越小; 数值例: p=4, m=4 → $\frac{3}{7} \approx 43\%$; p=8, m=32 → $\frac{7}{39} \approx 18\%$; p=8, m=64 → $\frac{7}{71} \approx 10\%$; m 足够大时趋近 0。interleaved(V 型/分块)调度把每个 stage 的层再分 c 块, 可把气泡再降约 c 倍, 但增加通信与显存开销。Ring-Attention 序列并行: 序列切成 N 段分到 N 卡, 每卡先算本地 chunk 的部分注意力, 再把 KV 块沿 ring 依次传递, N−1 轮后每卡收到全部 KV——单卡 KV 显存 O(1)(不随 seq 增长), 每层通信量 ≈ 2 × 全部 KV 大小(每卡完整读一遍 K 和 V), 带宽恒定、延迟随 N−1 次串行 hop 增长; 与 FlashAttention 的 online-softmax 技术同源。

训练 Infra 面试重灾区: “推导气泡率公式”“TP 为什么每层 2 次 AllReduce”“1F1B 比 greedy 好在哪里”“Ring-Attention 显存和通信怎么算”“175B 怎么配 3D”。

TP 让单层计算摊到多卡但带来高频小消息(只能靠节点内 NVLink 扛); PP 让不同 stage 各持部分层、显存随层数分摊, 但 idle 气泡把利用率吃掉——用 1F1B 压显存、用大 m 压气泡、用 interleaved 进一步折中; Ring-Attention 则把最长的维度(序列)也并行化, 让超长上下文训练不再受单卡 KV 显存限制。

检验攻克程度：针对「Megatron TP/PP 与气泡率」专属刷题排雷

做单选排雷题、推导选项机制，答错自动收录进专属错题本。