4. 模型显存占用分析
做GPU集群这么久，我见过最多的翻车现场就是——显存算错了。
你以为模型能塞进去，结果一跑就OOM。你以为显存够用，结果batch size调小到1还是炸。
说白了，显存占用不是拍脑袋估的。你得一笔一笔算清楚。
今天我就带你把这笔账算明白。参数、梯度、优化器状态、激活值，一个都不放过。
核心公式（记住这个）：
总显存 ≈ 参数显存 + 梯度显存 + 优化器状态 + 激活值显存 + 临时缓冲区
4.1 参数显存：最基础的那部分
参数显存最好算。说白了就是模型权重占的地方。
每个参数如果是FP32，占4字节。如果是FP16，占2字节。
计算公式：
参数显存 = 参数量 × 每个参数字节数举个例子：
- LLaMA-7B，FP32训练：7 × 10⁹ × 4 = 28 GB
- LLaMA-7B，FP16推理：7 × 10⁹ × 2 = 14 GB
我的经验： 我个人习惯把参数量先换算成GB心里有个数。7B模型FP32就是28GB，这个数字我闭着眼都能背出来。
4.2 梯度显存：反向传播的代价
梯度显存很多人会漏算。其实它跟参数显存一样大。
为什么？因为每个参数都有一个对应的梯度。
计算公式：
梯度显存 = 参数量 × 每个梯度字节数注意：梯度的精度通常跟参数一致。你用FP16训练，梯度也是FP16。
所以：
- LLaMA-7B，FP32训练：梯度显存 = 28 GB
- LLaMA-7B，FP16训练：梯度显存 = 14 GB
避坑指南： 我曾经遇到过一个问题——梯度累积的时候，显存会翻倍。因为你要存多份梯度。如果你梯度累积步数是4，那梯度显存就是原来的4倍。这个坑我踩过，疼。
4.3 优化器状态：Adam的"三倍暴击"
优化器状态这块，是新手最容易低估的。
你想想看，Adam优化器除了存参数和梯度，还要存什么？
- 一阶动量（momentum）
- 二阶动量（variance）
每个都是跟参数一样大小的张量。
所以Adam的显存占用是：
优化器状态显存 = 参数量 × 每个状态字节数 × 2（一阶+二阶）算一笔账：
- LLaMA-7B，FP32训练，Adam优化器：7 × 10⁹ × 4 × 2 = 56 GB
嗯，你没看错。优化器状态比参数本身还大。
| 优化器 | 额外状态数 | 显存倍数（相对参数） |
|---|---|---|
| SGD | 0 | 0x |
| SGD + Momentum | 1 | 1x |
| Adam | 2 | 2x |
| AdamW | 2 | 2x |
我的建议： 如果你显存吃紧，可以考虑用bitsandbytes的8-bit Adam。优化器状态直接从FP32压缩到INT8，显存直接砍掉75%。代价是训练速度稍微慢一点，但显存省太多了。
4.4 激活值显存：最"动态"的部分
激活值显存是最难算的。因为它跟batch size、序列长度、模型结构都有关。
说白了，每一层前向传播时产生的中间结果，都要存下来给反向传播用。
粗略估算公式：
激活值显存 ≈ batch_size × 序列长度 × 隐藏层维度 × 层数 × 字节数 × 系数这个系数一般在2~4之间，取决于你用了多少激活函数和归一化层。
举个例子：
- LLaMA-7B，batch_size=1，seq_len=2048，hidden_dim=4096，32层，FP16
- 激活值显存 ≈ 1 × 2048 × 4096 × 32 × 2 × 3 ≈ 1.5 GB
看着不大？你把batch_size调到64试试：
- 1.5 GB × 64 = 96 GB
直接爆炸。
注意： 激活值显存跟batch size是线性关系。你batch size翻倍，激活值显存也翻倍。但参数显存、梯度显存、优化器状态跟batch size无关。所以大batch size训练时，激活值往往是瓶颈。
4.5 一张图看懂所有
下面这张SVG图，把整个显存占用结构画清楚了。我建议你保存下来，每次分配显存前看一眼。
4.6 实战：算一笔完整的账
咱们拿LLaMA-7B，FP16训练，batch_size=4，seq_len=2048来算：
| 项目 | 计算过程 | 显存占用 |
|---|---|---|
| 参数显存 | 7B × 2字节 | 14 GB |
| 梯度显存 | 7B × 2字节 | 14 GB |
| 优化器状态（Adam） | 7B × 2字节 × 2 | 28 GB |
| 激活值显存 | 4 × 2048 × 4096 × 32 × 2 × 3 | ~6 GB |
| 临时缓冲区（估算15%） | (14+14+28+6) × 0.15 | ~9 GB |
| 总计 | ~71 GB |
嗯，71GB。一张A100是80GB，勉强塞得下。但如果你batch size调到8，激活值翻倍到12GB，总显存就奔着80GB去了，非常危险。
我的经验： 我一般留20%的显存余量。比如80GB的卡，我只敢用到64GB左右。不然训练过程中突然爆显存，几个小时的训练就白跑了。
4.7 几个省显存的技巧
最后分享几个我常用的省显存手段：
- 梯度检查点（Gradient Checkpointing）： 不存所有激活值，反向传播时重新算。省显存但慢。我一般用这个把激活值显存砍掉50%~70%。
- 混合精度训练（AMP）： 参数和梯度用FP16，优化器状态用FP32。显存省一半。
- ZeRO优化器： 把优化器状态、梯度、参数分散到多张卡上。显存占用跟卡数成反比。
- 8-bit Adam： 优化器状态从FP32压缩到INT8，省75%。
避坑指南： 我曾经试过把所有优化器状态都压到INT8，结果模型收敛变慢了。后来发现是某些层的梯度范围太大，8-bit精度不够。所以建议关键层（比如embedding层）保留FP32。
好了，显存计算这块就讲到这里。你只要记住那个公式，再结合我给的例子，基本不会算错。下次分配显存前，先拿笔算一算，别等到OOM了才后悔。