AI 原生开放教材

外观

AI 工程从零到一（中文）源码研读·中文原文 ↗原件 ↓仓库 ↗

ZeRO Optimizer State 分片 ​

Adam 给每个参数存两个 moment 估计，都是 float32。一个 70 亿参数的模型背着 56 GB 的 optimizer state。ZeRO stage 1 把这部分 shard 到 N 个 rank 上；每个 rank 拥有 1/N 的 optimizer。本地 step 之后，更新后的参数 shard 再 broadcast 回来，每个 rank 重建出完整模型，下一步开始。这一招换来的是：把训练栈里最大那块单次分配的内存做了线性下降。

类型： Build 语言： Python 前置要求： 阶段19 Track C 第42-49课 预计时间： ~90 分钟

学习目标 ​

  • 把 optimizer state（first moment、second moment、fp32 master copy）shard 到 N 个 rank 上，让每个 rank 拥有 1/N。
  • 用 reduce_scatter 只把每个 rank 自己那份 shard 的梯度和投递给它，再用 allgather 把更新后的参数 shard 广播回来。
  • 算出 stage 1、stage 2、stage 3 相对原版 DDP 的内存节省表。
  • 从模型规模和带宽预算出发，论证在 stage 1、stage 2、stage 3 之间该怎么选。

问题背景 ​

原版 DDP 把一切都复制一遍：参数、梯度、optimizer state 在每个 rank 上都完整存在。对一个 fp16 的 70 亿参数模型，这意味着每个 rank 14 GB 参数、14 GB 梯度、28 GB optimizer state。Optimizer state 是最大的那一项，也是最容易 shard 的，因为它只在 step 时被碰，forward 和 backward 时都不碰。

ZeRO stage 1 把 optimizer state shard 掉。每个 rank 持有 1/N 的 Adam moment。Backward 之后，ZeRO 不再 allreduce 完整梯度再本地 step，而是用 reduce_scatter 让每个 rank 只收到自己那份 shard 的求和梯度。这个 rank 把 optimizer step 应用到它那份 master parameter 的 shard 上。更新后的参数 shard 再 allgather 回来，于是每个 rank 都有完整模型供下一次 forward。Optimizer 内存按 N 下降。每步的线路流量和 DDP 一样：一次 reduce_scatter 加一次 allgather 在带宽上等于一次 allreduce。内存赢了，吞吐不变。

核心概念 ​

mermaid
flowchart TD
  A[在完整模型上 forward + backward] --> B[每个 rank 上梯度都算完]
  B --> C[reduce_scatter 梯度]
  C --> D[rank r 持有求和梯度 shard r]
  D --> E[用本地 optimizer state 对 shard r 做 Adam step]
  E --> F[更新后的参数 shard r]
  F --> G[allgather 参数 shard]
  G --> H[下一次 forward 又看到完整模型]

ZeRO 的各个 stage ​

Stageshard 了什么每 rank 内存每步通信
DDP什么都不params + grads + optim1 次 allreduce
ZeRO-1optimizer stateparams + grads + optim/N1 次 reduce_scatter + 1 次 allgather
ZeRO-2optim + gradsparams + grads/N + optim/N1 次 reduce_scatter + 1 次 allgather
ZeRO-3optim + grads + paramsparams/N + grads/N + optim/N每层 1 次 allgather + 每层 1 次 reduce_scatter

Stage 1 是最便宜的赢面，因为 optimizer state 在预算里占大头。Stage 2 需要梯度 shard 的累积逻辑，但带宽一样。Stage 3（FSDP）为每次 forward 和 backward 付每层通信的代价，换来参数 shard 的内存下降。本课完整实现 stage 1。

内存计算，真实数字 ​

对一个有 P 个参数、用 Adam 做混合精度训练的模型：

项原版ZeRO-1为什么
fp16 参数2P 字节2P 字节forward 需要
fp16 梯度2P 字节2P 字节backward 需要
fp32 master copy4P 字节4P/N 字节只有 optimizer 用它
fp32 first moment4P 字节4P/N 字节只有 optimizer 用它
fp32 second moment4P 字节4P/N 字节只有 optimizer 用它
合计16P 字节4P + 12P/N 字节

在 N=8 时：原版 16P，ZeRO-1 5.5P，降了 65%。在 N=64 时：原版 16P，ZeRO-1 4.19P，降了 74%。

为什么 reduce_scatter 胜过"allreduce 再 shard" ​

Allreduce 给每个 rank 完整的求和梯度。如果你只需要 shard r，那被归约的梯度里有 (N-1)/N 在 rank r 上是白算的。Reduce_scatter 精确地投递每个 rank 拥有的那份 shard；每 rank 字节数和 allreduce 一样（因为 allreduce 就是 reduce_scatter + allgather），但后半段被换成了之后的参数 shard allgather。净线路流量和 DDP 完全相同，内存被除以了 N。

figure
cd-zero-shard

动手构建 ​

code/main.py 实现了：

  • flatten_params(module) 和 unflatten_into(module, flat)，把模型参数打包进一个连续张量再解包回去。这种扁平布局让按 rank 做 shard 变成一次简单的切片。
  • ZeroOptimizer(model, world_size, rank, lr)，持有该 rank 那份 master copy 和 Adam moment 的 shard。
  • step()，在扁平梯度上跑 reduce_scatter，对该 rank 的 shard 应用 Adam，再 allgather 把更新后的参数取回来。
  • 一个 demo，把一个 3 层 MLP 训 20 步，并打印每步的内存预算，旁边附上原版 DDP 基线。

运行：

bash
python3 code/main.py

输出：每步 loss，以及一张内存表，展示 ZeRO-1 在每个 rank 上只持有 1/N 的 optimizer state，而 DDP 持有完整副本。

真实世界中的生产模式 ​

有三个模式能把 ZeRO 打磨到可以上线。

sharded checkpoint 很关键。 ZeRO-1 的 optimizer state 散在各个 rank 上；checkpoint 必须记录哪个 rank 拥有哪部分。第 80 课构建的就是那个 sharded checkpoint manifest，让一次 ZeRO 运行能在同样 world size 上恢复。没有它，存下来的状态在重启时根本读不了。

混合精度才是重点。 ZeRO 是一种混合精度技术；被 shard 的正是那份 fp32 master copy。不用混合精度跑 ZeRO，就会付 fp32 master 的内存税，却拿不到对应的 fp16 forward 收益。生产运行总是把 ZeRO 和 autocast 或 bf16 权重配在一起。

stage 1 是近乎免费的赢面。 通信在带宽上和 DDP 完全一样。内存节省与 N 成线性。唯一的代价是 optimizer shard 的记账。生产栈默认上 stage 1，除非参数 shard 的内存也成了问题；那时再用 stage 2 或 3 拿通信换内存。

实际使用 ​

生产模式：

  • DeepSpeed ZeRO。 参考实现。deepspeed_config.json 选 stage 1/2/3 和 partition 大小。
  • PyTorch FSDP。 PyTorch 原生的等价物。ShardingStrategy.SHARD_GRAD_OP 是 ZeRO-2；FULL_SHARD 是 ZeRO-3。
  • HuggingFace Accelerate。 用统一配置把 DeepSpeed 和 FSDP 都包起来。

拿去用 ​

第 79 课（pipeline parallel）是正交的另一个 sharding 轴：它不是在同一个模型上 shard optimizer state，而是把层 shard 到各个 rank 上。第 81 课在端到端 demo 上把 DDP + ZeRO 组装起来。

练习 ​

  1. 扩展到 ZeRO-2，做梯度 shard：每个 rank 只存自己那份 shard 的梯度，做法是 backward 之后把非 shard 部分清零。
  2. 加一个内存 profiler，在 rank 0 上打印实际的 fp32 字节用量，对比公式预测。
  3. 测量原版 DDP 与 ZeRO-1 的每步墙钟时间，并拆解成 forward、backward、通信。
  4. 在 ZeRO-1 下实现梯度裁剪：L2 norm 必须通过对本地 norm 平方做 allreduce 来跨所有 shard 计算。
  5. 实现一个用 allreduce 而不是 reduce_scatter 的"朴素 ZeRO"，测量线路时间差。用数字论证为什么选 reduce_scatter。

关键术语 ​

术语大家怎么说实际含义
ZeRO-1"shard optimizer"每个 rank 持有 1/N 的 fp32 master + Adam moment
ZeRO-2"梯度也 shard"reduce_scatter 之后每个 rank 还会丢掉非 shard 的梯度
ZeRO-3"shard 参数"每个 rank 持有 1/N 的 fp16 参数；forward 时每层 allgather
Master copy"fp32 权重"optimizer 更新的那份高精度参数副本
Reduce_scatter"把和拆开"只把每个 rank 自己那份 shard 的求和梯度投递给它

延伸阅读 ​