显存管理与 OOM 排查

一句话定位：显存 OOM 是 GPU 上最高频的故障，排查的关键是分清**”真的用满了”还是“分配器占着没还”**——这与 Spark Executor OOM 的排查思路完全同构。

1. 常见成因

• batch / 序列过长：显存占用随 batch × 序列长度增长；注意力与 KV Cache（3.4）尤其敏感，是最常见的直接原因。
• 显存碎片化：反复分配释放不同大小的张量后，空闲显存总量够但没有足够大的连续块，于是分配失败。表现为”明明还剩几个 G 却 OOM”。
• 中间张量未释放：Python 变量仍持有引用（如把 loss/激活存进 list 做日志、未 detach() 导致计算图与激活被一直保留）。
• 缓存分配器未归还：PyTorch 使用缓存分配器（caching allocator），释放的张量显存进入其内部缓存池而不立即归还驱动，以加速后续分配。这会让 nvidia-smi 看到的占用远高于实际张量用量，也可能挤压其他进程。

2. 排查手段

• nvidia-smi：看进程级总占用与显存上限，判断是否被其他进程/多卡分配占用。注意它反映的是分配器保留的显存，不等于张量实际使用量。
• PyTorch 显存 API：
  • torch.cuda.memory_allocated()：实际张量占用；
  • torch.cuda.memory_reserved()：分配器保留总量；
  • torch.cuda.memory_summary()：分类汇总，快速看出碎片与峰值；
  • 显存快照（torch.cuda.memory._record_memory_history() + snapshot 可视化）：定位是哪段代码分配了未释放的张量。
• 核心诊断动作：分析 reserved 与 allocated 的差值
  • 差值大 → 说明分配器缓存了大量碎片显存，属碎片化问题：可尝试 torch.cuda.empty_cache()（归还空闲缓存）、调整 PYTORCH_CUDA_ALLOC_CONF（如 expandable_segments）、统一 batch/序列长度以减少大小多样性。
  • 差值小但 allocated 已接近上限 → 是真实容量不足：只能降 batch、缩序列、梯度累积/checkpointing、量化（4）或换更大显存的卡。

参考：PyTorch 官方文档 CUDA semantics