第 4 章 训练状态分片：ZeRO 与 FSDP¶
开场问题：8B 模型权重只有 16 GB，80 GB 的卡为什么 DDP 起不来？
4.1 先看清楚重复在哪¶
第 1 章算过这笔账：bf16 训练下，每个参数长期占用约 16 字节——bf16 权重 2、bf16 梯度 2、fp32 主权重 4、AdamW 的一阶矩 4、二阶矩 4。8B 参数就是约 128 GB，再加上激活，单卡必然放不下。
关键观察是：DDP 下这 128 GB 在每张卡上一模一样。64 张卡存了 64 份完全相同的内容。既然如此，为什么不每张卡只存 1/64？
4.2 ZeRO 的三个阶段¶
- Stage 1：切优化器状态。每张卡只保存 1/N 的 fp32 主权重和 Adam 矩，更新时各自更新自己那片，再把更新后的参数 AllGather 回来。省掉了最大的那块。
- Stage 2：再切梯度。反向时用 ReduceScatter 代替 AllReduce——每张卡只收到自己负责那片的梯度求和结果，正好够它更新自己那片参数。
- Stage 3：连参数也切。每张卡平时只持有 1/N 的权重，要算某一层时才把这层的完整参数 AllGather 回来，算完立刻丢掉。
显存从 Stage 1 到 3 递减，通信递增。但增得没有想象中多：DDP 每步的通信量约 2 倍参数量，ZeRO-3 约 3 倍——只多了 1.5 倍，却换来显存除以 N。这个性价比是它能成为默认方案的原因。
4.3 FSDP 的实际运行过程¶
PyTorch 的 FSDP 是 ZeRO-3 思路的原生实现。把模型按模块包裹成若干 FSDP 单元，每个单元的生命周期是：前向到这个单元时 AllGather 出完整参数 → 计算 → 立刻释放分片以外的部分 → 反向时再 AllGather 一次 → 算出梯度后 ReduceScatter → 只保留自己那片。
包裹粒度是这里唯一重要的调参。包得太粗（比如整个模型一个单元），AllGather 会把全部参数拉回来，等于没切；包得太细（每个 Linear 一个单元），通信次数暴增、每次数据量太小，启动开销吃掉收益。常见做法是按 Transformer Block 包裹，一层一个单元。
旧版 FSDP 以被包裹模块的扁平化参数为分片单位，FSDP2 改成逐参数分片（per-parameter sharding），在混合精度、部分冻结、与编译器配合上都更干净。它们和 ZeRO-3 目标相同但实现不同，写文档时不要当成同义词。
4.4 Prefetch：让通信藏进计算¶
如果算第 i 层时才去 AllGather 第 i 层的参数，GPU 会干等。所以 FSDP 会在算第 i 层的同时预取第 i+1 层的参数，反向同理。
这里有个直接的权衡：预取深度越大，通信越容易被藏住，但同时驻留在显存里的完整参数层数也越多。预取两层就意味着峰值显存多出两层的完整参数——很多 FSDP 的 OOM 不是分片没生效，而是预取太激进。
4.5 Offload：把状态放到 CPU 或 NVMe¶
显存仍然不够时，可以把优化器状态甚至参数放到主机内存，需要时再搬回 GPU。它确实能让原本跑不起来的配置跑起来，但代价必须算清楚：PCIe 的带宽比 HBM 低一个数量级以上，容量问题被转化成了搬运问题。
实践中它适合两种场景：硬件严重受限时先把流程跑通；或者只 offload 更新频率低、访问模式规整的那部分状态。把它当成常规加速手段会很失望。
4.6 分片之后，checkpoint 长什么样¶
分片训练自然产生分片 checkpoint——每张卡存自己那一片，保存快、恢复快。但这带来两个必须提前想清楚的问题：
第一，换并行配置还能不能加载。用 64 卡训的分片 checkpoint，想用 32 卡续训，需要 resharding 能力。很多团队是在故障恢复时才发现自己没有。
第二，对外发布的权重是另一种格式。推理引擎要的是完整的、按层组织的权重文件。导出流程应该在训练早期就跑通一次，而不是等训练结束。
核心问答¶
Q：ZeRO-3 和 FSDP 是一回事吗？ 思路相同（参数、梯度、优化器状态全分片），实现和接口不同。FSDP 是 PyTorch 原生实现，FSDP2 采用逐参数分片。讨论时最好说清是哪一个。
Q：既然 Stage 3 最省显存，为什么不总是用它？ 通信更多、调度更复杂，小规模训练上额外开销可能超过收益。放得下的时候，简单方案通常更快。
Q：FSDP 和张量并行的本质区别？ FSDP 平时把参数存成分片，算的时候恢复成完整的，计算本身还是单卡完整计算；张量并行是让一个矩阵乘法本身就跨卡执行，任何时刻都没有哪张卡持有完整矩阵。
Q：开了 FSDP 还 OOM，先查什么？ 按顺序查：包裹粒度是不是太粗、预取深度是不是太大、激活重计算有没有开、有没有哪个不被包裹的大张量（比如巨大词表的 lm_head 输出 logits）。