深夜排查 OOM：显存不够时先找谁在偷吃

深夜排查 OOM显存不够时先找谁在偷吃一、凌晨两点的 CUDA Out of Memory场景回放训练已经跑了 12 小时第 3 个 epoch 刚过半突然报错CUDA out of memory。检查显存使用——A100 80GB 怎么可能不够但 PyTorch 确实报了 OOM。重启训练调整 batch_size 从 64 降到 32。又跑了 4 小时再次 OOM。时间到了凌晨两点。显存排查的困难在于PyTorch 的显存分配不是透明的。它使用缓存分配器Caching Allocator释放的显存不会立即归还给 CUDA而是留在缓存中备用。nvidia-smi显示的显存占用包含了缓存torch.cuda.memory_allocated()显示的只是当前活跃的分配。两者之间的差值就是缓存中的僵尸显存。更隐蔽的显存占用来源前向传播的中间激活Activations、优化器状态Adam 的 momentum 和 variance、梯度累积、以及 DataLoader 的预取缓存。这些每一项都可能吃掉几 GB 到几十 GB 的显存。见证奇迹的时刻是当你终于定位到那个在for循环中不断.append()却没有detach()的 tensor 列表——它默默吞掉了 15GB 显存而你在排查的前四个小时里完全没注意到它。二、显存占用的完整解剖graph TD A[显存总占用] -- B[模型参数] A -- C[优化器状态] A -- D[前向激活] A -- E[梯度] A -- F[其他] B -- B1[FP32 参数: P * 4 bytes] B -- B2[FP16 参数: P * 2 bytes] B -- B3[混合精度时两者同时存在] C -- C1[Adam: P * 8 bytesbr/(momentum variance)] C -- C2[SGD: P * 4 bytesbr/(仅 momentum)] C -- C3[FP16 时额外: P * 4 bytesbr/(FP32 master weights)] D -- D1[激活 f(batch_size, seq_len, hidden_dim, num_layers)] D -- D2[Transformer: ~12 * batch_size * seq_len * hidden_dim] D -- D3[梯度检查点可减少 50%~80%] E -- E1[FP32 梯度: P * 4 bytes] E -- E2[梯度累积: 与正常训练相同] F -- F1[DataLoader 预取缓存] F -- F2[CUDA Context: ~1GB 基础开销] F -- F3[未 detach 的中间 tensor] F -- F4[日志/输出缓冲区] style D fill:#fce4ec style F3 fill:#ff5252显存占用的四大部分模型参数Model ParametersFP32 下每个参数占 4 bytesFP16 下 2 bytes。混合精度训练时两者同时存在FP16 用于前向FP32 master weights 用于更新。优化器状态Optimizer States这是最大头的显存占用。Adam 优化器为每个参数维护两个状态momentum 和 variance各占 4 bytes总计 8 bytes/参数。加上 FP32 master weights4 bytesAdam 混合精度的优化器状态开销是 12 bytes/参数——是参数本身的 6 倍。前向激活Activations前向传播中每层的中间结果因为反向传播需要它们来计算梯度。这是训练显存的第二大来源。Transformer 模型的前向激活量级约为12 * batch_size * seq_len * hidden_dim。梯度GradientsFP32 下每个参数的梯度占 4 bytes。三、OOM 排查的系统工具箱以下代码展示了一套显存诊断工具。import torch import gc from typing import Dict, List, Tuple from dataclasses import dataclass dataclass class MemoryReport: 显存诊断报告 total: int # GPU 总显存 allocated: int # PyTorch 当前分配 reserved: int # PyTorch 缓存含碎片 free_in_cache: int # 缓存中未使用 nvidia_smi_used: int # nvidia-smi 显示的使用量 class OOMDetective: OOM 排查工具集 设计原因提供分层的显存分析工具 从总量诊断到逐层剖析帮助定位显存泄漏源。 staticmethod def memory_report(gpu_id: int 0) - MemoryReport: 获取显存使用全景 设计原因同时对比 PyTorch 内部统计和 nvidia-smi 数据 差异部分 缓存碎片 非 PyTorch 分配。 total torch.cuda.get_device_properties(gpu_id).total_mem allocated torch.cuda.memory_allocated(gpu_id) reserved torch.cuda.memory_reserved(gpu_id) # 缓存中空闲 缓存总量 - 实际分配 free_in_cache reserved - allocated # nvidia-smi 显示的占用 ≈ reserved近似 nvidia_used reserved return MemoryReport( totaltotal, allocatedallocated, reservedreserved, free_in_cachefree_in_cache, nvidia_smi_usednvidia_used, ) staticmethod def format_bytes(b: int) - str: 格式化字节数 for unit in [B, KB, MB, GB]: if b 1024: return f{b:.2f} {unit} b / 1024 return f{b:.2f} TB staticmethod def diagnose_oom(model: torch.nn.Module, optimizer: torch.optim.Optimizer, batch_size: int, seq_len: int, hidden_dim: int, num_layers: int, gpu_id: int 0) - Dict[str, str]: OOM 根因分析 设计原因分别估算模型参数、优化器状态、激活和梯度四个部分 的显存占用与总显存对比找出超标部分。 # 参数显存 total_params sum(p.numel() for p in model.parameters()) param_mem_fp32 total_params * 4 # 优化器状态显存假设 Adam 混合精度 # 设计原因Adam 的 momentum(4B) variance(4B) # FP32 master weights(4B) 12B/参数 optimizer_mem total_params * 12 # 前向激活显存Transformer 经验公式 # 设计原因这是粗略估计实际受模型结构影响。 # 精确值需要使用 torch.cuda.memory_summary() 获取。 activation_mem batch_size * seq_len * hidden_dim * num_layers * 2 # 梯度显存 grad_mem total_params * 4 total_estimated (param_mem_fp32 optimizer_mem activation_mem grad_mem) report OOMDetective.memory_report(gpu_id) return { 模型参数: OOMDetective.format_bytes(param_mem_fp32), 优化器状态: OOMDetective.format_bytes(optimizer_mem), 前向激活(估算): OOMDetective.format_bytes(activation_mem), 梯度: OOMDetective.format_bytes(grad_mem), 估算总计: OOMDetective.format_bytes(total_estimated), GPU 总显存: OOMDetective.format_bytes(report.total), 估算占用率: f{total_estimated / report.total * 100:.1f}%, } staticmethod def find_leaked_tensors() - List[Tuple[str, int]]: 查找未被释放的 tensor 设计原因遍历 Python 堆中的所有 CUDA tensor 找出那些没有 grad_fn 且未被引用的孤儿 tensor。 这些 tensor 通常是显存泄漏的元凶。 leaked [] for obj in gc.get_objects(): try: if (torch.is_tensor(obj) and obj.is_cuda and obj.grad_fn is None): # 检查是否有 Python 引用 ref_count len(gc.get_referrers(obj)) if ref_count 2: # 仅有循环引用和自身 leaked.append(( str(obj.shape), obj.element_size() * obj.numel() )) except Exception: pass # 按大小排序 leaked.sort(keylambda x: x[1], reverseTrue) return leaked[:10] staticmethod def clear_cache(gpu_id: int 0): 清理显存缓存 设计原因在 OOM recovery 时必须先清理缓存 然后手动调用 Python GC 释放所有中间 tensor 引用。 两者的顺序很重要——先 CUDA cache后 Python GC。 torch.cuda.empty_cache() gc.collect() torch.cuda.synchronize(gpu_id)四、显存管理的经验法则与异常边界排查 OOM 时有一些经过验证的经验法则作为参考。数量级估算模型规模参数显存优化器参数显存训练总显存估算7B28 GB (FP32)84 GB (AdamFP32)~160 GB13B52 GB156 GB~300 GB70B280 GB840 GB~1600 GB常见 OOM 陷阱DataLoader 的预取num_workers 0时DataLoader 在主进程外预取数据。预取的数据已经转换为 CUDA tensor占用了显存但不计入 PyTorch 的memory_allocated统计。动态扩展的 tensor 列表在循环中.append(tensor)但未.detach()这些 tensor 保留了完整的计算图导致显存线性增长。评估模式的激活缓存torch.no_grad()只阻止梯度计算不释放前向传播的中间激活。如果评估时 batch_size 设为训练时的 2 倍因为没有梯度消耗可能反而触发 OOM。见证奇迹的时刻在于不是找到了某个吃掉 50GB 的巨兽而是发现了三个各吃 5GB 的小偷——它们分散在不同模块中单独看都不起眼合在一起却占满了 80GB 显存。排查需要的是系统性思维而不是直觉。五、总结训练 OOM 的排查需要从模型参数、优化器状态、前向激活和梯度四个维度逐一分析。Adam 优化器的状态显存是参数本身的 3 倍加上 FP32 master weights 达到 6 倍是显存的最大消耗来源。梯度检查点Gradient Checkpointing以计算换显存可以将前向激活占用降低 50%~80%。DataLoader 预取缓存和未 detach 的中间 tensor 是常见的隐性显存泄漏源。显存排查应先从总量诊断对比分配 vs 缓存 vs nvidia-smi入手再逐层定位到具体的问题模块。必要时使用torch.cuda.memory_summary()获取分段统计信息。