第二讲：显存去哪儿了？模型参数、梯度、优化器状态、中间激活值——一张图看懂显存分配

大家好，我是你们的老朋友。

上一讲我们聊了显存溢出有多坑。今天咱们来点实在的——显存到底被谁吃了？

很多同学跑模型一爆显存就慌，其实你只要搞清楚显存里住了哪几位“房客”，问题就解决了一半。我个人习惯把显存想象成一个大房子，里面住着四类人：模型参数、梯度、优化器状态、中间激活值。今天我们就挨个拜访一下。

一、显存里的四大“房客”

先上一张图，这是我手绘的显存分配全景图。你看完就明白，为什么一张显卡那么贵，却总是不够用。

显存分配全景图 模型参数 权重矩阵 + 偏置项 约占总显存 20%~40% 梯度 反向传播计算的梯度值 与模型参数大小相同 优化器状态 Adam: 一阶/二阶动量 SGD: 几乎为0 中间激活值 前向传播的中间结果 显存大头！占 40%~60% 关键公式： 总显存 ≈ 模型参数 × (1 + 梯度系数 + 优化器系数) + 中间激活值 • 纯推理：总显存 ≈ 模型参数 + 中间激活值（无梯度、无优化器） • 训练（SGD）：总显存 ≈ 模型参数 × 2 + 中间激活值 • 训练（Adam）：总显存 ≈ 模型参数 × 4 + 中间激活值

嗯，这张图你保存下来，以后每次爆显存就拿出来对照一下。说白了，显存就这四块，我们一个一个拆解。

二、模型参数：最老实的房客

模型参数是最好算的。一个 7B 的模型，用 FP16 存储，参数占多少？

7 × 10⁹ × 2 bytes = 14 GB

就这么简单。但注意，这只是权重本身。我在项目中遇到过有人把 embedding 层和输出层的参数重复计算，结果显存预算翻了一倍。你想想看，一个 7B 模型，embedding 层可能就占了 0.5B，别小看它。

小技巧： 用 model.parameters() 遍历一下，把每个参数 tensor 的 numel() 加起来，再乘以 dtype 的字节数，就是精确的参数量。别信网上说的“7B 模型就是 14GB”，那是理想情况。

三、梯度：参数的双胞胎兄弟

梯度的大小，和模型参数一模一样。每个参数对应一个梯度值，dtype 也相同。

所以训练时，模型参数占一份，梯度占一份。这就是为什么纯推理只需要一半显存——因为不需要存梯度。

我记得有一次帮一个团队排查 OOM，他们用 Adam 训练 13B 模型，怎么算显存都不够。结果发现他们把梯度也存成了 FP32，而参数是 FP16。这一下就多出了 13GB 的额外开销。嗯，这里要注意：梯度的 dtype 最好和参数保持一致，除非你有特殊理由。

避坑指南： 我曾经见过有人用 torch.no_grad() 包住了整个训练循环，结果梯度一直为 None，模型根本不更新。排查了一整天，最后发现是上下文管理器用错了位置。所以，检查梯度是否存在，用 model.parameters() 里的 .grad 属性看一眼就行。

四、优化器状态：最贪吃的房客

优化器状态才是真正的显存大户。尤其是 Adam，它需要维护两个动量：一阶动量（m）和二阶动量（v）。

咱们算一笔账：

优化器 额外显存开销 7B 模型（FP16） 13B 模型（FP16）
SGD 几乎为 0 ≈ 0 GB ≈ 0 GB
SGD + Momentum 1 倍参数量 ≈ 14 GB ≈ 26 GB
Adam 2 倍参数量 ≈ 28 GB ≈ 52 GB
AdamW 2 倍参数量 ≈ 28 GB ≈ 52 GB

看到了吗？Adam 的优化器状态，比模型参数本身还大两倍。所以很多大模型训练用 SGD 或者 Adam 的变体，不是没有道理的。

我个人习惯，在显存紧张时，优先考虑 混合精度训练 + AdamW，把优化器状态也降到 FP16。虽然精度会有一点损失，但显存能省一半。你想想看，省下来的显存可以塞更大的 batch size，训练速度反而更快。

五、中间激活值：隐藏的显存黑洞

这是最容易被忽视的，也是最大的显存消耗来源。很多人以为显存大头是模型参数，其实错了。中间激活值才是真正的显存黑洞。

什么是中间激活值？就是前向传播时，每一层计算出来的中间结果。比如 attention 层的 Q、K、V 矩阵，还有 softmax 之后的概率分布。这些值在反向传播时还要用，所以不能丢。

咱们算一下：

假设一个 Transformer 层：
- hidden_size = 4096
- seq_len = 2048
- batch_size = 4

中间激活值 ≈ batch_size × seq_len × hidden_size × 层数 × 系数
            = 4 × 2048 × 4096 × 32 × 12（保守估计）
            ≈ 12 GB

这还只是一层！整个模型 32 层，中间激活值轻松超过 30 GB。而模型参数才 14 GB。所以，中间激活值才是显存的第一大用户。

重点： 中间激活值的大小和 batch_size × seq_len 成正比。想省显存，第一刀就砍 batch size 或者 seq_len。我见过有人为了省显存，把 batch size 从 8 降到 1，结果中间激活值直接少了 87.5%。代价是训练速度变慢，但至少能跑起来。

六、一张表总结：显存分配速查

组件 训练时 推理时 省显存技巧
模型参数 1 倍参数量 1 倍参数量 量化（INT8/INT4）
梯度 1 倍参数量 0 梯度累积、混合精度
优化器状态 0~2 倍参数量 0 用 SGD、优化器状态量化
中间激活值 batch_size × seq_len × hidden_size × 层数 batch_size × seq_len × hidden_size × 层数 减小 batch_size、梯度检查点

七、实战：如何快速定位显存去哪儿了？

光说不练假把式。我给你一个我常用的调试脚本，跑一下就知道显存被谁吃了。

import torch
import torch.nn as nn

def analyze_memory(model, input_ids, optimizer=None):
    """分析显存分配情况"""
    # 清空缓存
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    
    # 前向传播
    output = model(input_ids)
    loss = output.logits.mean()
    
    # 记录前向后的显存
    mem_after_forward = torch.cuda.memory_allocated()
    
    # 反向传播
    loss.backward()
    
    # 记录反向后的显存
    mem_after_backward = torch.cuda.memory_allocated()
    
    # 优化器更新
    if optimizer:
        optimizer.step()
        mem_after_optimizer = torch.cuda.memory_allocated()
    
    print(f"模型参数: {sum(p.numel() for p in model.parameters()) * 2 / 1e9:.2f} GB")
    print(f"前向后显存: {mem_after_forward / 1e9:.2f} GB")
    print(f"反向后显存: {mem_after_backward / 1e9:.2f} GB")
    print(f"中间激活值 ≈ {mem_after_forward - sum(p.numel() for p in model.parameters()) * 2 / 1e9:.2f} GB")
    print(f"梯度 ≈ {mem_after_backward - mem_after_forward:.2f} GB")
    
    if optimizer:
        print(f"优化器状态 ≈ {mem_after_optimizer - mem_after_backward:.2f} GB")

这个脚本我用了好几年，每次换模型或者调参，先跑一遍看看显存分布。你想想看，如果中间激活值占了 60%，那你再怎么压缩模型参数也没用，得从 batch size 或者序列长度下手。

我的习惯： 每次训练前，先用一个小 batch（比如 batch_size=1）跑一遍，看看显存基线。然后逐步增大 batch size，直到 OOM 的前一刻。这样你能精确知道你的显卡能承受的最大 batch size 是多少。

八、总结：显存分配的核心逻辑

好了，今天的内容就这些。总结一下核心要点：

  • 模型参数：最老实，大小固定，量化可省
  • 梯度：和参数一样大，训练时才存在
  • 优化器状态：Adam 最贪吃，SGD 最省
  • 中间激活值：隐藏的显存黑洞，和 batch_size × seq_len 成正比

记住一句话：显存不够，先看中间激活值；优化器太贵，换 SGD 或量化。下次你遇到 OOM，别慌，拿出今天这张图，对照着排查，问题基本就能定位。

想要更多资料，联系可可

← 上一章 📖 返回目录 下一章 →