大模型训练中 optimizer 的内存布局：fp32 参数 + fp32 状态 = 三倍模型大小的显存

你可能以为，训练大模型最吃显存的是模型本身。我一开始也这么想——一个 7B 参数的模型，光权重就有 14GB（fp16）。直到我第一次尝试在单机多卡上微调 7B 模型，结果还没跑完一个 step 就 OOM 了。盯着 torch.cuda.memory_summary() 的输出了半天，才发现罪魁祸首不是权重，而是 optimizer 状态。

具体来说，在典型的混合精度训练里，每个参数背后站着三个 fp32 张量：参数的 master copy、Adam 的一阶动量 m、二阶动量 v。每个张量对每个参数占用 4 字节，三个就是 12 字节。而你心里的"模型大小"如果是那个 fp16 版本，它只有 2 字节——光这三份就相当于 fp16 模型大小的 6 倍；如果以 fp32 模型大小为基准，则正好是 3 倍——这就是标题那个等式的由来。

这个发现当时颠覆了我的认知。我一直以为显存是"模型多大就占多少"，没想到真实账单里还有这么一笔巨款。所以今天我想跟你把这笔账算清楚：optimizer 为什么需要这么多状态？它到底怎么布局？以及为什么你听到的"3倍模型大小"既对又不全对。

Adam 为什么需要两个状态张量？

要理解内存，先看 Adam 的更新公式（Adam 论文）：

m_t = beta1 * m_{t-1} + (1-beta1) * g_t
v_t = beta2 * v_{t-1} + (1-beta2) * g_t^2
theta_t = theta_{t-1} - lr * m_t / (sqrt(v_t) + eps)

m 是梯度的一阶矩（带衰减的累加），v 是二阶矩。它们让每个参数拥有独立的自适应学习率：梯度平方大的方向迈小步，梯度平方小的方向迈大步。但代价是你必须从第一步开始，一路记住所有参数的 m 和 v，直到训练结束。这两个张量的形状和参数矩阵完全一样，而且以 fp32 存储。

你可能会问：梯度 g 不也是 N 个元素吗？确实是，但梯度是临时的：反向传播算出来、更新完就释放。m 和 v 不会释放，它们是跨 step 积累的"状态"。这就是 optimizer 状态内存的由来。

三倍的账是怎么算出来的

假设模型有 N 个参数，用 fp32 训练。模型本身是 4N 字节。optimizer 需要维护的东西，按持久性排列：一个 fp32 的参数主副本（如果是纯 fp32，它就是模型自己）、一个 fp32 的 m、一个 fp32 的 v。于是：

参数（fp32）:  4N
m（fp32）:     4N
v（fp32）:     4N
-------------------
合计:         12N

12N 除以 4N，就是 3。这就是"fp32 参数 + fp32 状态 = 三倍模型大小"的算术本质。

不同优化器的固定状态差异，在表格里一眼就能看出来：

优化器	额外持久状态数	总显存（以 fp32 参数 = 1 为单位）
SGD（无动量）	0	1
SGD + Momentum	1	2
Adam / AdamW	2	3

注意这里没有计激活值、梯度、临时缓冲区。实际训练时这些也要占，但 optimizer 状态是"常驻"的——从训练开始到结束一直占着，而且会随模型大小线性增长。

混合精度：为什么你看到的远不止 3 倍

训练 7B 以上的模型几乎不用纯 fp32。混合精度用 fp16/bf16 做 forward/backward，但更新参数时不能也丢到 fp16——梯度更新量常常小于 fp16 的最小精度，直接累加会"归零"。所以仍然需要一个 fp32 的 master weights。按照 NVIDIA 的混合精度训练文档，主权重用 fp32 保存，前向和反向用 half 精度副本。于是账单变成：

fp16/bf16 模型副本:  2N
fp32 master weights: 4N
Adam m:              4N
Adam v:              4N
------------------------------
固定占用合计:        14N

对比一下：fp16 模型大小是 2N，那 14N 是它的 7 倍；fp32 模型大小是 4N，那 14N 是它的 3.5 倍。为什么比"3倍"还高？因为多出来的一个 fp16 模型副本虽然小，但也是实打实的显存。所以"3倍"这个说法，只有在你以 fp32 模型大小为单位、并且只数那三个 fp32 张量的时候才精确。大家常说的"Adam 混合精度训练每个参数约 16 字节"，是这样凑出来的：fp16 权重 2N + fp16 梯度 2N + fp32 master 4N + Adam 状态 8N = 16N。

我在第一次 OOM 后，把 nvidia-smi 和 PyTorch 的 memory_summary() 对着看，才意识到：那个 7B 模型，fp16 权重只有 14GB，但加上 optimizer 状态轻松超过 70GB。8 张 A100（80GB）也放不下一个满参数微调，难怪数据并行要上 ZeRO。

ZeRO 是怎么把三倍拆回去的

核心思路很简单：数据并行时，每张卡都有一份完整的模型副本——这是冗余的。而 optimizer 状态不必每张卡都存完整。ZeRO-1 把 m、v（optimizer states）按数据并行的维度切块，每张卡只保留自己负责的那 1/N；更新时每张卡更新自己的分片，再通过 all-gather 把参数拼回完整。ZeRO 论文的目标，就是把这些本来被复制多份的模型状态，变成只存一份。