ZeRO 优化器：把显存“分”着用
说到大模型训练，显存永远是个绕不开的话题。我早期做分布式训练时，最头疼的就是显存爆炸。明明单卡能跑的小模型，一上多卡反而跑不起来了。后来接触了 ZeRO，才明白问题出在哪——数据并行里藏着大量冗余。
ZeRO 的核心思想很简单：既然每张卡都存一份完整的数据，那不如分着存。你想想看，数据并行时，每张卡都有完整的模型参数、梯度和优化器状态。这就像四个人各背一桶水，其实每人背四分之一就够了。
ZeRO 分三个 Stage，一个比一个“狠”。我们一个个来看。
ZeRO Stage 1：优化器状态分片
Stage 1 是最温和的。它只对优化器状态下手。
什么是优化器状态？以 Adam 为例，它需要保存每个参数的 momentum 和 variance。这两个东西加上参数本身，显存占用大概是：
- 参数：4 字节（FP32）
- 梯度：4 字节
- momentum：4 字节
- variance：4 字节
总共 16 字节。其中优化器状态占了 8 字节，也就是一半。
Stage 1 的做法是：把优化器状态切分到各张卡上。每张卡只维护自己那部分参数的 momentum 和 variance。更新参数时，各卡只更新自己负责的那部分，然后通过通信把更新后的参数广播出去。
显存节省：约 50% 的优化器状态显存。整体显存节省约 25%-30%。
我在项目中遇到过这种情况：一个 7B 模型用 8 卡训练，单卡显存 80GB。不开 ZeRO 时，光优化器状态就占了 28GB。开了 Stage 1，每卡只存 3.5GB，瞬间轻松不少。
小提示：Stage 1 几乎不影响计算效率。通信量只增加了一次 all-gather（收集更新后的参数）。如果你的模型已经能跑起来，只是想省点显存加 batch size，Stage 1 是最安全的选择。
ZeRO Stage 2：梯度分片
Stage 2 在 Stage 1 的基础上，把梯度也分片了。
数据并行中，每张卡计算完自己的 batch 后，会得到一份完整的梯度。然后通过 all-reduce 把梯度汇总。Stage 2 的做法是：不存完整梯度，只存自己负责的那部分。
具体流程是这样的：
- 前向传播正常进行
- 反向传播计算梯度
- 各卡只保留自己分片对应的梯度，其他梯度通过 reduce-scatter 通信丢弃
- 用分片后的梯度更新分片后的优化器状态
说白了，Stage 2 把梯度的存储和通信都做了优化。原来 all-reduce 需要先求和再广播，现在 reduce-scatter 只求和一次，省了一半通信量。
显存节省：相比 Stage 1，额外省掉了梯度的存储（每卡 4 字节 × 参数量）。整体显存节省约 50%。
我曾经在一个 13B 模型上试过 Stage 2。单卡 80GB 显存，不开 ZeRO 时 batch size 只能设 1。开了 Stage 2，batch size 直接翻到 4。训练速度反而快了，因为更大的 batch 让 GPU 利用率更高。
注意：Stage 2 需要修改反向传播的逻辑。如果你的代码用了 PyTorch 的 DDP，需要换成 DeepSpeed 或 Megatron-LM 的实现。自己手写的话，坑不少。
ZeRO Stage 3：参数分片
Stage 3 是最激进的——连模型参数都分片。
这意味着，每张卡只存一部分参数。前向传播时，需要用到哪些参数，就从对应的卡上拉过来。计算完再丢掉。反向传播同理。
你想想看，这相当于把模型“打散”了。每张卡像是一个分布式存储节点，随时按需取用参数。
流程大致是：
- 前向传播时，每层计算前先 all-gather 该层的参数
- 计算完后，丢弃不属于自己分片的参数
- 反向传播时，再次 all-gather 参数
- 计算梯度后，用 reduce-scatter 分发梯度
显存节省：理论上可以省掉参数、梯度、优化器状态的全部冗余。整体显存节省约 75%。
但代价也很明显——通信量暴增。每层都要做 all-gather 和 reduce-scatter，通信开销是 Stage 2 的好几倍。我试过在 64 卡集群上跑 Stage 3，网络带宽成了瓶颈。训练速度反而比 Stage 2 慢了 30%。
我的建议：Stage 3 适合模型大到单卡完全放不下的场景。比如 70B、100B 以上的模型。如果你的模型单卡能放下，优先用 Stage 1 或 Stage 2。
ZeRO-Offload：把显存“借”给 CPU
ZeRO-Offload 是 ZeRO 的一个扩展。它把优化器状态和梯度放到 CPU 内存里，GPU 只保留参数和计算中间结果。
为什么要这么做？因为 CPU 内存便宜啊。一台机器 512GB 内存很常见，但 GPU 显存 80GB 就顶天了。把优化器状态这种“计算少、存储多”的东西放到 CPU 上，能省出大量显存。
具体做法：
- 前向/反向传播在 GPU 上完成
- 梯度计算完后，异步拷贝到 CPU
- CPU 负责更新优化器状态
- 更新后的参数再拷贝回 GPU
这里有个关键点：CPU 更新参数比 GPU 慢很多。所以 ZeRO-Offload 用了异步流水线，让 CPU 和 GPU 并行工作。GPU 算下一轮的 forward，CPU 同时更新上一轮的参数。
显存节省：优化器状态和梯度全部 offload 到 CPU，GPU 只存参数。对于 Adam 优化器，显存节省约 50%。
我曾经在一个 33B 模型上试过 ZeRO-Offload。单卡 40GB 显存，不开 offload 时 batch size 只能设 2。开了 offload，batch size 直接拉到 8。虽然每步训练时间长了 20%，但整体吞吐量反而提升了，因为 batch 大了收敛更快。
避坑指南：我曾经遇到过 CPU 内存不够的情况。ZeRO-Offload 会把优化器状态放到 CPU，如果你的模型很大（比如 100B），CPU 内存至少需要 200GB。另外，CPU 和 GPU 之间的 PCIe 带宽也是瓶颈。建议用 NVLink 或 InfiniBand 连接，否则 offload 的延迟会让你怀疑人生。
核心逻辑图
下面这张图展示了 ZeRO 三个 Stage 和 Offload 的显存分布变化：
从图上可以清楚看到：Stage 1 只分片优化器状态，Stage 2 加上梯度分片，Stage 3 连参数都分片了。Offload 则是把部分数据挪到 CPU 内存。
怎么选？
我个人的经验是：
- 模型能单卡放下：用 Stage 1，省点显存加 batch size
- 模型略超显存：用 Stage 2，省一半显存
- 模型远超显存：用 Stage 3，但要注意网络带宽
- CPU 内存充裕：加 Offload，进一步省显存
嗯，ZeRO 的核心就是“分而治之”。把冗余的数据分到多张卡上，用通信换显存。至于怎么换、换多少，就看你的模型大小和硬件条件了。
想要更多资料，联系可可