模型显存分析：一块显卡到底能吃多少？

我们都知道训练大模型费显卡。但具体费在哪里？
是模型参数把显存撑爆了，还是中间激活把显存吃光了？
我见过不少团队，一上来就调大 batch size，结果 OOM 了，然后一脸懵地来问我。

其实，显存消耗是可以算出来的。搞清楚了这笔账，你就知道该从哪下手优化了。
我个人习惯先把参数、梯度、优化器状态、中间激活这四大块拆开看，哪个是瓶颈一目了然。

1. 模型参数与梯度：最直白的开销

假设你有一个 70 亿参数的模型，用 FP32 存储。每个参数 4 字节。
参数本身占多少？很简单：

参数量 × 每个参数的字节数 = 显存占用

7e9 × 4 = 28 GB

我刚开始做分布式训练的时候，总以为参数占大头就够了。
但别忘了，梯度是和参数一一对应的。训练时你存了一份梯度，它的大小和参数一模一样。
所以参数 + 梯度，在 FP32 下就是 28 GB + 28 GB = 56 GB。
一张 A100 才 80 GB，这还没算别的呢。

实际训练中，我们通常用混合精度训练（AMP）。
参数用 FP16 存，梯度也用 FP16 存。每个占 2 字节。

参数（FP16）：7e9 × 2 = 14 GB
梯度（FP16）：7e9 × 2 = 14 GB
合计：28 GB

嗯，这里要小心。虽然前向传播和梯度计算时用 FP16 存，但优化器更新参数时，还是需要一份 FP32 的 Master Copy。加上去，显存就又上去了。

2. 优化器状态：被很多人忽略的显存“大头”

优化器状态是什么？说白了就是优化器在训练过程中自己需要维护的额外变量。
比如 Adam 优化器，它会为每个参数保存两个动量（一阶动量 m 和二阶动量 v）。
如果你的模型参数是用 FP16 或 BF16 存的，优化器状态通常还是用 FP32 来维护。

我记得有个项目，参数只占了 14 GB，队友觉得显存够用，结果一跑就爆了。
为什么？因为优化器状态才是真正的显存“黑洞”。

以 Adam 为例，每个参数有 m 和 v 两个状态，每个状态是 FP32（4 字节）。

优化器状态 = 参数量 × 2（状态数） × 4（字节）
          = 7e9 × 2 × 4
          = 56 GB

加上参数（FP16）和梯度（FP16）的 28 GB，总共是 84 GB。
一张 80 GB 的 A100 根本放不下。

这就是为什么要用 ZeRO 优化的原因。ZeRO 阶段 2 把优化器状态和梯度分摊到多张卡上，单个卡的显存一下子就降下来了。

组件	FP32 纯精度	混合精度（AMP）
模型参数	28 GB	14 GB
梯度	28 GB	14 GB
优化器状态（Adam）	56 GB	56 GB
合计（不含激活）	112 GB	84 GB

数字很吓人是吧？没关系，这只是一张卡的情况。
分布式训练里，我们会把这些东西分配到多张卡上，显存压力就能平摊。

3. 中间激活：看不见的“隐形”杀手

前向传播的时候，每一层都会产生中间激活值。
比如 LayerNorm 的输出、Attention 的 QKV 计算结果、FFN 的中间隐藏层，这些都是临时产生的，但都占显存。
你的 batch size 越大，序列越长，中间激活就越多。

我见过最坑的一次，是有人把序列长度从 512 改到 4096 后，显存直接用光了。
中间激活的显存，大致和 batch size × sequence length × hidden size × num_layers 成正比。
一旦序列长度翻倍，中间激活可能翻好几倍。

中间激活的估算没有统一公式，因模型架构而异。但在 Transformer 模型里，它常常和参数占的显存差不多，甚至更多。
如果中间激活占了 40 GB，加上前面那 84 GB，你想想看总共要多少显存。

4. 总览：显存占用的结构图

下面这张图，把四大显存组件串了一遍。你看看这个结构，每个环节都能对应上我们刚才讲的内容。

5. 实战：动手算算你的显存

反正我个人习惯，在开始训练前一定会拿笔算一遍。
你可以用下面这个简单的脚本，估算一下自己的显存够不够。

import torch

# 假设你的模型参数
num_params = 7e9  # 7B 模型
dtype_bytes = 2   # FP16 / BF16

# 参数 + 梯度
param_mem = num_params * dtype_bytes
grad_mem = num_params * dtype_bytes

# 优化器状态：假设 Adam，FP32 存两个状态
optim_mem = num_params * 2 * 4

# 中间激活估算：这个需要根据层数和隐藏层大小来算
# 这里只是一个粗略估计，具体要看模型架构
# 假设 batch=1, seq=4096, hidden=4096, layers=32
batch_size = 1
seq_len = 4096
hidden_dim = 4096
num_layers = 32
activation_mem = batch_size * seq_len * hidden_dim * num_layers * 2  # 单位字节

print(f"参数: {param_mem / 1e9:.2f} GB")
print(f"梯度: {grad_mem / 1e9:.2f} GB")
print(f"优化器状态: {optim_mem / 1e9:.2f} GB")
print(f"中间激活估算: {activation_mem / 1e9:.2f} GB")
print(f"合计: {(param_mem + grad_mem + optim_mem + activation_mem) / 1e9:.2f} GB")

6. 优化策略：从哪下手最有效？

讲真，优化显存不能眉毛胡子一把抓。我们要找收益最大的方向。
我一般按这个优先级来：

1. 优化器状态：使用 ZeRO 阶段 2 或 3，把优化器状态分到多卡上。这是见效最快的一步。
2. 中间激活：启用 Gradient Checkpointing，前向传播时不存所有激活，反向传播时重算。用时间换空间。
3. 参数与梯度：用混合精度训练，把参数和梯度从 FP32 降到 FP16 或 BF16。
4. 模型架构：如果前面的方法都用过了还是 OOM，那就要考虑用更小的模型，或者减少序列长度、batch size。

我曾在一个 70B 模型的项目里，只用 ZeRO 阶段 3 + 混合精度 + Gradient Checkpointing，就把单卡显存从 80 GB 降到了 30 GB 左右。
效果立竿见影，你自己试一下就知道了。