你有个 7B 的模型，想在一张 48G 的卡上微调。单卡放不下，于是你上了数据并行：8 张卡，每张卡分不同的 batch。结果显存还是爆——因为每张卡都塞了一整套模型、梯度和优化器状态，加起来 112GB。为什么数据并行这么“浪费”？因为它把模型副本在每个 GPU 上都存了一份，而训练时这些副本的内容其实是一样的。
ZeRO 就是来消除这种冗余的。它来自微软 DeepSpeed 团队，全称是 Zero Redundancy Optimizer，零冗余优化器。思路很简单：既然所有 GPU 持有同样的优化器状态，为什么不把它们切成 N 份，每张卡只拿一份？原论文里把这个想法称为对数据并行的“存储重构”。
先算一笔账：显存到底被谁吃掉了？
混合精度训练 7B 模型，每个参数需要 16 字节：
- fp16 参数：2 字节
- fp16 梯度：2 字节
- fp32 参数副本：4 字节
- Adam 一阶动量：4 字节
- Adam 二阶动量：4 字节
合计 7B × 16B = 112GB。其中优化器状态（后三样）占 12 字节，总共 84GB，比例高达 75%。数据并行下，每张卡依然要全量存这些。唯一的好处是计算被摊到多卡上，显存一点没省。
Stage 1：优化器分片，先拆掉 75% 的冗余
训练中，每步更新后的优化器状态在所有 GPU 上是完全一致的。既然一致，就不必要每人存一份完整的。Stage 1 把优化器状态按参数范围切块，平均分给各卡。
每一步是这样工作的：
- 每卡做自己的前向和反向，因为参数是全量副本，所以能算出完整的梯度。
- 对梯度做 all-reduce，得到所有 batch 的平均梯度。
- 每卡用平均梯度中自己负责的那一块，更新自己的优化器分片，算出新参数。
- 通过 all-gather 把新参数广播给所有卡，保证下一轮前向参数一致。
这一步之后，梯度和参数仍然是全量的，但优化器状态从 84GB 降到 84/N。N=8 时，每卡显存大约从 112GB 降到 38.5GB，48G 卡就能跑了。代价是多了一次参数广播，通信量约为数据并行的 1.5 倍。ZeRO 论文的理论分析给出了这个数字。
Stage 2：梯度分片，顺手把 14GB 也省了
反向传播算出的梯度，真的需要每张卡都保存完整一份吗？不需要。Stage 2 把原本对梯度做的 all-reduce 换成了 reduce-scatter。
- all-reduce
- 每张卡都拿到全量平均梯度，代价是每张卡都要存全量。
- reduce-scatter
- 每张卡只拿到自己负责那部分参数的平均梯度，代价是额外一次通信。
工作流变成：
- 前向、反向照旧。
- 梯度做 reduce-scatter，每卡只保留局部梯度。
- 用局部梯度更新自己分片的优化器和参数。
- all-gather 广播新参数。
于是梯度也降到 14/N。N=8 时，每卡显存约 26.25GB。通信量依然大约 1.5 倍数据并行，因为 reduce-scatter 和 all-gather 加在一起，和原来的一次 all-reduce 差不多。
Stage 3：参数分片，让模型大小不再取决于单卡显存
既然优化器状态和梯度都能分片，参数为什么不行？Stage 3 把参数也切开了。但参数是前向/反向的必需输入，切开之后怎么算？
ZeRO 的答案是一句有点暴力的话：用的时候再凑回来，算完立刻丢掉。前向传播到某一层时，先通过 all-gather 把这层的所有权重收集到当前卡上；计算完这一层，只保留属于自己分片的参数，其他临时权重全部释放。反向传播时，再重新收集一遍权重，算出梯度后做 reduce-scatter。
这样，每张卡长期只需要保存 1/N 的参数、1/N 的梯度和 1/N 的优化器状态。N=8 时，每卡显存约 14GB。卡越多，每卡消耗越小，模型规模理论上可以随卡数线性扩展。
代价是通信量显著增加，每一层都要 gather 和 scatter，网络开销是 Stage 1/2 的好几倍。这也是为什么 Stage 3 在高速网络（如 InfiniBand）上表现很好，换成千兆以太网就会慢到怀疑人生。
一张表说透三个 Stage
| Stage | 分片对象 | 每卡仍需全量 | 每卡显存（N=8, 7B） | 通信开销 |
|---|---|---|---|---|
| Stage 1 | 优化器状态 | 参数+梯度 | 约 38.5GB | 约 1.5× 数据并行 |
| Stage 2 | 优化器状态+梯度 | 参数 | 约 26.25GB | 约 1.5× 数据并行 |
| Stage 3 | 优化器状态+梯度+参数 | 无 | 约 14GB | 明显更高，逐层通信 |
注意，这里的数字是纯训练状态，实际训练还要算上激活值、临时缓冲区，所以真实开销会更高，但三个 Stage 的相对关系不变。
我曾经的误解：Stage 3 不是模型并行
我最初看到“参数分片”四个字，脑子立刻想到模型并行——把不同的层放到不同 GPU 上，大家一起拼出一个大模型。后来读了 DeepSpeed 的实现细节，才发现完全不是一回事。
流水线模型并行把 Transformer 的层切开，设备 1 算第 1-8 层，设备 2 算第 9-16 层，数据像流水线一样流过。而 ZeRO Stage 3 里，每张卡都会完整地算每一个层，区别只是这一层的权重不一定在本地，需要临时从别的卡取。一个是按层分工，一个是按参数分工。
想通这一点后，我突然理解 ZeRO 真正的价值：它没有改变数据并行的计算方式，只是用通信换掉了数据并行里每卡重复存储的模型状态。所以它不需要处理模型并行的负载不均，只需要把通信管好。
实战中选择哪一档？
- 模型参数能放下一张卡，但优化器状态放不下 → Stage 1 足够了。
- 希望再把梯度省掉，好开更大的 batch → Stage 2 是甜点位。
- 单卡连完整参数都放不下 → 不得不 Stage 3。
- 如果你的集群只有普通以太网，优先考虑 Stage 1/2，别为了省显存把自己卡到通信瓶颈上。
很多人以为 Stage 越高越高级，其实不是。Stage 3 的显存节省确实逐级递增，但通信开销也在递增。你的训练瓶颈如果在计算，那 Stage 3 的通信额外开销会抵消掉显存优势；如果瓶颈在显存，那 Stage 3 才真正解决你的问题。
原创文章，作者：guanweilu，如若转载，请注明出处：https://guanweilu.cn/article/1460.html