3、常见OOM场景：模型过大、Batch Size过大、序列长度过长、中间激活值爆炸
做深度学习这几年，我见过最多的报错就是OOM。说白了就是显存不够用了。但有意思的是，很多人一看到OOM就慌了，然后盲目调小batch size。其实OOM背后有四种典型场景，每种场景的解法完全不同。
我个人习惯是，遇到OOM先问自己三个问题：模型多大？batch多大？序列多长？这三个问题问完，基本就能定位到根因了。
3.1 模型过大：参数撑爆显存
这是最直观的场景。模型参数量太大，光是把参数加载到显存里就爆了。
为什么会这样？ 你想想看，一个7B参数的模型，用FP16存储，光参数就要占14GB显存。再加上优化器状态、梯度，轻松突破24GB。
显存占用估算公式（训练时）：
总显存 ≈ 模型参数 × (4 + 4 + 8) # 参数、梯度、优化器状态
+ 中间激活值 × batch_size × 序列长度其中4+4+8对应FP32参数、FP32梯度、Adam优化器状态（2个动量项）。
我在项目中遇到过一位同学，用A100 80G跑LLaMA-13B，一启动就OOM。我让他算了一下：13B × 16字节 ≈ 208GB。嗯，一张卡肯定不够。
我的建议： 先算清楚模型占多少，再决定用几张卡。一张80G的A100，最多能装5B参数的模型（FP16训练）。
3.2 Batch Size过大：数据并行撑爆显存
Batch size过大导致的OOM，其实不是参数的问题，是中间激活值的问题。
举个例子：ResNet-50，参数才25M，很小对吧？但如果你把batch size设成1024，输入图片是224×224×3，中间每一层的特征图都会乘以1024。这些中间激活值加起来，比模型参数大得多。
| 模型 | 参数量 | batch=1激活值 | batch=256激活值 |
|---|---|---|---|
| ResNet-50 | 25M | ~200MB | ~50GB |
| ViT-Large | 307M | ~1.5GB | ~384GB |
你看，batch size从1涨到256，激活值涨了256倍。这就是为什么很多人调大batch size就OOM。
避坑指南： 我曾经在训练ViT时，为了加速把batch size从64调到128，结果直接OOM。后来才发现，ViT的中间激活值比CNN大得多，因为自注意力层要保存QKV矩阵。所以调batch size前，先看看你的模型结构。
3.3 序列长度过长：Transformer的显存杀手
做NLP的朋友应该深有体会。Transformer的显存占用和序列长度是二次方关系。为什么？因为自注意力机制要计算N×N的注意力矩阵。
具体来说：
- 序列长度1024：注意力矩阵大小 1024×1024 ≈ 1M 元素
- 序列长度4096：注意力矩阵大小 4096×4096 ≈ 16M 元素
- 序列长度16384：注意力矩阵大小 16384×16384 ≈ 268M 元素
序列长度翻4倍，注意力矩阵大小翻16倍。这就是二次方爆炸。
显存占用公式（Transformer单层）：
注意力显存 = batch_size × num_heads × seq_len² × 2 # QK^T 和 softmax 输出
+ batch_size × seq_len × hidden_dim × 4 # Q、K、V、输出投影当seq_len很大时，第一项占主导。
我记得有一次帮朋友调一个长文本摘要模型，序列长度设成了8192。他用的A100 40G，batch size=1都OOM。我让他把序列长度降到2048，batch size反而能开到8，训练速度还更快了。
实用技巧： 如果必须处理长序列，可以用这些方法：
- 梯度检查点：用时间换空间，中间激活值不保存，反向传播时重新计算
- Flash Attention：分块计算注意力，避免一次性生成N×N矩阵
- 序列切分：把长序列切成多个短序列，分别处理再拼接
3.4 中间激活值爆炸：被忽略的显存黑洞
这个场景最容易被忽略。很多人以为显存只用来存参数，其实中间激活值才是大头。
什么是中间激活值？就是前向传播时每一层的输出。反向传播时需要这些值来计算梯度。所以它们必须保存在显存里，直到反向传播结束才能释放。
我画了一张图，帮你理解显存里到底放了什么：
你看，中间激活值占了将近一半的显存。而且它和batch size、序列长度直接相关。这就是为什么有时候模型参数不大，但训练时还是OOM。
我曾经踩过的坑： 训练一个12层的Transformer，模型参数才300M，但batch size=32时OOM了。我一开始以为是参数太大，后来用torch.cuda.memory_summary()一看，好家伙，中间激活值占了22GB，参数才2GB。后来用了梯度检查点，激活值降到6GB，batch size直接开到128。
3.5 四种场景的快速诊断方法
遇到OOM时，别慌。按这个流程走一遍：
- 先看模型参数：用
model.parameters()算一下总参数量，乘以数据类型字节数，看看是不是参数就爆了 - 再看batch size：把batch size降到1试试，如果还OOM，说明是模型本身太大；如果不OOM了，那就是batch size的问题
- 检查序列长度：如果是Transformer，把序列长度减半试试，显存降为原来的1/4
- 最后看激活值：用
torch.cuda.max_memory_allocated()看峰值显存，减去参数和梯度，剩下的就是中间激活值
实用工具推荐：
torch.cuda.memory_summary()：查看显存分配详情nvidia-smi：实时监控显存使用pyTorch Profiler：分析每层显存占用
说白了，OOM不可怕，可怕的是不知道哪里爆了。把这四种场景记在心里，下次遇到OOM，你就能快速定位问题，对症下药。