ZeRO 优化器

上一章我们看到 DDP 的内存问题：为了保证训练一致性（通过 All-Reduce 同步梯度），每个 GPU 都需要存储完整的模型状态。4 个 GPU 就是 4 份完整副本（参数、梯度、优化器状态）。ZeRO（Zero Redundancy Optimizer）的核心思想很直接：既然最终状态是一致的，那就每个 GPU 只存一部分，需要的时候再通信取回。

训练状态的冗余分析

先量化一下 DDP 的浪费。以混合精度 + Adam 为例，$N$ 个 GPU 训练一个 $\Phi$ 参数的模型，每个 GPU 需要存储：

• 参数（fp16）：$2\Phi$ bytes
• 梯度（fp16）：$2\Phi$ bytes
• 优化器状态（fp32）：$12\Phi$ bytes（参数副本 + 一阶矩 + 二阶矩）

合计 $16\Phi$ bytes，其中优化器状态占了 75%。

$N$ 个 GPU 就是 $N$ 倍冗余：全局存储 $16N\Phi$ bytes，但实际只需要 $16\Phi$ bytes。ZeRO 的三个 Stage 就是按从大到小的顺序，依次消除这些冗余。