数据并行最浪费的地方不是算力，而是每张卡都存一份完整的优化器状态、梯度和参数。 ZeRO 把这三类状态依次切开：Stage 1 分优化器状态，Stage 2 再分梯度，Stage 3 连参数也分。 Stage 3 的显存按 Nd 下降，代价从 Stage 3 开始显现：论文算出每步通信量是 baseline DP 的 1.5 倍。
混合精度 Adam 下，每个参数要背 16 字节。 2Ψ 是 fp16 参数，2Ψ 是 fp16 梯度，剩下 12Ψ 是优化器状态：4Ψ fp32 主权重、4Ψ momentum、4Ψ variance（§3.1）。 1.5B 的 GPT-2 因此要 24 GB 才能放下模型状态，而 fp16 权重本身只有 3 GB。
ZeRO 不换优化器，也不改模型，只改“谁存哪一份”。 Stage 1 把 12Ψ 的优化器状态切成 Nd 份，每卡只留 1/Nd； Stage 2 再把梯度切成 1/Nd；Stage 3 把参数也切成 1/Nd，用到哪层再临时 all-gather 回来。 模型状态从 16Ψ 一路降到 16Ψ/Nd。
省显存不是免费的。 论文的通信账里，Stage 1/2 与 baseline DP 同为 2Ψ；Stage 3 需要前向、反向各一次参数 all-gather，再加一次梯度 reduce-scatter， 总计 3Ψ，是 baseline 的 1.5 倍（§1、§7.2.1、§7.2.2）。DeepSpeed 当前实现还在这条主线之外加了 CPU/NVMe offload、prefetch 与 bucket 化通信。
先建立一笔账：混合精度 Adam 下，一个参数在每张卡上要占 16 字节，其中权重只占 2 字节。
论文把训练显存分成两类：model states（优化器状态、梯度、参数）和 residual states（activations、临时 buffer、碎片）。大模型训练里，前者是主项。 Adam 在混合精度训练中除了 fp16 参数和梯度，还要保留 fp32 主权重、一阶矩 momentum 和二阶矩 variance（§3.1）：
前向和反向都要用的 fp16 权重。一张卡上没有它，这一步就算不了。
反向传播逐层产生；baseline DP 要对所有参数做一次全量 all-reduce。
fp32 主权重、momentum、variance 各 4Ψ。Adam 的“记忆”比权重本身贵 6 倍。
1.5B 参数的 GPT-2，按 16Ψ 计算至少需要 24 GB 模型状态；而只存 fp16 权重只要 3 GB。 论文指出，这样的模型在 32 GB 的 V100 上仍然训练不了，因为 residual states 还要占地方。 1.5B GPT-2 在 seq 1024、batch 32 下，activations 约 60 GB；activation checkpointing 能把它压到约 8 GB，代价是 33% 的重算；此外还有约 6 GB 的 fp32 梯度融合 buffer。 碎片也会让显存在“30% 以上仍空闲”时分配失败（§3.2）。
问题的核心于是变成：能不能让每张卡只存模型状态的一部分，同时保留 DP 的计算粒度和通信效率？ baseline DP 的做法是每张卡都存完整 16Ψ，再用 all-reduce 同步梯度；模型越大，这份冗余越贵。
理解 ZeRO 的取舍，先看它没有选的两条路：模型并行和流水线并行。
模型并行（MP）把每一层的参数和计算竖着切到多张卡。显存是省的，但计算粒度变细，层间还要频繁通信。 论文实测 40B 参数模型跨两个 DGX-2 节点跑 Megatron-LM，每张 V100 只有约 5 TFLOPs，不到硬件峰值的 5%。 跨节点带宽是硬约束：节点内 NVSwitch 每链路约 300 GB/s，跨节点 InfiniBand EDR 只有 12.5 GB/s。
流水线并行（PP）按层横着切，用 micro-batch 填 pipeline bubble。 它能省显存，但论文指出 PP 在 tied-weights、batch-normalization 等功能上难实现， 而且需要与流水线段数成比例的 batch size 来遮住 bubble，可能影响收敛（§2.1）。
CPU offload把状态搬到 CPU。论文引用的前作里，GPU-CPU-GPU 传输最多能吃掉 50% 的训练时间（§2.2.2）。 ZeRO 的判断是：DP 本身有更高的计算粒度、更低的通信量，问题只是“每张卡都存了完整状态”； 与其换并行方式，不如把 DP 的状态分片，同时保留 DP 的调度。
ZeRO-DP 不改变优化器、不改变模型结构，也不改变每步的数学结果；它只改变三件事的存放位置： 优化器状态、梯度、参数。分片之后，每个 rank 仍然按 DP 的方式算自己那份数据， 只在需要时通过 all-gather / reduce-scatter 把缺的部分补齐。
三个阶段是累积的。每多切一类状态，每卡显存再降一档，但参数的“临时拼回”也更频繁。
记模型有 Ψ 个参数、DP 度为 Nd、优化器状态乘数为 K（混合精度 Adam 为 12）。baseline DP 的每卡模型状态是 16Ψ。 三个阶段的公式直接来自 §5.1–§5.3：
每个 rank 只存 1/N_d 的 optimizer state，也只更新 1/N_d 的参数；step 末尾把更新后的参数 all-gather 回全量。梯度仍是全量。
梯度在反向时按 owner 做 reduce-scatter，每个 rank 只保留自己负责的那 1/N_d；参数仍全量，step 后 all-gather 更新。
每个 rank 只常驻 1/N_d 参数。做某一层前向/反向之前，把该层参数 all-gather 成临时全量，算完立刻释放。
把 16Ψ 拆开看，三个阶段的“地板”很清楚：Stage 1 至少留下 4Ψ（参数+梯度），Stage 2 至少留下 2Ψ（参数）， 只有 Stage 3 能把模型状态降到 16Ψ/Nd。这也是为什么论文说 Stage 3 让 DP 能训练“任意大小”的模型—— 前提是有足够多的卡来分摊（§5.3–§5.4）。
推断 精度切到 fp32、优化器切到 SGD 时，页面按同一字节记账替换 b 与 K；论文 Table 1 只给出混合精度 Adam（K=12）的组合， 这些扩展组合不在表内。默认组合（fp16/bf16 混合 + Adam）直接取自 §5.1–§5.3 的公式。 压缩比按所选精度/优化器的 Stage 0 记账：baseline = (2b + K)Ψ，b 为参数字节（混合精度 2、fp32 4）， K 为优化器状态乘数（Adam 8、SGD+m 4、SGD 0；混合精度另加 4 字节 fp32 主权重）。 本模拟只算 model states，未覆盖 activations、临时 buffer、碎片，也未模拟通信 overlap 与网络拓扑。
显示规则：≥100 GB 取整；10–100 GB 最多保留 1 位小数；<10 GB 最多保留 2 位小数； 整数不补零（例如 31.0 显示为 31，1.50 显示为 1.5）。GB 按 109 字节，不按 GiB。 分项先算后舍入，总计按未舍入值算，所以分项显示值相加可能与总计差 0.01–0.04。
Stage 1：4×7.5 + 12×7.5/64 = 30 + 1.40625 = 31.40625 → 31.4 GB
Stage 2：2×7.5 + 14×7.5/64 = 15 + 1.640625 = 16.640625 → 16.6 GB
Stage 3：16×7.5/64 = 1.875 → 1.88 GB
通信：baseline / Stage 1 / Stage 2 为 2×7.5×2 = 30 GB；Stage 3 为 3×7.5×2 = 45 GB。
| Stage | 参数 | 梯度 | 优化器状态 | 总计 | 通信（论文） | Stage 1（含 step 末 gather） |
|---|---|---|---|---|---|---|
| 0 | 15 | 15 | 90 | 120 | 30 | — |
| 1 | 15 | 15 | 1.41 | 31.4 | 30 | 45 |
| 2 | 15 | 0.23 | 1.41 | 16.6 | 30 | — |
| 3 | 0.23 | 0.23 | 1.41 | 1.88 | 45 | — |
单位 GB。Stage 1 的 45 GB 按论文 §5.1 / DeepSpeed 实现里 step 末的一次参数 all-gather 计入（+Ψ）；论文 §1、§7 的总账对 Stage 1 记 30 GB（见 04 节）。
Stage 1：4×128 + 12×128/64 = 512 + 24 = 536 GB
Stage 2：2×128 + 14×128/64 = 256 + 28 = 284 GB
Stage 3：16×128/64 = 32 GB
通信：baseline / Stage 1 / Stage 2 为 2×128×2 = 512 GB；Stage 3 为 3×128×2 = 768 GB。
| Stage | 参数 | 梯度 | 优化器状态 | 总计 | 通信（论文） | Stage 1（含 step 末 gather） |
|---|---|---|---|---|---|---|
| 0 | 256 | 256 | 1536 | 2048 | 512 | — |
| 1 | 256 | 256 | 24 | 536 | 512 | 768 |
| 2 | 256 | 4 | 24 | 284 | 512 | — |
| 3 | 4 | 4 | 24 | 32 | 768 | — |
A、B 两组的 Stage 1/2/3 总计分别对应 Table 1 的 7.5B 行与 128B 行（31.4 / 16.6 / 1.88，536 / 284 / 32），可逐项对照。
记账口径：b 为参数字节（混合精度 2、fp32 4），K 为优化器状态乘数 （Adam 8、SGD+m 4、SGD 0；混合精度另加 4 字节 fp32 主权重），Stage 0 = (2b + K)Ψ。
C1. Ψ = 7.5B，N_d = 64，fp32，SGD + momentum（b=4，K=4）：
Stage 0 (2×4+4)×7.5 = 90 GB；
Stage 1 2×4×7.5 + 4×7.5/64 = 60.46875 → 60.5 GB；
Stage 2 4×7.5 + (4+4)×7.5/64 = 30.9375 → 30.9 GB；
Stage 3 (2×4+4)×7.5/64 = 1.40625 → 1.41 GB；
压缩比 90 / 1.40625 = 64×。
C2. Ψ = 1.5B，N_d = 8，fp32，SGD（b=4，K=0）：
Stage 0 (2×4+0)×1.5 = 12 GB；
Stage 1 2×4×1.5 + 0 = 12 GB；
Stage 2 4×1.5 + 4×1.5/8 = 6.75 GB；
Stage 3 (2×4+0)×1.5/8 = 1.5 GB；
压缩比 12 / 1.5 = 8×。
C 组不在论文 Table 1 内，是本页按 (2b + K)Ψ 的 Stage 0 记账外推；与交互模拟使用同一公式，显示值可直接对照。
分片省了显存，问题是每一步要多搬多少字节。论文把通信量拆成 all-gather 与 reduce-scatter 两笔账。
baseline DP 的梯度同步是一次 all-reduce。按 §7.1 的拆法，all-reduce = reduce-scatter Ψ + all-gather Ψ， 所以每步每卡的数据移动量是 2Ψ。Stage 2 把梯度 all-reduce 换成 reduce-scatter（Ψ）， 参数更新后再 all-gather 一次（Ψ），合计仍是 2Ψ，与 baseline 相同（§7.2.1）。
Stage 3 多出来的那一笔，是参数的第二次 all-gather：前向按层把参数拼回来，反向按逆序再拼一次。 这两次 gather 不必同时存在——某一层算完就可以释放，所以显存不会回到 16Ψ，但通信账上是实打实的 3Ψ。 论文强调这不是“把整个模型同时广播”，而是把 gather 摊到整个前向/反向过程中（§7.2.2）。
论文 §1 对 Pos 的表述是 “4x memory reduction, same communication volume as DP”，
§7 的分项通信账也只把 Pos+g 记入 2Ψ、把 Pos+g+p 记入 3Ψ。
但论文 §5.1 在描述 Pos 时，自己写了 optimizer step 末尾的一次参数 all-gather；
DeepSpeed 的实现（all_gather_dp_groups）与 §5.1 的流程一致。
若按字节把这次 gather 计入，Stage 1 是 3Ψ；§1 与 §7 的总账仍记 2Ψ。三处来源并列如下，本页不替原文裁决。
§1 与 §7 的总账：Pos、Pos+g 都不增加通信量；只有 Pp 额外做参数 gather，总计 1.5×。
§5.1 与 stage_1_and_2.py:3180 都在 optimizer step 末尾调用参数 all-gather。Stage 2 用 reduce-scatter Ψ + 参数 all-gather Ψ 凑成 2Ψ；Stage 1 的梯度 all-reduce 已经占满 2Ψ，这次 gather 是额外的 Ψ。
实现层面，DeepSpeed 把这些 collective 做成 bucket 化、可 overlap 的调用：
Stage 2 的梯度按 partition 归并后 reduce-scatter；Stage 3 的梯度用
reduce_scatter_coalesced（deepspeed/runtime/zero/stage3.py:1948），
参数 gather 则走 deepspeed/runtime/zero/partition_parameters.py:1723 的 _all_gather，
由 coordinator 在每层 forward 前触发（deepspeed/runtime/zero/partitioned_param_coordinator.py:310）。
前向用到哪层，就把哪层的参数 all-gather 回来；算完立刻释放。反向按逆序再来一遍。
fetch_sub_module 在子模块 forward 前触发参数 all-gather
（partitioned_param_coordinator.py:310），并把后续几层放进 prefetch 队列（同文件 452–484 行）。
参数按反向顺序再次 all-gather；梯度按 partition 归并后 reduce-scatter
（stage3.py:1948），每个 rank 只保留自己负责的梯度分片。
optimizer step 只更新本地 fp32 参数分片与优化器状态（stage3.py:2913）；
下一步 forward 需要别的层时再 gather。常驻参数上限由 stage3_max_live_parameters 控制（默认 109 个元素，deepspeed/runtime/zero/config.py:245）。
源码锚点（DeepSpeed f47dc3e）：deepspeed/runtime/zero/stage3.py:152（Stage 3 类）、
:1282（_partitioned_params_swap_out，把用完的参数换出）、
:1948（reduce_scatter_coalesced）；
deepspeed/runtime/zero/partition_parameters.py:1723（_all_gather）、
:1765（_partition）、:1778（_partition_param 释放数据）。
Table 1 是每卡模型状态的解析结果，Table 2 把“理论最大模型”与“实测能跑的模型”放在一起。
| N_d | 7.5B | 128B | 1T | ||||||
|---|---|---|---|---|---|---|---|---|---|
| P_os | P_os+g | P_os+g+p | P_os | P_os+g | P_os+g+p | P_os | P_os+g | P_os+g+p | |
| 1 | 120 | 120 | 120 | 2048 | 2048 | 2048 | 16000 | 16000 | 16000 |
| 4 | 52.5 | 41.3 | 30 | 896 | 704 | 512 | 7000 | 5500 | 4000 |
| 16 | 35.6 | 21.6 | 7.5 | 608 | 368 | 128 | 4750 | 2875 | 1000 |
| 64 | 31.4 | 16.6 | 1.88 | 536 | 284 | 32 | 4187 | 2218 | 250 |
| 256 | 30.4 | 15.4 | 0.47 | 518 | 263 | 8 | 4046 | 2054 | 62.5 |
| 1024 | 30.1 | 15.1 | 0.12 | 513 | 257 | 2 | 4011 | 2013 | 15.6 |
Table 1 已完整转录。原文用加粗标出能装入 32 GB V100 的组合，本页未复刻加粗，可按数值读出。
N_d=64 时，7.5B 模型的 Stage 1 是 31.4 GB，刚好在 32 GB 以内；Stage 2 降到 16.6 GB，Stage 3 只要 1.88 GB。 同样 64 张卡，128B 模型只有 Stage 3 的 32 GB 能装下；Stage 1/2 分别是 536 GB 与 284 GB。 1T 模型更极端：N_d=1024 时 Stage 3 每卡 15.6 GB，Stage 1/2 仍有 4011 GB 与 2013 GB。 这不是线性缩放失效，而是 Stage 1 的 4Ψ 与 Stage 2 的 2Ψ 地板在模型变大时先撞墙。
| MP | GPUs | Max Theoretical Model Size | Measured Model Size | ||||
|---|---|---|---|---|---|---|---|
| Baseline | P_os | P_os+g | P_os+g+p | Baseline | ZeRO-DP (P_os) | ||
| 1 | 64 | 2B | 7.6B | 14.4B | 128B | 1.3B | 6.2B |
| 2 | 128 | 4B | 15.2B | 28.8B | 256B | 2.5B | 12.5B |
| 4 | 256 | 8B | 30.4B | 57.6B | 0.5T | 5B | 25B |
| 8 | 512 | 16B | 60.8B | 115.2B | 1T | 10B | 50B |
| 16 | 1024 | 32B | 121.6B | 230.4B | 2T | 20B | 100B |
Table 2 已完整转录。条件是按 MP 度 × DP 度的组合；P_os 的实测值与理论最大值吻合，说明内存分析可作为现实上界。
论文的实现叫 ZeRO-100B，只覆盖 Pos+g + ZeRO-R，不是完整 Stage 3（§10）。 它能在 400 张 V100 上跑 170B 参数，8B–100B 区间平均 15 PFLOPs，超过峰值 30%， 相对基线最多 10× 加速；不加 MP 时，128 张卡能跑 13B 参数，平均每卡超过 40 TFLOPs，而 baseline DP 最大只能跑 1.4B、每卡不到 20 TFLOPs。 Table 1 的 1T 行是解析结果，不是 1T 实测；Table 2 的实测列也只到 P_os 的 100B。把这两类数字分开看，才不会把“能装下”误读成“已经训练过”。
把 ZeRO 放回完整训练系统里看：它解决了显存的主项，但剩下的项和实现代价仍然决定能不能跑、跑得多快。
Stage 1–3 只切优化器状态、梯度、参数。activations、临时 buffer、碎片仍要占显存； 论文用 ZeRO-R 的 Pa、CB、MD 处理这些项（§6），但它们不在本页的 16Ψ 账里。 显存计算器给出的数字只是 model states 的下界。
3Ψ 是字节账。实际时间取决于 bucket 大小、prefetch 是否盖住 gather、跨节点带宽和 batch size。 论文说 Stage 3 的 gather 可以摊到整个前向/反向（§7.2.2），但没承诺在所有网络上都能完全 overlap。
Stage 1/2 有显存地板。无论 N_d 多大，Stage 1 至少留 4Ψ、Stage 2 至少留 2Ψ；只有 Stage 3 是 16Ψ/N_d。 1T 模型在 1024 张卡上，Stage 1/2 的每卡仍要 4011 GB / 2013 GB，只有 Stage 3 的 15.6 GB 落在单卡范围。 所以“ZeRO 能训练任意大小模型”这句话，准确说是“Stage 3 在卡足够多时能”。
offload 是另一条轴，不是论文 ZeRO-DP 的默认路径。
论文 §2.2.2 明确说 ZeRO 不把模型状态存到 CPU；只有 ZeRO-R 在少数情况下把 activation checkpoint 放到 CPU（§6.1）。
DeepSpeed 当前实现把 offload 做成配置项：offload_optimizer 支持 Stage 1/2/3，
offload_param 只支持 Stage 3（deepspeed/runtime/zero/config.py:164–177），
设备可以是 CPU 或 NVMe（deepspeed/runtime/engine.py:1581–1588；NVMe swap 见 deepspeed/runtime/zero/stage3.py:828）。
这能进一步降显存，但把瓶颈移到 PCIe / NVMe 带宽，和论文的 1.5× 通信账是两笔不同的成本。
PyTorch FSDP 文档写明它 “is inspired by Xu et al. as well as the ZeRO Stage 3 from DeepSpeed”。
两者都把参数、梯度、优化器状态分片到 data-parallel rank，并在使用前 all-gather 参数；
差别在实现单位：FSDP 默认把 module 的参数收进内部 FlatParameter 管理，
DeepSpeed ZeRO-3 以参数/分区为单位，并配 prefetch 与 stage3_max_live_parameters 限制常驻量。
这是同一设计目标下的两套工程实现，不是同一份代码。
引用：PyTorch docs, FullyShardedDataParallel（“A wrapper for sharding module parameters across data parallel workers. This is inspired by ... the ZeRO Stage 3 from DeepSpeed.”）
| 维度 | 论文（arXiv 1910.02054v3） | DeepSpeed f47dc3e |
|---|---|---|
| Stage 1 通信 | §1、§7 总账记 2Ψ（“same communication volume as DP”）；§5.1 描述 step 末参数 all-gather | 对应 §5.1：step 末 all_gather_dp_groups（stage_1_and_2.py:3180），按字节另计 Ψ |
| Stage 3 参数驻留 | 按层 broadcast / discard，gather 摊到前向与反向（§7.2.2） | hook + coordinator fetch/prefetch（deepspeed/runtime/zero/partitioned_param_coordinator.py:310,452）；stage3_max_live_parameters 默认 109 元素（deepspeed/runtime/zero/config.py:245） |
| 模型状态 offload | ZeRO-DP 不把 model states 放 CPU；只有 Pa+cpu 放 activation checkpoint（§2.2.2、§6.1） | offload_optimizer 支持 Stage 1/2/3；offload_param 仅 Stage 3；CPU 或 NVMe（deepspeed/runtime/zero/config.py:164–177） |
| 评测范围 | ZeRO-100B = Pos+g + ZeRO-R，最高 170B（§10）；1T 为解析 | 实现 Stage 0–3，含 offload、NVMe、prefetch 等工程路径 |
什么时候别用 Stage 3：