张量显存占用分析：模型参数、梯度、优化器状态、中间激活的内存计算方法
做深度学习训练，你最先撞上的瓶颈是什么？十有八九是显存不够。
我遇到过不少同学，模型明明能装下，一跑训练就OOM。问题出在哪？
说白了，训练时占显存的不止模型参数。还有梯度、优化器状态、中间激活——这四大家族才是显存大户。
咱们今天就把每个家族的身家算个清楚。以后你再碰到OOM，直接公式一套，心里就有底了。
模型参数：最直观的占位
每个参数占多少内存？就两个因子：参数量 × 数据类型字节数。
参数量 = 所有可学习张量的元素个数之和
字节数：float32 → 4 字节，float16/bfloat16 → 2 字节，int8 → 1 字节
模型参数显存 = 参数量 × 字节数举个例子：LLaMA-7B，参数量约 7×10⁹。
如果用fp16，参数显存 = 7e9 × 2 = 14 GB。如果用fp32，就是28 GB。
你看，一张A100也就80GB，纯参数就快吃一半了。
model.parameters()torch.numel() 算总参数量，再手动乘字节数。模型加载前就能估算。
梯度：训练时隐形的兄弟
梯度跟参数等大。前向算完，反向传播必须存一份梯度，才能更新参数。
梯度显存 = 参数量 × 字节数（通常与参数精度一致）所以训练时，参数 + 梯度就已经是 2倍的参数显存了。
如果你用fp16参数 + fp16梯度，那就是2×14=28 GB。
这里有个坑：混合精度训练时，梯度可能是fp16，但Loss Scaling之后会转成fp32再更新。你想想看，这块内存有临时分配，也可能造成峰值上升。我之前就踩过——以为自己算得很准，结果跑起来还是炸了，一查发现是梯度累加时没清临时buffer。
优化器状态：Adam这个吃内存大户
SGD的话，优化器只记录学习率，几乎不占额外内存。
但Adam就厉害了。它需要维护 一阶矩（动量） 和 二阶矩（自适应学习率）。每个参数对应两个额外变量。
纯fp32训练时：
参数: 4B/参数
梯度: 4B/参数
一阶矩: 4B/参数
二阶矩: 4B/参数
总计 = 16B/参数 → 即 4倍参数量
混合精度（常见做法）：
参数(fp16): 2B/参数
梯度(fp16): 2B/参数
fp32副本: 4B/参数
一阶矩(fp32): 4B/参数
二阶矩(fp32): 4B/参数
总计 = 16B/参数 → 还是 4倍？不对，因为参数本身是fp16 + fp32副本，实际是6倍？
嗯，这里要注意。混合精度下，参数保存两份：一份fp16用于前向/反向，一份fp32用于累积更新。所以内存分到优化器状态里的“参数副本”算作优化器开销。通常我们简化计算：
训练时总显存（不含激活） ≈ 4×模型参数量 × 字节数（以fp32为基准）
| 优化器 | 参数占内存 | 梯度占内存 | 优化器状态占内存 | 合计（以每个参数计，fp32单位） |
|---|---|---|---|---|
| SGD（无动量） | 4B | 4B | 0 | 8B |
| SGD（动量 0.9） | 4B | 4B | 4B（动量） | 12B |
| Adam（fp32） | 4B | 4B | 8B（一阶+二阶） | 16B |
| AdamW（混合精度） | 2B（fp16参数） | 2B（fp16梯度） | 12B（fp32副本+一阶+二阶） | 16B（与fp32 Adam一致） |
中间激活：最大的隐藏炸弹
这部分最容易被忽视。前向传播时，每一层的输入、输出、中间值都要保留，供反向使用。
我见过一个项目，模型参数只占30GB，激活却吃掉了50GB，一训练就OOM。
以Transformer为例，每层激活内存 ≈ batch_size × seq_len × hidden_dim × 字节数 × 层数 × 系数
系数包括：
- 注意力：Q/K/V + 注意力矩阵 + dropout mask ≈ 3~4倍hidden
- FFN：两层线性，中间4倍隐藏维度
- LayerNorm：统计量
粗略估算：batch_size=4, seq_len=1024, hidden=4096, layers=32, fp16
单层激活 ≈ 4×1024×4096×2B ≈ 32 MB
总共32层 ≈ 1 GB（注意力+FFN系数2~3）≈ 2~3 GB
如果batch_size=32呢？直接变成16~24 GB。激活会随batch线性增长。
实战中，激活占比经常超过参数。比如LLaMA-7B（参数14GB fp16），把batch从4调到8，激活从2GB跳到4GB，总显存就从20GB涨到22GB，而峰值可能更高（因为内存碎片）。
所以：激活优化是显存优化的大头。
总显存公式与实战计算
训练总显存 ≈ 参数显存 + 梯度显存 + 优化器状态显存 + 中间激活显存 + 额外开销（~5%）
或者用更通用的公式（以fp32为基准）：
总显存 = 参数量 × 4 × 4（Adam） + 激活显存
= 16×参数量（GB, 当参数量以10^9计） + 激活（gb）
举个例子：7B模型，fp16 Adam训练，batch=4，seq=1024，hidden=4096，layers=32。
参数+梯度+优化器（混合精度）≈ 16B/参数 × 7e9 = 112 GB ？等等，这里16B是按fp32每个参数4B算的。
实际混合精度下参数+梯度+优化器 = 14GB（fp16参数） + 14GB（fp16梯度） + 28GB（fp32副本+动量+方差） = 56 GB。
激活 ≈ 3 GB（上面估算）。
额外开销 ≈ 3 GB。
总计 ≈ 62 GB。一张A100 80G勉强能跑。
你看，简单一算就知道为什么大模型必须用ZeRO分片了。
可视化：显存占用的四大块
下面这张图帮你快速理解它们之间的关系。
实战代码：快速查看各块显存
PyTorch 里有个很好用的 API：torch.cuda.memory_summary()。但更细粒度的分析建议用 torch.cuda.memory.mem_get_info() 和 torch.cuda.memory_allocated()。
import torch
from transformers import AutoModel
model = AutoModel.from_pretrained("meta-llama/Llama-2-7b-hf", torch_dtype=torch.float16)
model.cuda()
# 只算模型参数
param_mem = sum(p.numel() * p.element_size() for p in model.parameters())
print(f"模型参数显存: {param_mem / 1024**3:.2f} GB")
# 模拟一次前向（生成激活）
batch_size, seq_len = 4, 1024
input_ids = torch.randint(0, 32000, (batch_size, seq_len)).cuda()
out = model(input_ids, output_hidden_states=False)
# 总分配显存
allocated = torch.cuda.memory_allocated()
print(f"当前总显存分配: {allocated / 1024**3:.2f} GB")
# 激活显存 ≈ 总分配 - 参数
act_mem = allocated - param_mem
print(f"中间激活（估计）: {act_mem / 1024**3:.2f} GB")torch.cuda.max_memory_allocated()，看峰值。如果峰值比总分配高很多，说明有临时buffer没及时释放。这时候可以用 torch.cuda.empty_cache() 手动释放，但治标不治本，还是得优化代码。
避坑与调优建议
最后总结几个我踩过的坑：
- 只算参数不算激活 – 相当于只带衣服没带行李，一跑就炸。
- 优化器状态不注意精度 – 混合精度下，fp32副本 + 动量 + 二阶矩 = 12B/参数，比参数本身还多。用bitsandbytes的8-bit Adam能省一半。
- 梯度累加时显存翻倍 – 如果你手动迭代累加，每个step都会保留梯度buffer。用PyTorch的accumulation_steps参数，或者用
no_grad包裹。 - 大batch激活炸掉 – 这时候 gradient checkpointing 是你的好朋友。代价是慢30%左右，但显存立减60%。