1、显存危机：为什么你的模型总是OOM？——显存分配机制与常见瓶颈分析

做深度学习这几年，我见过太多人卡在OOM（Out Of Memory）这一步了。模型写好了，数据准备好了，一跑就崩。说白了，显存就是一块有限的地皮，你盖的楼太高，地基不够大，自然就塌了。

我个人习惯，拿到一个新模型，第一件事不是调参，而是先算一笔账——这模型到底要吃多少显存？很多人觉得这是玄学，其实不是。显存的分配机制，说白了就那几块，搞清楚了，OOM就不再是玄学。

1.1 显存到底被谁吃了？

你想想看，一个模型跑起来，显存里塞了哪些东西？我列个清单给你看：

• 模型参数：权重、偏置，这些是固定的
• 梯度：反向传播时要用，大小跟参数一样
• 优化器状态：Adam有动量和方差，这玩意儿比参数还占地方
• 中间激活值：前向传播时每一层的输出，这个最要命
• 临时缓冲区：框架自己搞的一些缓存，比如cuDNN的workspace

我在项目中遇到过最坑的一次，是训练一个Transformer模型。我算好了参数和梯度，觉得32GB显存绰绰有余。结果一跑，直接OOM。查了半天才发现，中间激活值占了将近20GB。嗯，这里要注意，激活值的大小跟batch size和序列长度是线性关系，很多人忽略了这一点。

1.2 显存分配机制：框架到底在干嘛？

很多人以为，PyTorch或者TensorFlow会智能地管理显存。其实不是。框架的显存分配策略，说白了就两种：

PyTorch用的是第二种。它启动时会先向GPU申请一大块显存，然后自己内部做内存池管理。这样做的好处是快，坏处是——你明明只用了10GB，它可能已经占了20GB。为什么？因为它预分配了。

我记得有一次调试一个模型，用nvidia-smi一看，显存占了18GB，但实际模型只用了8GB。我当时就懵了。后来查了PyTorch的文档，发现有个环境变量PYTORCH_CUDA_ALLOC_CONF可以控制这个行为。设置max_split_size_mb之后，情况好多了。

1.3 常见瓶颈：OOM的三大元凶

根据我的经验，OOM的原因逃不出下面三个：

1.3.1 Batch Size 太大

这个最直观。batch size翻倍，激活值也翻倍。很多人为了训练稳定，把batch size设得很大，结果显存直接爆了。我建议，先设一个很小的batch size（比如2或4），跑通了再慢慢往上加。

1.3.2 序列长度太长

NLP模型尤其容易中招。Transformer的self-attention复杂度是O(n²)，序列长度翻倍，显存占用翻四倍。我曾经处理过一个长文本分类任务，序列长度从512改成1024，显存直接从12GB飙到了40GB。嗯，这就是为什么很多模型要限制输入长度。

1.3.3 优化器状态太大

Adam优化器需要保存动量和方差，每个参数对应两个额外的值。如果你的模型有1亿参数，Adam就要额外占用2.4GB（1亿 × 2 × 4字节 × 3倍）。如果你用FP16，能省一半，但精度可能会掉。

1.4 知识体系：显存危机的全景图

下面这张图，是我自己总结的显存危机分析框架。你看一眼，基本就能定位问题出在哪了。

1.5 实战：如何快速定位OOM？

如果你现在遇到了OOM，别慌。按这个步骤来：

1. 先看nvidia-smi：确认显存占用情况，是不是被其他进程占了
2. 用torch.cuda.memory_summary()：看PyTorch内部的内存分配详情
3. 逐步减小batch size：直到能跑通，然后算一下每样本的显存消耗
4. 检查序列长度：是不是有特别长的样本？可以截断或者用动态padding
5. 看优化器：是不是用了Adam？能不能换成SGD或者AdamW的变体？

# 显存预算：
# 参数: 2.1GB (1.2B参数 × FP16)
# 梯度: 2.1GB
# 优化器状态: 4.2GB (Adam)
# 激活值: 8GB (batch=16, seq=512)
# 总计: 16.4GB → 需要24GB显存

好了，这一章我们聊了显存危机的根源。说白了，OOM不是玄学，是数学。算清楚每一块显存去哪了，问题就解决了一半。下一章，我会讲怎么用显存池技术来复用显存，把峰值降下来。