★ FEATURED ARTICLE 凌晨两点训练跑到第 47 个 epochloss 曲线还挺漂亮然后进程被系统直接干掉日志最后一行是一串 OOM 报错。你回头翻 nvidia-smi 的历史记录才发现显存从刚启动时的 18G一路缓慢爬到 30G、40G最后撞上天花板。这个问题我踩过不止一次也帮别人定位过很多回它的迷惑性在于不是模型装不下而是训练过程中显存不断增大。如果一启动就爆那是配置问题五分钟能查出结论但如果是每跑几百步涨一点、跑十几个小时才炸那大概率是某处持有了一份本该释放的引用。这篇内容就是围绕训练时显存持续增长这件事展开的从现象判别、根因分类、定位手法到工程化防御把我自己的一套排查链路完整写出来适合已经能跑通训练脚本、但被长跑稳定性折磨的读者。1. 显存缓慢爬升和一次爆掉是两码事很多人一看到 OOM 就条件反射去调小 batch size这个动作在缓慢爬升场景下几乎无效因为根因根本不在峰值占用。先把问题分类能省掉一半排查时间。1.1 为什么显存曲线会呈现锯齿状上移正常的训练显存曲线应该是这样的启动时快速上升到某个水位之后在一个固定区间内上下抖动抖动幅度取决于当前 step 的输入长度、临时张量大小但不会系统性地抬高基线。如果你把nvidia-smi每隔 10 秒采一次样画出来的线应该是一条贴着水平线波动的锯齿。而泄漏型增长的曲线特征是锯齿依然存在但锯齿的下沿在稳步抬高。这个区别非常关键。锯齿下沿代表每步结束后释放完临时内存的稳定水位它抬高说明每步都有一小部分内存没有被回收。这批没被回收的内存累积速度取决于泄漏量和步频所以表现出来就是跑得越久越危险。我一般会做一件很土但很有效的事在训练循环里每 N 步打印一次torch.cuda.memory_allocated()和torch.cuda.memory_reserved()同时把前 20 步和前 200 步的数值对比。如果 allocated 的基线在 200 步内涨了超过 5%基本可以宣判是泄漏。1.2 判断是泄漏还是容量不足的三个观测点判别维度可以整理成一张表遇到问题先对号入座观测点容量不足的表现泄漏型增长的表现崩溃时机启动后不久或第一次遇到最长样本时运行数小时后且与训练时长正相关显存基线稳定不随时间变化下沿持续抬升换小 batch 后明显缓解或消失只是推迟崩溃时间曲线依旧上涨第三行是最有诊断价值的。如果你把 batch size 砍半问题从第 10 小时推迟到第 20 小时但依然会挂那说明泄漏速率没变只是池子变大了——这就是典型的泄漏。还有第四个观测点很多人会忽略换个模型结构是不是还复现。如果换成同规格但结构简单的模型就不涨了那泄漏大概率藏在你自己的模型代码或 loss 计算里如果换成任何模型都涨那要往数据管道、日志系统、训练框架封装层去找。1.3 一个有用的量化口径每步增量定性判断之后我更倾向于把它量化。做法很简单在一个稳定阶段取两个时间点# 在第 k 步和第 kN 步各记录一次 alloc_before torch.cuda.memory_allocated() # ... 训练 N 步 ... alloc_after torch.cuda.memory_allocated() per_step_growth (alloc_after - alloc_before) / N如果per_step_growth是几百 KB 到几 MB 量级稳定复现且为正那基本锁定泄漏。这个数字还有个用处估算剩余安全时间。假设剩余可用显存是 2G每步增长 200KB那大概一万步之后会挂。这能帮你判断是必须马上修还是先加个监控扛过这一轮实验。反过来如果per_step_growth在正负之间随机跳动均值接近 0那说明不是泄漏而是分配器的正常行为或者只是采样噪声——这种情况下去优化显存基本是白费力气应该去看是不是训练配置本身留的余量太少。提示memory_allocated()反映的是张量实际占用memory_reserved()反映的是分配器向驱动申请的池子大小。判断泄漏要看 allocated因为 reserved 只增不减是设计如此不是问题。2. 计算图被意外延长最常见的一类显存爬升如果让我给训练显存持续增长的根因排序计算图相关的引用泄漏绝对是第一名。这类问题的共同点是本该在backward()之后被释放的中间激活值因为某个 Python 变量还持有它的引用活到了下一个 step 甚至更久。2.1 loss累加为什么会让整张大图活到最后最经典的写法长这样total_loss 0 for step, batch in enumerate(loader): out model(batch) loss criterion(out, label) total_loss loss # 问题在这里 loss.backward() optimizer.step() print(fepoch loss: {total_loss / len(loader)})单看逻辑毫无毛病数值也对。但total_loss loss累积的不是一个 float而是一个带 grad_fn 的张量。每次加法都会生成一个新的张量节点并把它挂到原来的计算图上于是第 1 步到第 N 步的计算图被这条加法链串成了一条超长的图。这些图节点持有的激活值一个都释放不掉。修正方式几乎零成本total_loss loss.detach().item() # 或者 loss.detach().item()把它变成纯 Python float.detach()切断计算图引用。两种都行.item()更省心因为它彻底脱离了张量体系但会触发一次 GPU 到 CPU 的同步拷贝如果在意同步开销用.detach()累积成张量最后再.item()。这个坑的隐蔽之处在于小规模测试根本发现不了。跑 100 步显存涨几十兆你以为是正常波动等到跑 50 万步就是几十 G。所以我的建议是任何进入累加器的量只要它来自带梯度的张量第一反应就是 detach。2.2 记录指标时那些顺手的操作上面说的是故意累加还有一种更隐蔽的为了做日志或者可视化顺手把中间结果存起来了。history {loss: [], logits: []} for step, batch in enumerate(loader): logits model(batch) loss criterion(logits, label) loss.backward() history[logits].append(logits) # 灾难logits是带 grad_fn 的张量列表持有它它持有的整条反向路径就跟着一起活下来。一个 step 存一份几百步就把显存吃干净。即使是history[loss].append(loss)这种看起来无害的只要loss没 detach同样会累积计算图。我见过一个真实的案例某位同学为了画 loss 曲线把每一步的 loss 张量都存进了 list跑了一天之后以为是模型太大其实是这个 list 里塞了几万张计算图。正确的做法很明确任何要进入持久容器list、dict、DataFrame的张量必须先.detach().cpu()。如果只是想记数值直接存 float如果确实需要保留张量做后续分析.detach().cpu()之后它就变成了普通 CPU 张量不再持有 GPU 上的计算图。2.3 容器里塞tensor的隐蔽写法比显式 append 更隐蔽的是通过闭包、默认参数、或者数据类间接持有。class Meter: def __init__(self): self.values [] def update(self, v): self.values.append(v) # v 可能是张量调用方可能觉得我只是把 loss 传进去统计一下完全没意识到 Meter 内部原样存下了张量。这类代码在多人协作的项目里特别容易出问题因为写 Meter 的人和调用它的人对这里该传 float 还是 tensor没有明确约定。我的习惯是统计类组件的接口一律只接受 Python 原生数值在update里强制float(v)。这样即使调用方传了张量也会在边界处被转成 float把泄漏挡在组件外面。这是一个很小的防御设计但省了很多事后排查。还有一种来自namedtuple或 dataclass 的持有。如果你用 dataclass 组织每步的输出然后把这些 dataclass 实例塞进一个列表做回放那所有字段里的张量都会被保留。这类问题用静态检查工具也不容易发现只能靠约定输出对象如果是短生命周期的就别往长生命周期容器里放。2.4 自定义autograd与中间结果保留如果你写了自定义的torch.autograd.Functionforward里存到ctx上的张量会一直活到backward结束。这本身是设计如此但有两个容易踩的点。第一ctx.save_for_backward只应该保存反向真正需要的张量。有些人图省事把 forward 里所有中间结果都 save 一遍反正以后可能用得上结果一个模块就多吃几百兆。我一般会在写完 backward 之后回头检查一遍这里用到的每个 saved tensor 是不是都有使用删掉不用的通常能省 20% 到 40% 的激活显存。第二forward里创建了临时张量但只用于计算、不需要保留的要确保它没被闭包捕获。比如在 forward 里定义了一个 lambda 并挂到 ctx 上这个 lambda 会持有它引用的所有局部变量。还有一种情况是retain_graphTrue。这个参数在梯度累积、多次 backward 的场景下会用但一旦用了这张图在显式释放前都不会被回收。如果你在循环里每次都retain_graphTrue而没有配套的释放逻辑那就等同于每步都留一张图。排查时可以直接搜代码里所有retain_graph出现的位置。3. 用二分法和快照把泄漏点逼出来上面讲的是常见模式但真实项目里的泄漏往往藏在你想不到的地方。硬猜效率很低我更推荐用系统化的方式定位。3.1 先构造最小复现脚本第一步永远是把问题从大项目里剥离出来。做法是复制训练主循环砍掉所有非必要逻辑——日志、可视化、checkpoint、评估、指标上报全部注释掉只留前向 → loss → 反向 → step这四步。然后跑 500 步观察 allocated 是否还在涨。如果不涨了泄漏在被砍掉的某块逻辑里进入第 3.2 节的二分流程。如果还在涨泄漏在核心训练路径上重点看模型定义、loss 计算、优化器配置。这个剥离过程本身就有价值因为它帮你排除了大量干扰项。我见过不少案例最后定位到的都是某个第三方日志库在后台缓存了每一步的张量跟模型和 loss 完全无关。3.2 逐步砍代码的排查顺序从最小复现开始往回加功能每次只加一块加完跑 500 步看显存。这个顺序我建议按可疑程度排而不是按代码结构排指标统计与日志上报——最可疑优先验证模型内部的辅助输出注意力权重、中间特征图评估/验证路径——经常漏no_grad数据增强与预处理——尤其是用了 GPU 侧增强的优化器与调度器——检查有没有往 state 里塞额外东西分布式相关封装——如果是多卡训练整个流程大概要花一两个小时但比盲目读代码高效得多。关键是每次只加一块不要图快一次加两块否则涨了也说不清是谁的问题。3.3 内存快照和引用链查看如果二分到某个函数还是看不出问题就该上快照工具了。PyTorch 自带的torch.cuda.memory_snapshot()能给出当前所有已分配块的详细信息包括大小和分配时的调用栈import torch, pickle snap torch.cuda.memory_snapshot() # 找出占用最大的若干块 big sorted(snap, keylambda b: b.get(allocated_size, 0), reverseTrue)[:10] for b in big: print(b.get(allocated_size), b.get(frames, [])[:5])frames里就是 Python 层的调用栈能直接告诉你这块内存是在哪一行申请、被谁持有的。我靠这个手段定位过一次某数据处理库内部缓存了整批样本的问题读代码是绝对看不出来的。另一个思路是看 Python 侧的引用。用gc.get_referrers(obj)可以找出谁在引用某个对象import gc, torch # 假设你怀疑某个张量被意外持有 t model.some_module.weight for r in gc.get_referrers(t): print(type(r), str(r)[:200])这个方法对定位张量被谁抱着不放特别有效。缺点是输出可能很长需要有点耐心过滤。3.4 排查清单表把上面这些经验浓缩成一张对照表遇到问题时逐条过一遍可疑点典型症状快速验证方式loss 未 detach 累加每步稳定涨几十 KB 到几 MB改成.item()重跑 500 步列表存了带梯度张量增长量与日志频率正相关注释掉所有 append 语句验证路径缺 no_grad每个验证周期后跳升一截在 eval 循环加装饰器retain_graph每次 backward 后不降搜代码关键字DataLoader worker 堆积随 epoch 数递增设num_workers0对比日志库后台缓存与训练逻辑无关的稳定增长换成 print 对比这张表覆盖了我遇到的八成以上案例剩下两成通常是框架版本相关的分配器行为放到第 5 节讲。4. 框架细节里的隐性持有者核心训练路径干净之后剩下的泄漏往往来自框架的边角料。这些地方平时根本不会去看出问题时又特别难怀疑到它们头上。4.1 DataLoader的worker与pin_memory数据管道引起的显存增长有两种机制。第一种是pin_memoryTrue配合num_workers0时pin memory 线程会维护一个预取的缓冲区。正常情况下这个缓冲区大小是有限的但如果主进程消费速度慢于生产速度缓冲区会堆积。判断方法很简单把num_workers设为 0、pin_memory设为 False跑 500 步对比。第二种更隐蔽自定义的Dataset里如果持有 GPU 张量或者跨 worker 共享的对象每个 worker 都会有自己的一份而且生命周期和 worker 绑定。如果你用了persistent_workersTrue这些对象会一直活着。我之前遇到过一次某份预计算的 embedding 表被当作 Dataset 的属性四个 worker 各存一份凭空多吃了好几 G。正确的设计是把 Dataset 做成无状态或轻量状态大的数据要么放共享内存要么在 worker 内部按需加载。数据增强里的随机数生成器也建议用 worker 局部的避免跨 worker 共享导致的隐性同步和状态堆积。4.2 hooks、回调、日志工具持有引用register_forward_hook和register_backward_hook是调试利器但 hook 函数如果捕获了张量并保存起来就会一直持有计算图。这类代码经常是临时加的调试完忘了删然后这个临时代码在线上跑了几周。排查方式是搜代码里所有register_*_hook逐个确认 hook 内部有没有往外部容器写东西。如果有改成只记录标量或者只做即时统计。日志和可视化工具的坑更隐蔽。某些实验追踪工具会缓存最近 N 步的完整输出用于可视化这个缓存如果按张量存而不是按数值存就会保留计算图。我的一般做法是任何上报接口的入参先用.detach()或float()处理一遍在上报层做一次强制清洗比在调用点一个个改可靠。回调机制也有类似问题。如果你实现了一个每 N 步执行一次的回调回调函数里捕获了当前 step 的输出张量而这个回调对象被训练器长期持有那就形成了持有链。设计回调时接口参数最好只传数值和元信息不传张量。4.3 验证与推理路径漏了no_grad这是最经典的新手坑但老手也会犯尤其是在验证逻辑比较复杂的项目里。model.eval() with torch.no_grad(): # 外层加了 for batch in val_loader: with autocast(): # 但内部某处又开了梯度 pred model(batch) metric.update(pred) # pred 如果没 detachmetric 又持有它外层no_grad确实能防止计算图构建但如果pred被塞进了会被长期持有的容器内存依然会增长——只是这次存的是不带梯度的张量增长速度慢一些症状更隐蔽。另一个常见情形是混用了torch.inference_mode()和no_grad。前者更激进跑出来的张量不能参与后续自动求导但如果验证完把结果和训练结果拼在一起做某些操作可能会报错或者产生意料之外的保留。我的经验是验证阶段统一用no_grad输出立即转成 CPU 或数值别留任何张量引用。至于评估指标的计算如果用的是第三方库务必确认它在update里有没有做 detach。有些库的update会保留 logits 用于后续计算更复杂的指标这时要看清它的文档必要时手动传.detach()之后的张量。4.4 优化器state、EMA、梯度累积优化器的 state 是显存里一块固定的开销正常情况下第一步之后就稳定了。但如果你的参数量在训练过程中发生变化比如动态添加模块、从 checkpoint 恢复时结构不一致state 也可能跟着变。这种情况少见但排查时要确认一下len(optimizer.state)在不同 step 是否稳定。EMA指数移动平均是另一个容易出问题的地方。实现 EMA 时如果直接对参数做深拷贝参数有多少份显存就多占一份。更麻烦的是如果拷贝发生在计算图还没释放的时候复制出来的是带 grad_fn 的张量那泄漏规模会成倍放大。正确写法是在no_grad下做更新并且用detach()或者直接操作.datatorch.no_grad() def update_ema(ema_model, model, decay): for e, p in zip(ema_model.parameters(), model.parameters()): e.mul_(decay).add_(p.detach(), alpha1 - decay)梯度累积本身不增加显存但有个配套问题很多人为了梯度累积会把 loss 乘上累积步数如果不 detach乘出来的中间张量也可能被持有。另外梯度累积配合retain_graph使用时要注意每轮循环的最后一次 backward 必须不 retain否则图会一直留着。最后提一下混合精度。GradScaler的内部状态是固定开销不会增长。但如果你的 loss 在 inf/nan 检查里被存进了一个异常样本集那这些样本的 loss 张量会一直被保留。异常样本集最好只存样本索引别存张量。5. 缓存分配器、碎片与empty_cache的真实作用讲完引用泄漏还得讲另一半真相有时候显存增长不是 bug而是分配器的正常行为被误读了。这两类问题混在一起会导致排查方向完全跑偏。5.1 reserved和allocated的区别PyTorch 的 CUDA 分配器是分层的张量申请的显存先记在已分配allocated分配器为了减少系统调用会一次性向驱动申请更大的块记为已保留reserved。reserved 只增不减是设计如此因为释放回驱动再申请回来代价很高。所以看到 reserved 一路上涨不用慌真正要看的是 allocated。判断标准很明确allocated 稳定 → 没问题allocated 持续上涨 → 有泄漏。至于两者之间的差值那是可以复用的空闲块属于暂时用不上但留着以后用的部分。我见过有人因为 reserved 涨到接近上限就开始优化结果折腾半天发现 allocated 一直很稳纯属白费功夫。所以排查第一步永远是分清这两个指标。5.2 碎片化的判定与缓解碎片化是指虽然总空闲量够但都是小块凑不出一块连续的大空间于是申请大张量时失败。它的症状很有意思报 OOM但 nvidia-smi 显示还有不少空闲。判断方法是看memory_reserved()和memory_allocated()的差值以及分配器统计里的大块分布。如果差值很大比如 reserved 40Gallocated 只有 15G而且持续存在那就是严重的碎片化或者大块闲置。缓解手段有几个层次。最轻的是把 batch size 调成更整的数或者把输入尺寸对齐到 8 的倍数减少奇形怪状的分配。中等的是调整PYTORCH_CUDA_ALLOC_CONF里的max_split_size_mb限制分配器把大块切得太碎。STFW 之前也被这个参数坑过——设得太小反而会让大张量申请失败一般从 512 开始试比较稳。更深层的解法是改架构把变长输入按长度分桶让同一批内的序列长度接近避免一个大张量配一堆小张量的组合。这个改动对扩散模型、序列模型的效果都很明显。5.3 empty_cache什么时候有用什么时候白费torch.cuda.empty_cache()的作用是把 reserved 里空闲的块释放回驱动它不会释放任何被张量引用的内存。理解这一点非常关键否则会出现我每步都调了 empty_cache 为什么还涨的困惑。所以它的适用场景很有限训练中进行一次性的大操作比如保存完整模型、做一次全量评估之前需要腾出连续空间——有用训练中每个 step 都调——没用而且有害因为每次释放都要重新申请反而拖慢速度存在真实引用泄漏时调——完全没用泄漏的张量还被引用着empty_cache 碰不到我的一般建议是训练循环里永远不要调 empty_cache。只在阶段切换点训练转评估、评估转保存调一次这样既能清理碎片又不影响吞吐。如果发现不调就 OOM那说明有真实泄漏应该去修泄漏而不是靠 empty_cache 硬撑——这种缓解会让问题越滚越大最后连缓解的空间都没有。注意empty_cache有同步开销会让 GPU 空转等待。在分布式训练里频繁调用还可能引发所有卡之间的等待拖慢整体进度。6. 长跑训练里的工程化防御单次修好一个问题只是治标。真正省心的是建立一套机制让这类问题在早期就暴露、在崩溃前能自救。6.1 显存监控埋点我会在训练循环里固定加一小段监控代码成本极低但收益很高if step % 100 0: alloc torch.cuda.memory_allocated() / 1024**3 reserved torch.cuda.memory_reserved() / 1024**3 logger.info(fstep{step} alloc{alloc:.2f}G reserved{reserved:.2f}G) # 与历史基线比较异常时告警 if alloc baseline_alloc * 1.3 and step 500: logger.warning(allocated 较基线增长超过 30%疑似泄漏)关键点是记录一个启动后稳定期的基线值比如第 200 步的 allocated之后持续对比。这样不管是哪种机制的泄漏只要表现为 allocated 上涨都能被抓住。我一般把阈值设在 20% 到 30%低于这个数的波动通常是正常的批次差异。日志格式上要注意别把整个张量打出来只打标量。见过有人写logger.info(fout{out})调试日志文件几秒钟涨到几个 G同时这个格式化字符串本身就持有张量引用。6.2 结构上的约束代码结构层面有几条硬约束能大幅降低泄漏概率。第一训练循环里不允许出现对张量的 append 或赋值到外部变量。指标统计一律走传入即转数值的通道。我把这条写进过团队的代码规范效果立竿见影。第二模型 forward 的返回值只用一次立刻消耗。不要在 forward 之外再持有输出做稍后处理因为稍后往往意味着跨 step。第三所有跨 step 存活的对象必须显式声明生命周期。比如 EMA 模型、teacher 模型、缓存的特征库这些是故意要长期存活的那就明确写出来并注释原因其余的都不应该跨 step 存活。第四验证和推理路径统一用装饰器强制no_grad而不是靠每个函数自己记得加。写一个torch.no_grad()包在外层比在每个 eval 函数里加可靠得多。6.3 checkpoint与异常恢复既然无法百分百保证不 OOM那就得让崩溃的代价可控。checkpoint 的策略有几个细节值得说。保存频率上我一般按时间而不是按步数比如每 30 分钟存一次。按步数存的话样本长度变化大的任务里不同步之间的耗时差异可能有十几倍容易出现存得太频繁或者存得太稀疏。保存内容上除了模型参数和优化器 state随机数状态一定要存。不然恢复之后数据顺序、增强结果全变了实验无法复现。CUDA 的随机状态和 Python 的随机状态都要存。保存时机上建议在 epoch 边界或者验证之后存而不是在训练 step 中间存避免打断计算流的连续性。同时保存动作最好在empty_cache之后做这样大 checkpoint 的序列化过程不会因为碎片而失败。恢复逻辑上要处理模型结构变了但 checkpoint 没变的情况。我习惯在 checkpoint 里存一个版本号加载时校验版本不匹配就走参数名匹配的宽松加载跳过不匹配的层并打印警告。这比直接报错或者静默加载失败要好排查得多。6.4 我的几条硬性习惯最后分享几条我在长期实践里固化下来的习惯都是踩过坑之后加上的。每次改完训练代码先跑 1000 步看显存曲线。不用等长跑一千步足够暴露大部分稳定泄漏。这个动作现在是我提交前的固定流程成本三分钟。监控指标里永远包含 allocated 的绝对值而不只是使用率。使用率会掩盖很多问题因为分母是总显存而绝对值的趋势更敏感。遇到怀疑是泄漏的情况第一步先关掉所有日志和可视化重跑一次。这一条帮我解决过至少五次问题大部分神秘泄漏都出在日志链路上。不要迷信torch.cuda.empty_cache。如果它能让问题消失说明是碎片问题如果不能让问题消失说明是引用泄漏。两种情况都不该靠它硬撑。多卡训练时显存监控要打到每张卡上。只打 rank 0 会漏掉其他卡上的不均衡。有时候不是整体泄漏而是某张卡因为数据分布不均导致多吃了内存这种情况在 rank 0 上完全看不出来。给长时间任务加一个自动降级路径。比如检测到 allocated 超过阈值的 85%就自动保存 checkpoint 并触发优雅退出而不是等着系统 OOM 杀掉进程——后者可能来不及保存损失半天进度。这套东西配齐之后训练长跑基本就不会再有跑到一半突然挂掉不知道为啥的情况了即使出问题日志里也有足够信息让你在十分钟内定位到大致范围。
阅读完成 · 觉得有帮助？
