ZeRO 2020 — 把训练状态切成 N 份让万亿参数成为可能

待复核

是什么

ZeRO（Zero Redundancy Optimizer）是微软 2020 年在超算顶会 SC 提出的一组显存优化方法：训练一个大模型时，把每张 GPU 都重复存的那份”训练状态”切成 N 份，每张只留 1/N。日常类比：原本 N 个学生每人都抄一整本笔记，现在改成每人只保管 1/N 页，要看时再相互借。

论文里”万亿参数”不是噱头——它给出的内存账本说明：1024 张 V100（每张 32GB）配合 ZeRO-3，理论上能装下 1 万亿参数的模型训练。这是首次有人把”训万亿”这件事算到能落地的程度。

为什么重要

现代大模型训练框架几乎都是 ZeRO 的徒孙。理解这篇论文能帮你看懂：

为什么 PyTorch 的 fsdp-2023 与 ZeRO Stage 3 语义等价（实现独立，不是代码 fork）
为什么 megatron-lm 和 DeepSpeed 经常组合出现：一个切计算、一个切状态，正交互补
为什么 gpipe-2019 / pipedream-2019 这些流水线并行方案常被叠加在 ZeRO 之上，凑成 “3D Parallelism”
为什么大模型微调脚本里 zero_optimization.stage 这个旋钮值得花时间调

核心要点

ZeRO 的洞察是：训练状态有三块，逐块切分，每切一块多省一截内存。

论文给出的内存账本（以 7.5B 参数模型 + 64 张 GPU 为例，单位 GB/卡）：

配置	优化器状态	梯度	参数	总计
标准 DDP	120	30	15	165
ZeRO Stage 1	1.9	30	15	47
ZeRO Stage 2	1.9	0.5	15	17
ZeRO Stage 3	1.9	0.5	0.2	2.6

三个 Stage 逐级切：

Stage 1（Pos）：切优化器状态
- Adam 训练时，每个参数额外存 momentum + variance（fp32，共 8 字节/参数）
- 加上 fp32 master copy（4 字节）总共 12 字节/参数——是模型参数本身的 6 倍
- 切完每张卡只保留 1/N 份，省 4 倍以上
Stage 2（Pos+g）：再切梯度
- 关键技巧：把传统 all-reduce 替换为 reduce-scatter
- 每张卡只拿到自己负责那段的梯度，不收集完整梯度
- 类比：4 个学生改卷，原本每人都拿到 4 份完整成绩，现在每人只拿自己那一份
Stage 3（Pos+g+p）：连参数本身也切
- forward / backward 时按层动态 all-gather：算这层就把参数借齐，算完立刻丢
- 通信量比 Stage 2 多 50%，但内存压到 1/N
- 类比：每人只背 1/4 课本，上课讲到哪一页谁背就传给谁

关键设计：Stage 2 用 reduce-scatter 替代 all-reduce 的等式 all-reduce = reduce-scatter + all-gather 是 ZeRO 的灵魂——把一次 collective 拆两步，不增加通信量却能”白拿”一份内存优化。

为什么是 12 字节/参数：fp16 训练时，一份参数本身是 fp16（2 字节），但 Adam 优化器要保留 fp32 master copy（4 字节）+ fp32 momentum（4 字节）+ fp32 variance（4 字节）共 12 字节。这就是 OS 体积是参数本身 6 倍的来源——也是 Stage 1 收益最大的原因。

实践案例

案例 1：用账本算清「万亿能不能装下」

论文把 1T 参数 + 1024×V100(32GB) 写成可核对的内存账：

每卡训练状态 ≈ (1e12 参数 × 16 字节) / 1024 ≈ 16 GB
余下 ≈ 16 GB 留给 activation / 通信缓冲

逐部分解释：16 字节 约等于 fp16 参数 + Adam 的 fp32 状态摊到每参数的均值；/ 1024 是 Stage 3 把 OS、梯度、参数都切成 N 份。同一套账本下，论文还报了 100B 模型在 400 张 V100 上约 38 TFlops/卡——说明「能装下」之外吞吐也可落地。

案例 2：4×24GB 卡微调 7B（DeepSpeed + accelerate）

{
  "fp16": { "enabled": true },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu" }
  }
}

# 需已安装 deepspeed，且 train.py 用 accelerate 包装 Trainer
accelerate config    # 选 DeepSpeed，指向上面的 ds_config.json
accelerate launch train.py

逐部分解释：stage: 3 切 OS+梯度+参数，单卡不再存完整 7B；offload_optimizer 把 Adam 状态丢到 CPU，显存再降一截，但吞吐会明显变慢——适合微调，不适合大规模预训练。不开 ZeRO 时 7B+Adam 每卡常要 ~80GB 级完整副本；开 Stage 3+offload 后 4 张 RTX 4090（24GB）也能跑通。单机多卡微调若显存尚可，Stage 2 往往是更快的甜点区。

案例 3：Megatron 切计算 + ZeRO 切状态

一层 Attention → Megatron tensor-parallel 拆到多卡算
跨副本的 OS / 梯度 / 参数 → DeepSpeed ZeRO 分摊

逐部分解释：Megatron-Turing NLG 530B 把「算」和「存」正交叠加——前者降单层算力墙，后者降每卡状态墙。这是后来 Data × Tensor × Pipeline × ZeRO 的 3D/4D 并行起点；只开 ZeRO 不够时，再叠 TP/PP。

踩过的坑

小模型上 Stage 3 反而更慢：模型 < 1B 参数时，多一次 all-gather 的通信开销 > 内存收益。这种情况用 Stage 1 或纯 DDP 才对。
ZeRO 不解决 activation 内存：训练时 activation 常占总内存 40%+，ZeRO 完全切不动这块——必须配合 gradient checkpointing 一起用。
跨 node 比单 node 慢得多：ZeRO 频繁的 all-gather 在单 node NVLink（900GB/s）上很快，跨 node InfiniBand（约 25GB/s）就成瓶颈。1024 GPU 训练时 Stage 3 通信占比能到 30%+。
配置参数多，新人容易调爆：stage / offload / bucket_size / overlap_comm 一堆旋钮，调不好可能比 DDP 还慢。这也是 PyTorch 推出 fsdp-2023（更易用版 Stage 3）来收编社区的原因之一。

适用 vs 不适用

适用：

单卡装不下完整模型，但有多张 GPU（4–8 张消费级卡微调 7B–70B）
用 Adam / AdamW（OS 是大头；单机多卡微调常先试 Stage 2）
PyTorch / HuggingFace Transformers + DeepSpeed / FSDP 生态

不适用：

单卡能装下整个模型（< 1B）→ 纯 DDP 通常更快
用无状态/弱状态优化器（如朴素 SGD）→ Stage 1 收益接近零
推理 → ZeRO 只管训练；推理用量化 / vllm 等
千卡级跨 node 只开 Stage 3 → 通信占比过高，需叠 TP / PP

历史小故事（可跳过）

2019 年 10 月：微软 DeepSpeed 团队挂出 ZeRO 论文 arXiv 预印本
2020 年 11 月：在 SC 2020 正式发表，DeepSpeed 开源到 GitHub 立刻成为大模型训练事实标配
2021–2022 年：同作者发 ZeRO-Infinity（NVMe offload）与 ZeRO++（量化通信 / hierarchical partitioning）
2023–2024 年：fsdp-2023 与 Stage 3 语义对齐；FSDP2 / 主线默认路径巩固，ZeRO 思想成事实标准

学到什么

“复制是必须的”可被挑战——DDP 全量复制是惯例不是物理定律；规模一大，复制本身就是瓶颈。
all-reduce = reduce-scatter + all-gather 总成立——拆一步 collective，就能”白拿”一份内存优化。
工程化的胜利：ZeRO 没发明新算法，只是把切分贯彻到 OS、梯度、参数三层——极致组合往往比新算法更有影响力。
算账比写代码更重要：那张内存账本把”万亿训练”从口号变成可核对的工程目标；通信换内存没有银弹。

关联

megatron-lm —— Megatron 切计算 / ZeRO 切状态，正交互补
fsdp-2023 —— 与 ZeRO Stage 3 语义等价的 PyTorch 原生分片
gpipe-2019 —— 流水线并行常和 ZeRO 叠加
pipedream-2019 —— 1F1B 流水线，3D Parallelism 的另一维
alpa-2022 —— 把数据/张量/流水/ZeRO 统一成搜索问题

反向链接

（暂无反向链接）