3、ZeRO核心思想：内存冗余问题，分片策略，通信与计算权衡
在数据并行训练中，最容易被忽视的成本不是算力，而是显存冗余。假设使用 8 张 GPU 训练一个 7B 参数模型，每张卡都保存一份完整的模型参数、梯度和优化器状态，那么同一份数据被复制了 8 次。ZeRO 的核心思想正是：把原本每张卡都重复保存的状态切分开，只在需要时通过通信临时聚合。它用通信换显存，用显存换可训练规模。
一、内存冗余从哪里来
以混合精度 Adam 训练为例，每个可训练参数在单卡上通常对应以下状态：
- FP16 参数：2 字节
- FP16 梯度：2 字节
- FP32 主参数副本：4 字节
- Adam 一阶动量：4 字节
- Adam 二阶动量：4 字节
合计约 16 字节/参数。一个 7B 模型仅模型状态就约 112 GB，远超单张 80 GB 显卡容量。若再叠加激活值、临时缓冲和通信桶，显存压力会进一步放大。传统 DDP 的问题在于：每张卡都持有完整副本，显存随模型线性增长，而计算却已经被切分。
二、分片策略：ZeRO-1/2/3 的递进切分
ZeRO 将模型状态分为三部分，并按阶段逐步切分：
| 阶段 | 切分对象 | 单卡显存节省 | 关键通信 |
|---|---|---|---|
| ZeRO-1 | 优化器状态 | 约 4 倍 | 梯度 AllReduce + 参数 AllGather |
| ZeRO-2 | 优化器状态 + 梯度 | 约 8 倍 | 梯度 ReduceScatter + 参数 AllGather |
| ZeRO-3 | 优化器状态 + 梯度 + 参数 | 约 N 倍（N 为卡数） | 前向/反向按需 AllGather 参数 |
ZeRO-1 只切优化器状态，参数和梯度仍保持完整副本，因此实现简单、通信增量小。ZeRO-2 进一步切梯度，反向传播时用 ReduceScatter 替代 AllReduce，每张卡只保留自己负责的那一片梯度。ZeRO-3 最彻底，连参数也切分，前向和反向计算到某一层时，才通过 AllGather 把该层参数临时拼齐，用完即释放。
三、通信与计算的权衡
分片不是免费的。ZeRO-1 在反向传播后需要对梯度做 AllReduce，再在优化器更新后对参数做 AllGather；ZeRO-2 把 AllReduce 拆成 ReduceScatter，通信量约为原来的 1.5 倍；ZeRO-3 在前向和反向中反复 AllGather 参数，通信量进一步上升，但通过通信桶和计算重叠可以隐藏大部分开销。
关键权衡点如下：
- 显存 vs 通信：切得越细，单卡显存越低，但跨卡通信越多。
- 通信量 vs 重叠能力：ZeRO-3 通信量大，但可与计算流水重叠，实际吞吐未必线性下降。
- 分片粒度 vs 实现复杂度：分片越细，参数收集、释放和生命周期管理越复杂。
- 集群带宽：NVLink、InfiniBand 带宽越高，ZeRO-3 的收益越明显；低带宽环境可能更适合 ZeRO-1/2。
stage3_prefetch_bucket_size、stage3_param_persistence_threshold 等参数调优。
四、案例：7B 模型在 8 卡上的显存变化
假设 8 卡训练 7B 模型，单卡模型状态约 112 GB。开启不同阶段后，单卡模型状态近似变化如下：
ZeRO-1：优化器状态分片
参数 14 GB + 梯度 14 GB + 优化器状态 84/8 ≈ 10.5 GB
合计约 38.5 GB
ZeRO-2：优化器状态 + 梯度分片
参数 14 GB + 梯度 14/8 ≈ 1.75 GB + 优化器状态 10.5 GB
合计约 26.25 GB
ZeRO-3：参数 + 梯度 + 优化器状态全分片
(14 + 14 + 84) / 8 ≈ 14 GB
合计约 14 GB（未计激活值和临时缓冲）可以看到，ZeRO-3 把模型状态压到了约 1/8，但代价是前向和反向中频繁的参数 AllGather。若集群带宽充足，这部分通信可以被计算掩盖；若带宽不足，训练速度可能明显下降。
五、常见误区与复盘
误区二：忽略激活值显存。ZeRO 主要优化模型状态，激活值仍需靠梯度检查点、序列并行等手段控制。
误区三：只调 stage 不调桶大小。通信桶过小会导致大量小消息，降低带宽利用率；过大则增加显存峰值。
复盘时建议关注三个指标：单卡显存峰值、每步通信耗时占比、吞吐量（tokens/s）。如果显存下降但吞吐骤降，说明通信未重叠好；如果显存没降多少，说明分片未生效或仍有完整副本残留。ZeRO 的本质不是“免费省显存”，而是把显存压力转化为可调度的通信与计算重叠问题。