7.1 显存估算速查与 OOM 排查大全

本章目标：把 OOM 从”玄学崩溃”变成”看一眼就知道谁在偷显存”。这是卷7 工程手册的开卷章——你未来两年报障最多的问题没有之一。学完标准：能给任意”模型×卡×配置”算出显存账，拿到 OOM 按决策树 10 分钟定位。

1. 训练显存总账（一张表吃遍天）

$\text{总显存} = \text{权重} + \text{梯度} + \text{优化器状态} + \text{激活} + \text{KV/临时}$

组件	每参数字节	条件
权重	2（bf16）/ 4（fp32）/ 0.5（4bit）	由加载 dtype 决定
梯度	2（bf16）/ 4（fp32）	只算 requires_grad 的参数（LoRA 只占增量部分！）
AdamW 状态	8（fp32 两个动量）	同上，只算可训练参数
fp32 主权重	4	混合精度训练才有

两种典型配置的算例（7B）：

全参 bf16 混合精度：2+2+8+4 = 16 B/param → 112 GB + 激活 → 单卡无解；
LoRA（rank 32，可训练 0.5%）：权重 14 GB（冻结，只存不更新）+ 梯度/状态 ≈ 0.7 GB + 激活 → 24 GB 卡上还剩 ~8 GB 给激活。这就是全书的”LoRA 甜点”底层数学。

激活值的账（大头且常被忽略）：约等于 seq_len × batch × hidden × layers × 常数(10~30) × bytes——随序列长度线性、随 batch 线性。视觉 token 多（高分辨率）或长文本时它先爆。gradient_checkpointing 用重算换存储，激活降到约 1/√层倍数级别，代价 ~30% 速度。

2. 推理显存账（2.3/2.8 的合并速查）

权重：同上表；
KV cache = 2 × layers × d_kv × seq × batch × bytes（GQA 后 d_kv 缩小数倍）；
临时峰值：采样/解码时的 logits（vocab_size × batch，大词表模型不小）。

例：7B（L=32, d=4096, MHA, bf16）4k 上下文单条 KV ≈ 2.1 GB（2.3 已算）——并发服务里 KV 通常比权重先爆，所以 vLLM 把 gpu-memory-utilization 和 max-model-len 当第一杠杆。

3. OOM 排查决策树（按此顺序，10 分钟定位）

第 0 步：先测量，别猜。torch.cuda.max_memory_allocated()/1e9（训练循环里打印峰值）、nvidia-smi（谁在占卡）。多进程残留（上次崩掉的 python 僵尸、未释放的 vLLM）是”明明没跑啥却 OOM”的第一嫌疑——nvidia-smi 里找进程，kill 干净。

降 batch / 加梯度累积（激活线性降，等效更新不变）——成本最低，先动它；
gradient checkpointing（激活换时间）——LoRA/微调场景标配；
查 dtype 混用：某个模块没被 autocast 包住、偷偷 fp32；4bit 底座配了 fp16 的 LoRA——统一 bf16；
查”忘记 no_grad 的循环”：训练中途的验证循环没包 inference_mode() → 激活全存梯度 → 峰值翻倍（新手 OOM 榜第一）；
查”tensor 收集器”：日志列表里存了 loss tensor / logits（应立刻 .item() / 只存标量）；
序列与分辨率：cutoff_len 降档、图片分辨率降档、视觉 token 压缩（3.3）——长尾样本（最长那条文本/最大那张图）引起的 OOM：先跑 max(len) 统计；
碎片化：峰值远小于总量却 OOM → PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True；
最后手段：FSDP/DeepSpeed ZeRO（7.3 的主题——优化器状态分片，单机多卡或多机）。

4. 速查卡片（贴墙上）

卡	全参训	LoRA 微调	4bit LoRA	推理
24G（4090）	≤2B	≤7B	≤13B（紧）	≤13B bf16
48G（A6000/双4090）	≤4B	≤14B	≤30B	≤30B 4bit
80G（A100）	≤13B	≤70B LoRA	70B	70B 4bit / 13B bf16

（量级参考，实际再留 20% 余量给激活与框架开销。）

5. 自测题

全参 bf16 混合精度每参数 16 字节怎么组成？（权重2+梯度2+AdamW 8+fp32主权重4）
为什么 LoRA 的梯度/优化器状态只占增量参数？（只有 requires_grad 的参数才有梯度和动量）
激活显存随什么线性增长？两个降它的手段？（seq×batch×hidden×layers；降 batch/降 seq、gradient checkpointing）
“训练中途验证就 OOM”的机理？（验证循环没包 no_grad/inference_mode，激活带梯度图）
峰值远小于显存总量却 OOM，最可能是什么？（碎片化——expandable_segments；或多进程残留）
长尾样本 OOM 的预防？（先统计 seq len / 分辨率分布，cutoff 或分桶）
推理服务先爆的是权重还是 KV？（并发场景 KV 先爆——max-model-len×并发的预算要算）

6. 动手作业

写 mem_report.py：打印权重/梯度/优化器分段显存 + 峰值（torch.cuda.memory_summary），在 2.6 的 LoRA 训练里挂上，对照 §1 的账验证误差 <10%；
故意制造五种 OOM（忘 no_grad / 存 tensor 列表 / batch 拉满 / dtype 混用 / 僵尸进程），每种用 §3 决策树定位修复——你的第一次 OOM 演习；
给 8.9 选定题目的模型配置算一遍显存账，写进 proposal 的”算力与时间”栏。

图解 · OOM 排查决策树

flowchart TD
  O["OOM！"] --> M0["第 0 步：先测量<br>max_memory_allocated + nvidia-smi"]
  M0 --> Z{"僵尸进程 / 残留占用？"}
  Z -->|是| K["kill 干净，重跑"]
  Z -->|否| B["① 降 batch + 梯度累积"]
  B --> GC["② 开 gradient checkpointing"]
  GC --> D["③ 查 dtype 混用 / 忘 no_grad / tensor 收集器"]
  D --> L["④ 降 seq len / 分辨率 / 压视觉 token"]
  L --> FR{"峰值 ≈ 总量仍 OOM？"}
  FR -->|是，碎片| FS["PYTORCH_CUDA_ALLOC_CONF<br>=expandable_segments:True"]
  FR -->|否，真不够| FSD["FSDP / DeepSpeed ZeRO（7.3）"]

使用方式：OOM 时从顶部往下走，每走一格只改一个变量并重测——“一次改三处再跑”会让你永远不知道是哪招起效的（这也是所有调试的元规则）。

7. 延伸资源

PyTorch 官方《Understanding GPU Memory》博客——官方版账本
HuggingFace《GPU inference》文档页——推理侧速查
DeepSpeed / FSDP 文档——单卡装不下时的下一站（7.3 展开）

云镜检索

稍后阅读