训练显存账本：参数、梯度、优化器状态各占多少

摘要

训练一个模型，显存里装的远不止模型本身。以最常见的混合精度 + Adam 配置为例，每个参数需要约 16 字节显存，其中：

优化器状态 ≈ 75%（Adam 的两个动量 + FP32 主权重副本）
模型参数 ≈ 12.5%（FP16/BF16 工作副本）
梯度 ≈ 12.5%（FP16）

本文先逐项把这 16 字节的账算清楚（含”为什么还要存一份 FP32 主权重”的推导），再给出三种精度配置下的占比对照，讲 ZeRO 分片顺序为什么恰好按这个占比从大到小来，然后单开一节（§4）用 Megatron 公式逐项估算激活显存——这笔账与参数量无关、随序列长度平方增长，最后是 8-bit Adam、RL 训练等边界情况。

0. 记号约定

记号	含义
$\Psi$	模型参数量（个数）。7B 模型即 $\Psi = 7\times10^9$
FP32	32 位浮点，1 个参数占 4 字节
FP16 / BF16	16 位浮点，1 个参数占 2 字节
$m$ 、 $v$	Adam 的一阶矩、二阶矩
$s$	序列长度（token 数）
$b$	micro-batch 大小（一次前向处理的样本数）
$h$	隐藏维度（hidden size）
$a$	注意力头数
$L$	Transformer 层数

后文所有字节数都是”每个参数平均”的口径： $2\Psi$ 读作”每个参数 2 字节，总参数 $\Psi$ 个，共 $2\Psi$ 字节”。

1. 逐项算账：16Ψ 怎么来的

要点：混合精度 + Adam 下，每个参数身上有 4 项显存开销：FP16 参数 2 字节、FP16 梯度 2 字节、Adam 动量 8 字节、FP32 主权重 4 字节，合计 16 字节。

1.1 模型参数：2Ψ

混合精度训练里，前向和反向计算走 FP16/BF16（省显存、快），所以参数的工作副本按 2 字节存。这一项最直观：7B 模型参数本身就是 $2 \times 7\times10^9 = 14$ GB。

1.2 梯度：2Ψ

反向传播算出的梯度同样用 FP16/BF16 存储。梯度和参数一一对应，所以也是 $2\Psi$ 。

1.3 Adam 状态：8Ψ

Adam 为每个参数维护两个量：

$m$ （一阶矩）：梯度的指数滑动平均，相当于”最近梯度的方向惯性”；
$v$ （二阶矩）：梯度平方的滑动平均，用来归一化步长——哪个参数的梯度一直很大，它的更新步长就被压小。

两个量都按 FP32 存，各 $4\Psi$ ，合计 $8\Psi$ 。这是显存里最大的一笔单项开销。

1.4 FP32 主权重：4Ψ

参数已经有 FP16 副本了，为什么还要再存一份 FP32 的？因为FP16 的参数根本”加不动”小更新。看一个具体场景：

假设某个权重 $w = 1.0$ ，一步更新量 $\Delta w = 1\times10^{-4}$ （比如梯度 0.1 × 学习率 $10^{-3}$ ，量级合理）。FP16 的尾数只有 10 位，在 1.0 附近相邻两个可表示数间隔约 $2^{-10} \approx 9.8\times10^{-4}$ 。 $\Delta w$ 比这个间隔小一个量级， $1.0 + 10^{-4}$ 四舍五入后还是 1.0——更新被”吞掉”了。而 FP32 在 1.0 附近间隔约 $2^{-23}\approx 1.2\times10^{-7}$ ， $10^{-4}$ 可以精确加上。

所以标准做法是：更新在 FP32 主副本上做，定期拷回 FP16 工作副本。这份 FP32 主权重占 $4\Psi$ 。BF16 同理——它的尾数只有 7 位，比 FP16 更粗，只是指数范围与 FP32 一致，更需要 FP32 主副本兜底。

1.5 汇总

把四项加起来，就是著名的”每参数 16 字节”总账：

M = \underbrace{2\Psi}_{\text{参数}} + \underbrace{2\Psi}_{\text{梯度}} + \underbrace{8\Psi}_{m,v} + \underbrace{4\Psi}_{\text{主权重}} = 16\Psi \tag{1}

(1)

一个参数背的 16 字节（混合精度 + Adam）
┌──────────────┬──────────────┬───────────────────────────────┐
│   参数 FP16   │   梯度 FP16   │        优化器状态 FP32          │
│     2 Ψ      │     2 Ψ      │    m: 4 Ψ   v: 4 Ψ   主权重: 4 Ψ │
│    12.5%     │    12.5%     │              75%                │
└──────────────┴──────────────┴───────────────────────────────┘

占比结论：

优化器状态（ $m$ 、 $v$ 加主权重） $= 12\Psi = 75\%$
模型参数（FP16） $= 2\Psi = 12.5\%$
梯度 $= 2\Psi = 12.5\%$

可以把它记成：每个参数出门背 16 字节行李，其中 12 字节是 Adam 家的（两个动量各 4、主权重 4），参数自己和梯度各只占 2 字节。显存的大头从来不是模型本身，而是优化器。

2. 三个口径的占比对照

要点：占比数字取决于两件事——主权重算给谁、用什么优化器。下面三种配置互相对照。

配置	参数	梯度	优化器状态	合计
混合精度 + Adam	2Ψ (12.5%)	2Ψ (12.5%)	12Ψ (75%)	16Ψ
纯 FP32 + Adam	4Ψ (25%)	4Ψ (25%)	8Ψ (50%)	16Ψ
纯 FP32 + SGD	4Ψ (50%)	4Ψ (50%)	0	8Ψ

两种 FP32 配置总量恰好也都是 16Ψ 或 8Ψ，但构成完全不同。三个细节：

主权重算给谁：若把 FP32 主权重算进”参数”（很多人这么归类），混合精度 + Adam 的占比变成参数 $6\Psi$ =37.5%、梯度 12.5%、优化器动量 $8\Psi$ =50%。本文 §1 的口径把主权重算进优化器状态，因为它是 Adam 式更新流程的一部分——两种归类都对，关键是别重复计数。
纯 FP32 + SGD：SGD 没有动量状态，最省，但训练效果在大模型上通常不如 Adam 系。
共性：只要用了 Adam，优化器那一栏就不会低于一半。优化器是压缩显存的第一靶子，这直接引出 §3 的 ZeRO。

3. 两个著名数字与 ZeRO 的分片顺序

要点：GPT-2 1.5B 需要 24GB、7B 需要 112GB——这些数字都来自 16Ψ；ZeRO 三级分片恰好按占比从大到小拆。

验算两个经典数字：GPT-2 1.5B 训练要 $16 \times 1.5\times10^9 = 24$ GB，当年单卡 16GB 装不下，这是分布式训练研究的重要起点；7B 模型要 $16 \times 7\times10^9 = 112$ GB，一张 80GB 的卡也放不下。

ZeRO 的三级分片顺序，恰好就是按 §1 的占比从大到小来：

ZeRO-1：分片优化器状态（最大头，75%）；
ZeRO-2：再分片梯度（12.5%）；
ZeRO-3：连参数也分片（12.5%）。

为什么这个顺序？每级分片都按”谁占得多先拆谁”的贪心——先打掉 75% 的大头，边际收益最大。分片后的原理可以把它记成：每张卡只存自己那一份（1/N）的优化器状态或梯度，计算需要时再通信取回来，用带宽换显存。

4. 激活显存：第二本账（与参数量无关）

要点：16Ψ 只算了”模型状态”——参数、梯度、优化器。前向传播的中间结果还要占一份”激活显存”，它和参数量无关，随 batch × 序列长度线性增长、其中注意力部分随序列长度平方增长。长序列训练时，这笔账常常反超 16Ψ。

4.1 为什么反向要用前向的中间结果

反向传播是链式法则从最后一层往回推。看两个具体操作：

线性层 $y = xW$ ：算 $\partial L/\partial x$ 需要 $W$ ，算 $\partial L/\partial W$ 需要 $x$ ——前向的输入 $x$ 必须在反向时还在；
注意力里的 softmax： $\partial P/\partial S$ 是 $P$ 自己参与计算的，所以反向要 $P$ ——而 $P$ 又是 $S$ 算出来的。

所以前向每算一步，都要把中间结果留一份，等反向来取。这些”留饭”的中间结果就是激活（activations）。层越深、序列越长、batch 越大，饭留得越多。

4.2 逐项清单：线性项 + 平方项

Megatron 论文（Korthikanti et al.，arXiv:2205.05198）把单层激活分成两类：

项	形状	元素数	字节数
线性项：17 个逐元素张量（两个 LayerNorm 的输入、各线性层的 GEMM 输入、GELU 激活前值、dropout 残差掩码等）× 2 字节	$b\times s\times h$	$17\,b\,s\,h$	$34\,b\,s\,h$
平方项：注意力分数 $QK^\top$ （2 字节）	$b\times a\times s\times s$	$b\,a\,s^2$	$2\,b\,a\,s^2$
平方项：注意力概率 $P$ （2 字节）	$b\times a\times s\times s$	$b\,a\,s^2$	$2\,b\,a\,s^2$
平方项：dropout 掩码（布尔，1 字节）	$b\times a\times s\times s$	$b\,a\,s^2$	$1\,b\,a\,s^2$

合并起来就是著名的 Megatron 公式（单位：字节，单层）：

M_{\text{act/层}} = s\,b\,h\left(34 + 5\,\frac{a\,s}{h}\right) \tag{2}

(2)

两个系数记法：34 = 17 个逐元素张量 × 2 字节；5 = $QK^\top$ 的 2 + $P$ 的 2 + dropout 掩码的 1，单位都是”每个 $b\,a\,s^2$ 元素”。整份模型（ $L$ 层）就是 $L \times M_{\text{act/层}}$ ，再乘以同时在飞的 micro-batch 数（流水线并行 + 梯度累积时，多个 micro-batch 的前向会同时占着激活）。

一眼能看出的规律：序列长度 $s$ 翻倍，线性项翻倍，平方项翻四倍。注意力那个 $s\times s$ 才是激活显存真正的火药桶。

4.3 代入 7B 模型算一遍

取 7B 模型的典型超参： $h = 4096$ ， $a = 32$ ， $L = 32$ 层， $b = 1$ 。按公式 (2)：

序列长度 $s$	系数 $34 + 5as/h$	单层激活	32 层合计	对比模型状态 112GB
2048	$34 + 80 = 114$	0.96 GB	≈ 31 GB	约 1/4
4096	$34 + 160 = 194$	3.25 GB	≈ 104 GB	几乎持平
8192	$34 + 320 = 354$	11.9 GB	≈ 380 GB	3.4 倍

7B 模型：模型状态 112GB vs 激活显存
s=2048  模型状态 ████████ 112GB
        激活     ██        31GB    ← 还能接受
s=4096  模型状态 ████████ 112GB
        激活     ███████▌ 104GB   ← 反超在即
s=8192  模型状态 ████████ 112GB
        激活     ██████████████████████████████ 380GB ← 主角换人

这张表的结论很反直觉：在 s = 8192 的长序列训练里，占显存大头的已经不是那 16Ψ，而是注意力激活。ZeRO 论文原文也印证了这一点：1.5B 的 GPT-2 在序列 1024、batch 32 下训练，总共需要约 60GB 显存——比 24GB 的模型状态多出一倍多，多出来的主要就是激活。

4.4 怎么压激活

Activation checkpointing（重计算）：前向不存中间激活，反向要用时把那一层的前向重跑一遍现算。用算力换显存。最省的做法只留每层输入（平方项直接清零），但每层前向跑两遍；Megatron 的选择性重计算（selective recomputation）只重算注意力等大头，实测激活省 70%、FLOPs 只多 2.7%。
FlashAttention： $S$ 、 $P$ 根本不落地 HBM（见 FlashAttention 那篇）——公式 (2) 里的平方项直接清零，相当于免费拿到了 checkpointing 在注意力上的收益。
序列并行：把 LayerNorm、dropout 这些”逐 token”激活按序列维度切到各卡，线性项除以并行度；配合张量并行正好覆盖公式 (2) 的两项（见并行训练那篇）。

5. 边界与例外

要点：16Ψ 与激活之外，优化器本身也可以压缩；强化学习训练还有额外负担。

优化器状态可以压缩。 两个方向：

8-bit Adam：把 $m$ 、 $v$ 量化到 8 位，动量从 $8\Psi$ 压到 $2\Psi$ ，加上主权重 $4\Psi$ ，优化器状态共 $6\Psi$ ，总账从 $16\Psi$ 降到 $10\Psi$ ，优化器占比降到 60%。
Adafactor：二阶矩用行列分解存，只存两个向量而不是完整矩阵，状态量降到 $O(\sqrt{\Psi})$ 量级，几乎可以忽略——它最早就是为训练 T5 设计的。

强化学习训练有额外负担。 RL 阶段（比如 GRPO 这类组内相对优化的算法）除了训练模型本身，还要常驻一份参考模型（算重要性比率或 KL 的基准网络，至少 $2\Psi$ ~ $4\Psi$ ），rollout 生成的激活也要另算。所以同样的模型，RL 阶段的显存压力通常比 SFT 更大。

6. 结论

两本账分开记：

模型状态（§1~§3）：混合精度 + Adam 下每个参数背 16 字节，其中 75% 是优化器状态，参数和梯度各 12.5%。显存优化的第一优先级永远是优化器——ZeRO-1 只分片优化器状态就能省掉一大半；退而求其次才轮到梯度和参数。
激活（§4）： $s\,b\,h\,(34 + 5as/h)$ 字节/层，和参数量无关；短序列下只是零头，长序列下（7B 模型 s=8192 时 ≈ 380GB）反超模型状态三倍多。压法按性价比排序：FlashAttention 清平方项（免费）→ 选择性重计算（省 70%、多算 2.7%）→ 序列并行（线性项分摊）。

参考文献

Rajbhandari et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. arXiv:1910.02054.
Korthikanti et al. Reducing Activation Recomputation in Large Transformer Models. arXiv:2205.05198.
Micikevicius et al. Mixed Precision Training. arXiv:1710.03740.
Dettmers et al. 8-bit Optimizers via Block-wise Quantization. arXiv:2110.02861.
Shazeer & Stern. Adafactor: Adaptive Learning Rates with Sublinear Memory Cost. arXiv:1804.04235.

摘要