DP、ZeRO-1/2/3、FSDP：每卡显存多少、通信多多少？
Megatron distributed optimizer 是哪一级？
先弄清楚：每个概念是什么
- DDP（data parallel）
- 每张卡一份完整模型，各吃不同的数据，backward 后 all-reduce gradient 求平均。gradient 分成 bucket，一个 bucket 算完就异步 all-reduce，和剩下的 backward 同时跑。
- Model states：16Ψ bytes
- mixed precision Adam 每个参数：bf16 param 2 + bf16 grad 2 + fp32 master、m、v 共 12 bytes（推导见 optimizer）。例：7B 模型 DDP 每卡一整份，112 GB，一张 80 GB 的卡放不下。
- reduce-scatter 与 all-gather
- all-reduce = reduce-scatter（求和，每卡只拿到结果的 1/N）+ all-gather（把各卡的 1/N 拼回完整的），详见 all-reduce。
- ZeRO-1 / 2 / 3
- 把 model states 按 DP rank 切成 N 份，每张卡只保管、只更新自己那 1/N。stage 1 切 optimizer state，2 再切 gradient，3 再切 parameter。
- FSDP
- PyTorch 版的 ZeRO-3：模型分成若干 unit，算到某个 unit 前 all-gather 出完整参数，算完马上释放；backward 后 reduce-scatter gradient。每卡显存 ≈ 分片后的模型 + 最大的那个 unit。
看图：ZeRO 每级切掉什么，FSDP 一个 unit 怎么走
30 秒回答
In DDP every GPU holds a full replica, and gradients are all-reduced bucket by bucket during backward, overlapping communication with compute. Mixed-precision Adam needs 16 bytes per parameter, replicated on all N ranks. ZeRO shards these model states across the data-parallel ranks: stage 1 shards optimizer states, 4Ψ + 12Ψ/N bytes per GPU; stage 2 adds gradients, 2Ψ + 14Ψ/N; stage 3 adds parameters, 16Ψ/N. Stage 2 communicates 2Ψ, a reduce-scatter plus an all-gather, same as DDP, and the paper says stage 1 adds no communication either; stage 3 gathers parameters again in backward, 3Ψ or 1.5×. FSDP is PyTorch's ZeRO-3; Megatron's distributed optimizer is ZeRO-1-style optimizer-state sharding.
推导 / 关键公式
每卡 model states（bytes）
没切的部分每卡一整份，切了的部分 \(\div N\)。以 ZeRO-1 为例（ZeRO-2、3 见下表）：
\[ \text{ZeRO-1: } \underbrace{2\Psi + 2\Psi}_{\text{param, grad}} + \underbrace{\frac{12\Psi}{N}}_{\text{master, m, v}} \]每卡通信量（元素数，ring 下再乘 \((N-1)/N\)）
DDP：all-reduce \(= \Psi + \Psi\)。ZeRO-2：gradient reduce-scatter \(\Psi\) + 更新后 param all-gather \(\Psi\)，还是 \(2\Psi\)。论文说 ZeRO-1 同样不增加通信，但没有单独推导。ZeRO-3：forward、backward 各 all-gather 一次 param，再 reduce-scatter gradient：
\[ \Psi + \Psi + \Psi = 3\Psi = 1.5 \times \text{DDP} \]手算例子：7B、N = 8
| 做法 | 每卡 bytes | 7B、N = 8 | 通信 |
|---|---|---|---|
| DDP | \(16\Psi\) | 112 GB | \(2\Psi\) |
| ZeRO-1 | \(4\Psi + 12\Psi/N\) | 28 + 10.5 = 38.5 GB | \(2\Psi\) |
| ZeRO-2 | \(2\Psi + 14\Psi/N\) | 14 + 12.25 = 26.25 GB | \(2\Psi\) |
| ZeRO-3 / FSDP | \(16\Psi/N\) | 14 GB | \(3\Psi\) |
| Megatron distributed optimizer | \(6\Psi + 12\Psi/N\) | 42 + 10.5 = 52.5 GB | \(2\Psi\) |
Megatron 这行是 bf16 param + fp32 grad（gradient 用 fp32 累加，不切时 18Ψ = 126 GB）：param 2 + grad 4 = 6 bytes 每卡一整份，只切 fp32 main param、m、v 这 12 bytes。表里不含 activation。
必答点
- DDP：bucket + backward hook，all-reduce 和 backward overlap
- 16Ψ：bf16 param 2 + grad 2 + fp32 master、m、v 12
- ZeRO-1/2/3 依次切 optimizer state、gradient、param：\(4\Psi + 12\Psi/N\)、\(2\Psi + 14\Psi/N\)、\(16\Psi/N\)
- 通信：DDP 是 \(2\Psi\)；ZeRO-2 是 RS + AG，也是 \(2\Psi\)；论文说 ZeRO-1 同样不增加通信；ZeRO-3 是 \(3\Psi\)（1.5×）
- FSDP ≈ ZeRO-3（按 unit all-gather、用完释放、prefetch 下一个）；Megatron distributed optimizer ≈ ZeRO-1
追问
ZeRO-1 和 ZeRO-2 的通信量有区别吗？
论文只推导了 ZeRO-2。实现里如果 ZeRO-1 的 gradient 仍用 all-reduce，再 all-gather param，就是 \(3\Psi\)，换成 reduce-scatter 才是 \(2\Psi\)（按论文 §5.1 推算，牛客题问的 gap 指的就是这里）。memory 上 ZeRO-2 不归自己的 gradient 用完就释放，ZeRO-1 要留一整份。
FSDP 怎么把通信藏起来？
all-gather 放在单独的 CUDA stream 上，backward 默认 prefetch 下一个 unit 的 all-gather。forward 的 prefetch 默认关，CPU-bound 时才开，细节见 overlap。
FSDP2 和 FSDP1 有什么区别？
FSDP1 把一组参数 flatten、拼接成一个 FlatParameter 再切，FSDP2 用 DTensor 把每个参数各自按 dim-0 切，表示更简单、容易和 TP 组合，sharded state dict 不需要 all-gather，也不用像 limit_all_gathers 那样阻塞 CPU。
跨机时全切太慢怎么办？
用 hybrid sharding，机内切、机间复制，all-gather 和 reduce-scatter 只走 NVLink，见 Megatron-FSDP v2。
手写 FSDP 的一层 linear 怎么写？（xAI 题）
每卡存 \(W\) 的 1/N 行。forward 前 all-gather 出完整 \(W\)，算本地 batch 的 \(XW\) 后释放。backward 再 all-gather 一次算 \(dX\)，本地算完整形状的 \(dW\)，reduce-scatter 成各自那 1/N 行。
来源Rajbhandari et al. 2019（ZeRO）§3.1（pp.7–8）、§5、Table 1（pp.10–11）、§7（pp.13–14）· Zhao et al. 2023（PyTorch FSDP）§1（p.2）、§3、§3.2（pp.3–6）、§3.3（pp.6–7）· Megatron Core: Distributed Optimizer · PyTorch DDP notes · PyTorch fully_shard（FSDP2）