深度学习显存机制：张量存储、梯度、优化器状态

显存到底被谁吃掉了？很多人以为只有模型参数占显存，其实远不止。我刚开始做训练优化时，也这么认为。结果一跑大模型，OOM 了，才明白水有多深。今天咱们就把这件事彻底拆清楚。

显存里的东西，按生命周期可以分成三类：持久驻留的模型参数和优化器状态，每次迭代创建并释放的前向激活缓存，还有反向传播时临时出现的梯度。另外还有输入数据、临时工作区等等。下面我一个个说。

1. 张量存储：模型参数与前向激活

模型参数是最直观的。假设你有一个 7B 参数的模型，每个参数用 FP32 存——那就是 7B × 4 bytes = 28 GB。换成 FP16 或者 BF16，直接减半到 14 GB。所以混精度训练能把模型塞进更小的显存里。

但真正的大头往往不是参数，而是前向激活缓存。反向传播需要用到每一层的中间输出，这些输出要一直保留到该层梯度算完才能释放。我见过一个 12 层 Transformer，激活缓存的大小是参数量的 2～3 倍。如果你用超长序列，激活甚至能吃掉 80% 的显存。

举个例子：一个 LLaMA-7B 在 batch size 1、序列长度 2048 时，激活占用约 16 GB，参数 + 梯度 + 优化器状态才 28 GB。所以有时候你以为调大 batch size 没问题，其实激活先爆了。

import torch
print(f"当前分配: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
print(f"最大分配: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB")
print(f"缓存总量: {torch.cuda.memory_reserved() / 1e9:.2f} GB")

2. 梯度：动态分配，及时清理

梯度是在反向传播过程中产生的，形状和对应参数一样。所以如果参数是 28 GB，梯度也是 28 GB（FP32）。不过梯度不用一直保留——每个参数梯度算完后，就可以被覆盖或释放。正常训练中，梯度和激活释放的时机是交错的，峰值显存通常出现在反向传播的中后期。

我曾经犯过一个错：为了方便，把所有梯度都 retain 下来做梯度剪裁，结果显存直接翻倍。后来改用 in-place 的 clip_grad_norm_，这问题就没了。记住：不需要的梯度别留着。

还有一个点：梯度是否显式占用显存？是的。在 PyTorch 里，只要设置了 requires_grad=True，张量就会有一块独立的梯度存储空间（默认是 None，只有调了 backward 才会分配）。如果你用 torch.no_grad() 包住某些计算，就不会为那些输出分配梯度空间。这个技巧在推理或特征提取时很实用。

3. 优化器状态：被忽视的显存黑洞

优化器需要保存额外的状态才能更新参数。不同的优化器，额外开销不一样。我整理了一个表格，方便你看清楚：

看到没？用 AdamW 训练一个 7B 模型，参数 28 GB，优化器状态 56 GB，加上梯度 28 GB，总共已经 112 GB 了！这还没算激活。所以我自己做预训练时，要么用 Adafactor 节约显存，要么用 BF16 优化器状态（比如 bitsandbytes 的 8-bit Adam）。

4. 一张图看清显存构成

下面这张 SVG 把整个显存占用拆解成几个主要部分，箭头表示数据流和依赖关系：

5. 实战：如何计算你的工作集大小

知道了这些组件，我们就能估算一个训练 step 的显存峰值。假设参数用 FP16，激活用 FP16，梯度用 FP16，优化器动量用 FP16：

1. 模型参数量：N（每个参数 2 bytes）→ 2N
2. 激活缓存：根据模型架构和序列长度，通常是参数的 2～4 倍（经验值）
3. 梯度：和参数大小一样 → 2N
4. 优化器状态：Adam 额外 2×2N = 4N（一阶 + 二阶，每个 2 bytes）
5. 输入数据 + 临时缓冲区：约 1～2 GB

所以总显存 ≈ 2N (参数) + (2～4)×2N (激活) + 2N (梯度) + 4N (优化器) + 1～2 GB。如果 N=7B，那就是 2×7 + 3×14 + 2×7 + 4×7 + 2 = 14 + 42 + 14 + 28 + 2 = 100 GB。一块 A100 80G 装不下，得用两张或换优化器。

最后强调一个很容易忽略的点：框架的缓存分配器（比如 PyTorch 的 caching allocator）会多预留一部分显存，避免频繁 cudaMalloc。你看到的 memory_allocated 是实际使用的，而 memory_reserved 是总占用的。调 batch size 时，盯着 max_memory_reserved 会更准。