ZeRO 与 FSDP：把静态显存三份按 Rank 切片

训练基础 给出静态显存账：bf16 + AdamW 每参数 16 bytes，7B 模型 112 GB。普通数据并行（DDP）让每张卡都完整存一份——N 卡集群存 N 份一模一样的 Θ、∇W、m、v，这是巨大的浪费。ZeRO（Zero Redundancy Optimizer）的核心观察：这三份静态数据在任意时刻只有当前 rank 负责更新的那一小片是必需的，其余都可以分片存放、用时取回。

1. 逐级切片

级别	切什么	单卡静态显存	通信原语	通信量（每步每 rank，以 P 计）
DP baseline	无分片	16P	all-reduce ∇W	≈ 2P
ZeRO-1	optimizer states（8P 中的 12P：m、v、master）	4P + 12P/N	all-reduce ∇W	≈ 2P
ZeRO-2	+ 梯度 ∇W	2P + 14P/N	reduce-scatter + all-gather	≈ 2P
ZeRO-3 / FSDP	+ 权重 Θ	16P/N	逐层 all-gather Θ + reduce-scatter ∇W	≈ 3P

显存列的读法（bf16 + AdamW，P = 参数量）：ZeRO-1 下每卡仍存全量 bf16 权重（2P）和梯度（2P），优化器状态（12P）只存 1/N；ZeRO-3 下三份全按 1/N 存。

2. 通信原语怎么变

• DP / ZeRO-1：反向结束后对梯度做一次 all-reduce，各卡拿到全量平均梯度，各自更新自己负责的分片（ZeRO-1 只更新 1/N 的 master），再广播更新后的权重分片。
• ZeRO-2：all-reduce 拆成两步——reduce-scatter 让每卡只收到自己那份梯度的归约结果，更新本地 master 分片后 all-gather 拼回全量权重。总通信量与 all-reduce 相同，但梯度不用全量落盘。
• ZeRO-3 / FSDP：权重也不全量存。前向进入某层前先 all-gather 该层权重分片，用完即弃（或按 prefetch 流水）；反向同理再 gather 一次；梯度 reduce-scatter。参数相关通信从每步 2P 升到约 3P（前向 gather 1P + 反向 gather 1P + 梯度 scatter 1P）。

3. FSDP：PyTorch 原生实现

FSDP（FullyShardedDataParallel）把 ZeRO-3 的思想做成 PyTorch 原生 API：模型按 FSDPUnit/module 边界包成扁平的 FlatParameter 分片，前反向通过 hook 在 module 边界触发 all-gather 与释放，sharding_strategy 提供 FULL_SHARD（=ZeRO-3）、SHARD_GRAD_OP（≈ZeRO-2）等档位。FSDP2（fully_shard）把管理粒度从 FlatParameter 改成 DTensor/每参数分片，简化了混合精度与 checkpoint 的交互。与 DeepSpeed 的选择主要看生态：已在 PyTorch 原生栈就选 FSDP，需要 ZeRO 全家桶（offload、Infinity）再看 DeepSpeed。

4. 什么时候用哪档

• 显存够装静态三份：ZeRO-1 是免费午餐（通信与 DP 相同，省 12P×(1−1/N)）；ZeRO-2 再上梯度分片，通信仍不涨。
• 静态装不下：ZeRO-3/FSDP，权重逐层 gather。此时通信成为关键路径，要确认 all-gather 与前向计算重叠（prefetch），否则 GPU 空等；跨节点带宽弱时优先用节点内 TP 减小 DP 域，再套 FSDP。
• 别忘了激活：ZeRO 不切激活。显存大头是激活时，先开 activation checkpointing 或上 TP/PP（见 Megatron 并行）。

相关页面

• 训练基础：一个 Step 的执行图16 B/param 静态账的出处
• 混合精度训练m、v、master 为什么占 12P
• Activation CheckpointingZeRO 管不到的激活那份
• Megatron 并行TP/PP 与 ZeRO 的组合
• nccl-test 运行all-gather / reduce-scatter 带宽的实测方法