11 · ZeRO 与 FSDP

⏱ 28 分钟 ｜ ⭐ 最重要的一个分布式技术

🎯 一句话

数据并行的浪费很明显：N 张卡存着 N 份一模一样的优化器状态。 ZeRO 的想法简单到几乎是显然的：别存 N 份，切开，每卡存 1/N，用的时候再凑。 它让"数据并行"从省不了显存变成能训万亿参数。

🗑️ 一、被浪费的显存

🔑 ZeRO 的核心洞察： 这些副本在大部分时间里都是闲置的。 优化器状态只在 step() 那一瞬间用到 —— 既然如此，为什么每张卡都要一直存着完整的一份？

📊 二、三个 Stage

（8 卡，7B 模型）

🔑 最反直觉也最重要的一点： Stage 1 和 2 的通信量和普通 DDP 完全一样。

为什么？回忆第 5 章那个恒等式： AllReduce = ReduceScatter + AllGather

   DDP：     AllReduce(梯度)                    = RS + AG
   ZeRO-2：  ReduceScatter(梯度) → 各自更新自己那份
             → AllGather(参数)
             = RS + AG   ⭐ 一模一样！

ZeRO-2 只是把一次 AllReduce 拆成两半，在中间插入了"各自更新自己那片"。 所以它是几乎免费的显存节省。

🔄 三、Stage 3 怎么工作（FSDP 的核心）

💡 这就是为什么它叫 Fully Sharded Data Parallel： 从"每卡一份完整模型"变成"每卡一片，用时凑齐"。 代价是多了 1.5 倍的通信 —— 但换来的是 N 倍的显存。

⭐ 预取：让通信藏起来

# 🧩 骨架：`model` 来自你自己的代码，这一段只看写法
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy, MixedPrecision
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
import functools

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.SHARD_GRAD_OP,   # ⭐ = ZeRO-2
    # ShardingStrategy.FULL_SHARD                       #   = ZeRO-3
    # ShardingStrategy.HYBRID_SHARD                     #   机内切、跨机复制 ⭐
    auto_wrap_policy=functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls={TransformerBlock},       # ⭐ 按层包，非常重要
    ),
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,      # ⭐ 通信也用 BF16，通信量减半
        buffer_dtype=torch.bfloat16,
    ),
    device_id=torch.cuda.current_device(),
)

⚠️ auto_wrap_policy 不设置是最常见的错误： 不设置的话整个模型被当成一个单元，AllGather 时要凑出全部参数 —— 显存节省几乎为零，而且会 OOM。 必须按 Transformer 层来包。

💡 HYBRID_SHARD：大规模的最优解

⚖️ 四、ZeRO vs 张量并行：怎么选

⭐ 实用建议的顺序：

   ① 先上 ZeRO-2（几乎免费）
   ② 不够 → ZeRO-3 / FSDP FULL_SHARD
   ③ 多机 → HYBRID_SHARD
   ④ 单层都装不下 → 加张量并行（第 12 章）
   ⑤ 还不够 → 加流水线并行（第 13 章）

⚠️ 五、四个坑

# 🧩 骨架：`model` 来自你自己的代码，这一段只看写法
# ⭐ 分片保存：每卡只存自己那片，不需要凑齐
from torch.distributed.checkpoint.state_dict import get_state_dict
import torch.distributed.checkpoint as dcp

dcp.save({"model": model, "optim": opt}, checkpoint_id="ckpt/step1000")

💥 "保存 checkpoint 时 OOM" 是一个非常常见的翻车点： 训练跑得好好的，一到保存就炸 —— 因为默认的 FULL_STATE_DICT 会把所有分片汇总到 rank 0。大模型必须用分片保存。

🔗 和站内其他章的关系

✅ 检查点

1. ZeRO 的核心洞察是什么？8 卡训 7B 浪费了多少显存？
2. 三个 Stage 各切什么？显存各降多少？
3. 为什么 ZeRO-1 和 2 的通信量和 DDP 完全一样？（说出那个恒等式）
4. Stage 3 前向传播时怎么工作？峰值显存是多少？
5. 什么是预取？为什么它是 FSDP 性能的关键？
6. auto_wrap_policy 不设置会怎样？
7. HYBRID_SHARD 解决什么问题？
8. ZeRO-3 和张量并行的关键区别？什么时候必须用后者？
9. "保存 checkpoint 时 OOM"是怎么回事？

🛑 可以停在这里

⚡ 走神救援

⭐ 核心洞察：N 张卡存着 N 份一模一样的优化器状态，而它们只在参数更新那一瞬间用得到。 绝大部分显存是副本。

三个 Stage 逐级切：先切优化器状态、再切梯度、最后切参数。

⭐⭐ 最重要也最反直觉的一条：前两级的通信量和普通数据并行完全一样。 因为 AllReduce 本来就等于 ReduceScatter 加 AllGather——它只是把一次 AllReduce 拆成两半、中间插进「各自更新自己那一片」。⭐ 于是显存节省几乎是免费的，默认就该开。

第三级平时每卡只存一片，用到某一层时临时凑出完整参数、算完立刻丢——⭐ 峰值降到「分片参数加最大的那一层」，代价是通信量涨一半。

⭐ 预取是性能关键：算这一层的时候后台去取下一层，⚠️ 不预取 GPU 就干等，配置没调好可能比不切还慢。 ⚠️⭐ 最常见的错误是没设包装策略——整个模型被当成一个单元，显存几乎没省下来还会 OOM；必须按 Transformer 层包。

⭐ 多机的最优解是混合切分：机内全切（走高速互联）、机间用普通数据并行——显存降下来了，而跨机通信量和不切时一样。

⭐ 和张量并行的分工：这条路几乎不用改代码、能跨机；张量并行要改结构、基本不能跨机，⭐ 但单层都装不下时只有它能救。

⭐ 上手顺序：先开第二级 → 第三级 → 混合 → 再叠张量并行 → 最后流水线并行。 💥⚠️ 一个常见翻车点：保存 checkpoint 时 OOM——默认会把所有分片汇总到一张卡上，大模型必须用分片保存。

下一节 👉 12-张量并行.md