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 显存限制。

手算气泡率: p 个 stage、m 个 microbatch → $\text{bubble}=\frac{p-1}{m+p-1}$; 给出数值例(p=4,m=4 → 3/7≈43%; p=8,m=32 → 7/39≈18%)并说明如何降低气泡(增大 m、interleaved 分块)。

描述 TP 列/行并行与双 AllReduce: 列并行输出为什么需要 AllReduce 求和、行并行反向为什么也需要; 每 transformer 层每迭代共几次(前向 2 + 反向 2 = 4)。

对比 greedy 与 1F1B 调度: 为什么 1F1B 峰值显存更低(in-flight microbatch ≤ p)、反向启动更早; 气泡为何没有因此消除(stage 间的流水线等待)。

推导 Ring-Attention: 序列切 N 段后每卡 KV 显存为何 O(1); 每层通信量 ≈ 2 × 全部 KV; 带宽/延迟特性(带宽恒定、N−1 次串行 hop)与 FlashAttention 的关系。

3D 整合题: 按通信层次解释 DP/TP/PP 分工(TP 节点内高频小消息、PP 节点间低频大消息、DP 梯度 AllReduce), 为什么 3D 并行能同时扩大并行度、降低单卡显存与通信开销。