3、显存瓶颈剖析：显存不足导致OOM、显存碎片化、模型状态与优化器状态显存占用分析

显存瓶颈，说白了就是你的GPU“内存”不够用了。很多同学跑大模型训练，上来就报CUDA out of memory，然后一脸懵。我刚开始做分布式训练时，也经常被OOM搞得焦头烂额。后来发现，其实大部分显存问题，都是可以提前算出来的。

今天我们就来彻底拆解一下，显存到底被谁吃了？怎么吃？能不能省着点吃？

3.1 显存都去哪了？——模型状态与优化器状态

很多人以为显存只存模型参数。其实不对。你想想看，训练过程中，GPU不仅要存参数，还要存梯度、优化器状态、中间激活值。我习惯把这部分叫做“训练四件套”。

咱们先看一个经典公式：

显存占用 ≈ 模型参数 + 梯度 + 优化器状态 + 激活值 + 临时缓冲区

其中，模型参数和梯度是必须的。优化器状态呢？如果你用Adam，那就要额外存两份：一份是动量（momentum），一份是方差（variance）。

为什么会这样？我们来算一笔账：

组件	数据类型	每参数占用	说明
模型参数	FP16	2 bytes	必须
梯度	FP16	2 bytes	必须
优化器状态（动量）	FP32	4 bytes	Adam需要
优化器状态（方差）	FP32	4 bytes	Adam需要
合计		12 bytes	参数量的12倍

嗯，这里要注意：如果你用SGD优化器，那优化器状态就少很多，只需要一份动量（FP32），总占用降到8倍。但Adam是主流，所以12倍是个很现实的数字。

我在项目中遇到过，一个7B参数的模型，光参数和优化器状态就要吃掉：7B × 12 bytes ≈ 84 GB。单张A100（80GB）根本放不下。怎么办？要么用模型并行，要么换优化器。

3.2 显存碎片化——看不见的“黑洞”

除了显存不够用，还有一个更隐蔽的问题：显存碎片化。

你可能会问：明明显存总量还有20GB，为什么分配一个10GB的张量却报OOM？

这就是碎片化在作祟。GPU显存分配和内存分配类似，频繁申请释放会导致“外碎片”。比如你连续分配了4个4GB的张量，然后释放了第2个和第4个。此时空闲空间是8GB，但被分割成两个4GB的碎片。你想分配一个6GB的张量？对不起，放不下。

怎么避免？我建议：

• 预分配大块显存：比如用torch.cuda.empty_cache()不如一开始就分配好。
• 减少动态张量创建：尽量复用已有的张量，用inplace操作。
• 使用显存池：像DeepSpeed的ZeRO-3就自带显存管理，能有效减少碎片。

3.3 激活值——被低估的显存杀手

很多人只盯着模型参数，却忽略了激活值。其实，激活值才是训练时的显存大户。

举个例子：一个Transformer层，输入序列长度1024，隐藏维度4096，batch size为8。这一层的激活值有多大？

激活值 ≈ batch_size × seq_len × hidden_dim × 层数 × 精度
       = 8 × 1024 × 4096 × 32 × 2 bytes
       ≈ 2.1 GB

注意，这只是一层！如果你有32层Transformer，那激活值就是67 GB。加上参数和优化器状态，显存直接爆炸。

3.4 显存优化三板斧

既然知道了显存被谁吃了，那怎么优化？我总结了三个方向：

1. 减少模型状态：用混合精度训练（FP16/BF16），参数和梯度减半。或者用ZeRO优化器，把优化器状态分散到多个GPU上。
2. 减少激活值：用激活值检查点（activation checkpointing），以计算换显存。简单说就是前向时不保存所有激活值，反向时重新计算。能省50%-80%的激活显存。
3. 减少碎片化：统一张量生命周期，避免频繁申请释放。用torch.cuda.memory_summary()查看碎片情况。

下面这张图，是我自己整理的显存优化决策流程，你可以参考：

3.5 实战：如何精确估算显存

最后，分享一个我常用的显存估算方法。你可以在训练前先算一遍，避免跑起来才发现OOM。

def estimate_memory(model_size_params, batch_size, seq_len, hidden_dim, num_layers, use_adam=True):
    # 参数和梯度
    param_mem = model_size_params * 2  # FP16
    grad_mem = model_size_params * 2   # FP16
    
    # 优化器状态
    if use_adam:
        opt_mem = model_size_params * 8  # FP32动量+方差
    else:
        opt_mem = model_size_params * 4  # FP32动量
    
    # 激活值（粗略估计）
    act_mem = batch_size * seq_len * hidden_dim * num_layers * 2  # FP16
    
    total = param_mem + grad_mem + opt_mem + act_mem
    return total / (1024**3)  # 转换为GB

# 示例：7B模型，batch=4，seq=2048，hidden=4096，layers=32
mem = estimate_memory(7e9, 4, 2048, 4096, 32)
print(f"预估显存: {mem:.2f} GB")

这个估算值通常偏保守，但足够让你判断是否需要上模型并行或者ZeRO。

好了，显存这块我们就聊到这。记住一句话：显存不是无限的，但优化空间是无限的。只要搞清楚谁在吃显存，你就能对症下药。