给定模型和 cluster，怎么选 parallelism 配置？
（含 trillion-parameter training system design）
先弄清楚：每个概念是什么
- 3D parallelism：卡数 \(= t\cdot p\cdot N\)
- 先把一份模型切成 \(t\times p\) 块放下（TP 切每层，PP 按层分段），再整套复制 \(N\) 份，各吃不同的数据（DP）。
- Communication 有轻有重
- TP 每层 forward、backward 各 2 次 all-reduce，每卡共 \(8bsh\,(t-1)/t\)，而且挡在计算前面。PP 每个 micro-batch 只在 stage 边界 point-to-point 发一次 \(bsh\)。DP 每个 step 才同步一次 gradient。
- NVLink domain
- 一台机器的 8 张卡由 NVLink/NVSwitch 互连，跨机走 InfiniBand，慢得多。所以越重的 communication 放得越靠里，TP 用到单机 GPU 数为止，再大的模型靠 PP 跨机。Megatron 实测 TP 和 PP 组合没选好，吞吐可以差 2 倍。
- Memory 预算
- model states 16 bytes/param，每卡分到 \(16\Psi/(t\cdot p)\)，再加 activation（见 training memory）。
看图：TP、PP、DP 在 1024 张卡上怎么摆
30 秒回答
I start with a memory budget, then place each kind of parallelism by how much it communicates. Model states take about 16 bytes per parameter plus activations, which tells me how many GPUs one replica needs. Tensor parallelism all-reduces inside every layer, so I keep it within a node on NVLink, usually 8. If the model still doesn't fit, I add pipeline parallelism across nodes, with many more micro-batches than stages to keep the bubble small. Context parallelism handles long sequences, expert parallelism MoE. The remaining GPUs go to data parallelism with a sharded optimizer. Then I tune micro-batch size and recomputation and check MFU.
怎么选：先算账，再按 communication 从重到轻往外套
- 算账：\(16\Psi\) + activation，除以单卡显存 → 一份模型至少几张卡。
- TP 机内：\(t\) ≤ 每台 GPU 数（通常 8），再大通信会盖过计算（TP communication）。
- PP 跨机：取放得下的最小 \(p\)。bubble 占 \((p-1)/m\)，所以保证 \(m \gg p\)，用 interleaved 1F1B（PP）。
- 按需：sequence 很长加 CP；MoE 加 EP。
- 剩下给 DP：\(N = \text{卡数}/(t\cdot p)\)，配 distributed optimizer / ZeRO（DP、ZeRO、FSDP）。
- 调优：micro-batch size、recompute、overlap，看 MFU。
手算例子：175B GPT，1024 张 A100-80GB
\(h = 12288\)、\(L = 96\)、\(s = 2048\)，global batch 1536 条，每台 8 卡。
- model states \(16 \times 175\text{B} = 2.8\) TB → 一份模型至少 35 卡。
- \(t = 8\)，占满一台机器。
- \(p = 4\)：每卡 \(2.8\text{ TB}/32 = 87.5\) GB，放不下。\(p = 8\)：43.75 GB，加 activation（SP + selective recompute、\(b = 1\)，第一个 stage 存满 \(L\) 层：\(34sbhL/t \approx 10.3\) GB）共约 54 GB ✓
- \(N = 1024/64 = 16\)；每个 pipeline 分到 96 条 → \(m = 96\)，bubble 约 7.3%（\(7/96\)）
| 真实配置 | GPU | t | p | CP | N | 利用率 |
|---|---|---|---|---|---|---|
| GPT 145.6B（Megatron） | 1536 | 8 | 8 | – | 24 | 47%* |
| GPT 1T（Megatron） | 3072 | 8 | 64 | – | 6 | 52%* |
| Llama 3 405B，seq 8K | 16384 | 8 | 16 | 1 | 128 | MFU 41% |
| Llama 3 405B，seq 128K | 16384 | 8 | 16 | 16 | 8 | MFU 38% |
* 占峰值的比例，FLOPs 含 recompute。Llama 3 把维度从内到外排成 [TP, CP, PP, DP]，DP 用 FSDP。
Trillion-parameter system design 还要讲
- 算时长：\(6\Psi D\) ÷（卡数 × 峰值 × MFU），见 FLOPs 与 MFU；Megatron 估 1T 模型、450B token、3072 卡约 84 天。
- 容错：Llama 3 405B 54 天里中断 466 次 → 定期存 distributed checkpoint（异步写、能换并行度重新切分），自动发现坏卡并重启。
- 数据与监控：data pipeline 不能让 GPU 等；盯 MFU、loss spike、慢卡。
必答点
- 先算 memory（\(16\Psi\) + activation），决定 \(t\cdot p\)
- TP 机内 ≤ 8；PP 跨机，\(m \gg p\)；剩下全给 DP + ZeRO/FSDP
- CP 管长 sequence，EP 管 MoE
- 例：175B、1024 卡 → \(t = 8\)、\(p = 8\)、\(N = 16\)
- 1T 规模还要 checkpoint、自动重启、监控
追问
用了 distributed optimizer（ZeRO-1）之后呢？
12 bytes 的 optimizer state 再按 \(N\) 切。上例 \(p = 8\) 时每卡 model states 从 43.75 GB 降到约 13 GB，\(p\) 可以降到 2–4，stage 少、bubble 小，communication 量和普通 DP 一样。
micro-batch size 怎么选？
\(b\) 大，GEMM 效率高，但 \(m = B/(N\cdot b)\) 变小，bubble 变大，activation 也变多。总时间约 \((m + p - 1)\,(t_f(b) + t_b(b))\)，要实测，Megatron 的 1B 模型例子里最优是 \(b = 4\)。
DP 能一直加吗？
受 global batch 限制。Llama 3 从 8K 卡扩到 16K 卡时 global batch 固定为 16M token，每个 DP group 的 batch 减半，MFU 从 43% 降到 41%。
来源Narayanan et al. 2021 §1（p.2）、§2.2（pp.3–4）、§3.2–3.4 Takeaway #1–#3（pp.5–6）、§5.1 Table 1（p.8）、§5.4.1（p.10）· Korthikanti et al. 2022 §4–5 eq. 5–6（pp.8–9）· Rajbhandari et al. 2020, ZeRO §3.1（p.8）、§7（p.13）· Llama Team 2024, The Llama 3 Herd of Models §3.3.2 Table 4（pp.10–12）、§3.3.4（p.13）