【Bug已解决】CI again often fails with torch.OutOfMemoryError: CUDA out of memory 解决方案一、现象长什么样仓库的 CI 里有一条训练冒烟测试跑在一张T416GB上。它不是每次都挂而是经常挂——同一个 commit有时候全绿有时候在第二步 evaluate 时红torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.24 GiB. GPU has a total capacity of 15.75 GiB of which 0 bytes is free. Process has 14.91 GiB reserved, of which 0 bytes is reserved for allocation.奇怪的点有三个时好时坏同样的代码、同样的输入失败率大约三成不是必现。失败位置不固定有时挂在trainer.train()的第一个 step有时挂在trainer.evaluate()有时挂在generate()。本地复现难在本机A100上几乎从不挂只有 CI 的小卡才暴露。这类偶发 OOM最折磨人因为不能靠把 batch 调小硬扛调小之后它只是更低频地挂而且本地依旧复现不出来。必须找到那条随时间累积、把显存慢慢吃掉直到某一步越界的路径。二、背景CUDA 显存不是用多少占多少那么简单。PyTorch 的 caching allocator 会缓存已经释放的块方便下次复用所以显存占用通常是阶梯式上升后在一个平台震荡而不是随每个 step 线性归零。这就给排查带来两个坑你以为某一步del掉了张量实际上那块显存还在 allocator 的缓存池里没还给驱动偶发的峰值如果超过了缓存池里可用连续块就会触发真正的OutOfMemoryError而这种峰值往往来自临时对象日志里的.detach().cpu()之前在 GPU 上排了一串、或者torch.cat中间变量、或者generate的 past_key_values 没及时清。CI 上经常失败而不是必败通常意味着显存基线已经很高接近上限再叠加一个偶发的额外峰值就崩而这块额外峰值的大小受数据顺序、随机种子、DataLoader worker 抖动、甚至 Python GC 时机影响于是表现为概率性。常见制造偶发峰值的来源评估时不设torch.no_grad()导致 autograd 历史被建起来峰值翻倍。generate()没有限制max_new_tokens且past_key_values在长输出时累积。张量在进 logger / wandb 之前没.cpu()在 GPU 上排了 N 份副本。retain_graph误用或重复backward。梯度检查点没开forward 激活全留着。CUDA 缓存碎片大量不同形状的中间张量让 allocator 找不到连续块即使总空闲够也分配失败。三、根因针对我们这个冒烟测试逐步排查后定位到三个叠加因素因素 A评估漏了no_grad。测试里有一段手写评估def quick_eval(model, loader): outs [] for batch in loader: logits model(**batch) # 建了 autograd 图 outs.append(logits) # logits 留在 GPU 上图也留着 return outsmodel(**batch)默认在train()模式下还会做 dropout而且logits带着整条反向图。这一步的显存占用是训练 step 的近两倍。又因为它发生在训练若干 step 之后此时缓存池已经被训练撑到一个高位叠加起来就超过 T4 上限。因素 B长尾样本。DataLoader 偶尔吐出一个超长序列的样本数据里有几条没截断干净的generate()的past_key_values随长度线性增长偶发地把峰值顶破上限。因素 C碎片。训练 step 产生的激活形状很多样allocator 缓存里全是碎片遇到一个稍微大点的连续申请就失败——即使nvidia-smi看总量还有空。三者单独出现都不一定崩但 CI 小卡 高位缓存 偶发长样本 评估双图凑一起就概率性红。四、最小可运行复现下面用纯 PyTorchCPU 模拟 计数复现评估漏no_grad导致峰值翻倍这一核心机制不需要真实 GPU 也能看清原理真正确认显存数字请在 CUDA 上用torch.cuda.memory_allocated()打印import torch import torch.nn as nn class Tiny(nn.Module): def __init__(self, d512, layers4): super().__init__() self.stack nn.Sequential(*[nn.Linear(d, d) for _ in range(layers)]) def forward(self, x): return self.stack(x) def peak_with_grad(model, x): # 模拟评估时忘了 no_grad建了图并保留输出引用 out model(x) return out def peak_without_grad(model, x): with torch.no_grad(): out model(x) return out def demo(): model Tiny() x torch.randn(64, 512) o1 peak_with_grad(model, x) # 带图时同样的输入再前向一次相当于又一份激活在鲲鹏 o2 peak_with_grad(model, x) print(带 autograd 图时保留了, type(o1).__name__, 引用且图未释放) # 正确做法 o3 peak_without_grad(model, x) print(no_grad 下 out 不需要反向图峰值更低, o3.shape) if __name__ __main__: demo()在 CUDA 上把torch.cuda.max_memory_allocated()打在peak_with_grad和peak_without_grad之后对比会发现前者大约是后者的 1.8–2.2 倍视层数和是否保留输出引用。这就是 CI 偶发 OOM 的放大器。五、解决方案第一层评估与生成严格no_grad 不保留 GPU 引用第一层也是收益最大的一层手写评估/推理一律包torch.no_grad()并且不要长期持有 GPU 上的大张量torch.no_grad() def quick_eval(model, loader, device): total 0 correct 0 model.eval() for batch in loader: batch {k: v.to(device) for k, v in batch.items()} logits model(**batch) preds logits.argmax(dim-1) # 立刻在 GPU 上算完指标只保留标量绝不持有 logits correct (preds batch[labels]).sum().item() total preds.numel() # 不要 outs.append(logits) —— 这会把整张图和大张量留进列表 model.train() return correct / total if total else 0.0关键改动model.eval()torch.no_grad()双重保险避免 dropout 和建图指标在循环内就地约简成 Python 标量.item()列表里不再堆 GPU 张量循环结束前不持有logits让它出作用域即可被 allocator 回收虽然还在缓存池但不再占用逻辑峰值。六、解决方案第二层给generate()上长度与显存护栏第二层针对因素 B长尾样本。无论训练还是评估只要调generate都显式限制长度并清缓存torch.no_grad() def safe_generate(model, input_ids, max_new_tokens128, devicecuda): input_ids input_ids.to(device) with torch.backends.cuda.sdp_kernel(enable_flashTrue): generated model.generate( input_ids, max_new_tokensmax_new_tokens, # 硬上限防止长尾样本无限增长 pad_token_idmodel.config.pad_token_id, do_sampleFalse, ) # generate 结束后主动让出缓存降低后续训练 step 的碎片基线 if device cuda: torch.cuda.empty_cache() return generated def demo_generate(): # 伪代码用一个小模型验证 max_new_tokens 生效 import torch ids torch.randint(0, 1000, (1, 16)) capped ids # 真实场景替换为 safe_generate 返回值 assert capped.shape[1] 16 128, 生成长度应受 max_new_tokens 限制max_new_tokens把偶发长样本的天花板钉死empty_cache()在生成这种一次性大峰值之后把缓存池还一部分给驱动避免它和训练缓存叠在高位。注意empty_cache()有开销不要在每步训练里调只在生成/评估这种明显峰值之后调一次即可。七、解决方案第三层碎片治理与 CI 显存预算第三层针对因素 C碎片和 CI 整体预算开梯度检查点降低训练峰值model.gradient_checkpointing_enable()对 transformer 类模型这能砍掉大量激活显存用重算换空间对长序列尤其明显。固定输入长度 / 截断长尾在 collator 里把序列截到max_length从根上消除超长样本顶破峰值def collate(batch, max_length512): # 截断 padding 到统一长度避免形状抖动带来的碎片 ...设置 CUDA 显存上限做预算护栏CI 专用不进生产import os os.environ[PYTORCH_CUDA_ALLOC_CONF] max_split_size_mb:128max_split_size_mb限制 allocator 切分大块的上限能显著缓解碎片导致的总空闲够但连续块不够。配合expandable_segments:True较新 CUDA效果更好os.environ[PYTORCH_CUDA_ALLOC_CONF] expandable_segments:TrueCI 加显存监控在测试里打印峰值作为回归信号def test_smoke(): ... peak_gb torch.cuda.max_memory_allocated() / 1e9 print(f[mem] peak {peak_gb:.2f} GB) assert peak_gb 15.0, f显存峰值 {peak_gb:.2f}GB 接近上限需优化这样以后任何把峰值推高的改动会在还没 OOM时就被断言拦下。八、验证修复是否生效把以上三层合并后在 CI 的小卡上跑同一 commit 多次建议 10 次以上因为原本是概率性失败for i in $(seq 1 12); do python -m pytest tests/test_smoke.py -k cuda_oom || echo run $i FAILED done修复前这 12 次里平均挂 3–4 次修复后应当 12/12 通过且日志里[mem] peak稳定在 11–12GB 区间留出余量应对碎片与抖动。如果仍偶发优先检查是否有别的代码路径如可视化、额外 logger还在 GPU 上持有大张量。九、排查清单遇到CI 经常 OOM 但本地不复现按这个顺序查先确认是不是概率性 小卡专属本地大卡不挂、CI 小卡时好时坏基本锁定高位基线 偶发峰值。搜model(/logits是否在no_grad外所有评估、推理、生成路径都必须torch.no_grad()model.eval()。查是否append了 GPU 张量评估循环里别把logits/hidden_states堆进列表指标就地.item()约简。查generate是否限长max_new_tokens必须有硬上限长尾样本是偶发峰值的主要来源。查数据有没有超长样本collator 里截断到max_length统一形状减少碎片。开梯度检查点gradient_checkpointing_enable()直接砍训练峰值。设 allocator 护栏PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True,max_split_size_mb:128。加显存峰值断言把max_memory_allocated打印并设阈值作为回归护栏别等 OOM 才暴露。十、小结CI 经常 OOM和必现 OOM是两回事。必现是逻辑错误batch 真的大到放不下而偶发往往是显存基线被推到高位后再叠加一个偶发峰值越界——这个峰值可能来自漏写的no_grad、没限长的generate、或一条没截断的超长样本。因为受种子、数据顺序、GC 时机影响它才表现为时好时坏。修复分三层第一层收益最大给所有评估/生成严格no_grad且不持有 GPU 大张量直接削掉近一半峰值第二层给generate上max_new_tokens硬上限并适时empty_cache钉死长尾样本的天花板第三层用梯度检查点、统一长度、allocator 护栏和显存峰值断言从根上压低基线并防止碎片。核心心法是不要等 OOM 才处理把显存峰值变成可观测、可断言的指标这样这类概率性失败在 CI 里就再也没法偷渡上线。