大模型训练ZeRO内存优化原理详解

0. 引言

Zero Redundancy Optimizer (ZeRO)，主要目标是减少内存使用并加速大规模模型的训练过程。它通过在多个 GPU 或者节点之间分散模型的状态（如梯度和参数）来实现这一目标。这种分散减少了每个计算节点上存储的冗余数据量，从而降低了内存占用。

论文：《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》

1. GPU 内存分布

1.1 模型状态

模型状态包括：
（1）优化器状态（Optimizer States），例如使用 Adam 优化器时的动量和梯度方差
（2）梯度（Gradients）
（3）参数（ Parameters）

上面的模型状态通常占据了大部分的内存，在混合精度训练中，还需要额外的内存来存储 fp32 的参数和优化器状态。

比如 GPT-2（具有 1.5B 参数）模型，模型状态的保存要求至少 24 GB 的内存。

1.2 剩余内存

除了模型状态外，剩余的内存包含：
（1）激活内存。用于正向传播以执行反向传播的存储，可以通过激活检查点（checkpointing）来减少，但会提升计算量；
（2）临时缓冲区。用于存储中间结果，其大小随着模型大小的增加而增加
（3）不可用的碎片化内存

以上统称为除保存模型状态之外的剩余内存。

2. ZeRO 优化

2.1 ZeRO-DP 优化

ZeRO-DP（ZeRO 数据并行），优化三个阶段的内存消耗情况：

$\Psi$ 为模型大小（参数个数），$K$ 为优化器状态的内存乘数，$N_d$ 为数据并行度，可以理解为 GPU 卡数。

在本例中，假设基于 Adam 优化器的混合精度训练，模型大小为 7.5B， $N_d=64$，$K=12$。

下面分别介绍 ZeRO-DP 优化的三个阶段的具体情况。

2.1.1 ZeRO-Stage1: 优化器状态划分

（1）$P_{os}$（Optimizer State Partitioning，优化器状态划分）
ZeRO 通过将优化器状态划分为 $N_d$ 个数据并行进程，每个进程仅存储、更新其对应分区的优化器状态，即整体优化器状态的$\frac{1}{N_d}$，从而减少了每个设备上所需的内存量。在每个训练步骤结束时，再收集每一个进程的结果，以获取整体更新后的状态参数。

（2）ZeRO-Stage1 内存优化后的结果，主要针对优化器状态（请参考上图）：
$$(2+2)\Psi +\frac{K\ast \Psi }{N_d}$$
可见，优化器状态内存在原始基础上有一个 $N_d$ 的除数。

（3）举例
在 7.5B 的模型上，标准的情况下要求 120GB 的内存，但是使用 $P_{os}$ 后，$N_d=$