PyTorch训练显存缓慢爬升：OOM根因定位与长跑防御

凌晨两点，训练跑到第 47 个 epoch，loss 曲线还挺漂亮，然后进程被系统直接干掉，日志最后一行是一串 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 size 砍半，问题从第 10 小时推迟到第 20 小时，但依然会挂，那说明泄漏速率没变，只是池子变大了——这就是典型的泄漏。

还有第四个观测点很多人会忽略：换个模型结构是不是还复现。如果换成同规格但结构简单的模型就不涨了，那泄漏大概率藏在你自己的模型代码或 loss 计算里；如果换成任何模型都涨，那要往数据管道、日志系统、训练框架封装层去找。

1.3 一个有用的量化口径：每步增量

定性判断之后，我更倾向于把它量化。做法很简单，在一个稳定阶段取两个时间点：

# 在第 k 步和第 k+N 步各记录一次 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(f"epoch 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.Function，forward里存到ctx上的张量会一直活到backward结束。这本身是设计如此，但有两个容易踩的点。

第一，ctx.save_for_backward只应该保存反向真正需要的张量。有些人图省事，把 forward 里所有中间结果都 save 一遍，反正"以后可能用得上"，结果一个模块就多吃几百兆。我一般会在写完 backward 之后回头检查一遍：这里用到的每个 saved tensor 是不是都有使用？删掉不用的，通常能省 20% 到 40% 的激活显存。

第二，forward里创建了临时张量但只用于计算、不需要保留的，要确保它没被闭包捕获。比如在 forward 里定义了一个 lambda 并挂到 ctx 上，这个 lambda 会持有它引用的所有局部变量。

还有一种情况是retain_graph=True。这个参数在梯度累积、多次 backward 的场景下会用，但一旦用了，这张图在显式释放前都不会被回收。如果你在循环里每次都retain_graph=True而没有配套的释放逻辑，那就等同于每步都留一张图。排查时可以直接搜代码里所有retain_graph出现的位置。

3. 用二分法和快照把泄漏点逼出来

上面讲的是常见模式，但真实项目里的泄漏往往藏在你想不到的地方。硬猜效率很低，我更推荐用系统化的方式定位。

3.1 先构造最小复现脚本

第一步永远是：把问题从大项目里剥离出来。做法是复制训练主循环，砍掉所有非必要逻辑——日志、可视化、checkpoint、评估、指标上报全部注释掉，只留前向 → loss → 反向 → step这四步。

然后跑 500 步，观察 allocated 是否还在涨。

• 如果不涨了：泄漏在被砍掉的某块逻辑里，进入第 3.2 节的二分流程。
• 如果还在涨：泄漏在核心训练路径上，重点看模型定义、loss 计算、优化器配置。

这个剥离过程本身就有价值，因为它帮你排除了大量干扰项。我见过不少案例，最后定位到的都是"某个第三方日志库在后台缓存了每一步的张量"，跟模型和 loss 完全无关。

3.2 逐步砍代码的排查顺序

从最小复现开始往回加功能，每次只加一块，加完跑 500 步看显存。这个顺序我建议按"可疑程度"排，而不是按代码结构排：

1. 指标统计与日志上报——最可疑，优先验证
2. 模型内部的辅助输出（注意力权重、中间特征图）
3. 评估/验证路径——经常漏no_grad
4. 数据增强与预处理——尤其是用了 GPU 侧增强的
5. 优化器与调度器——检查有没有往 state 里塞额外东西
6. 分布式相关封装——如果是多卡训练

整个流程大概要花一两个小时，但比盲目读代码高效得多。关键是每次只加一块，不要图快一次加两块，否则涨了也说不清是谁的问题。

3.3 内存快照和引用链查看

如果二分到某个函数还是看不出问题，就该上快照工具了。PyTorch 自带的torch.cuda.memory_snapshot()能给出当前所有已分配块的详细信息，包括大小和分配时的调用栈：

import torch, pickle snap = torch.cuda.memory_snapshot() # 找出占用最大的若干块 big = sorted(snap, key=lambda b: b.get("allocated_size", 0), reverse=True)[: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 排查清单表

把上面这些经验浓缩成一张对照表，遇到问题时逐条过一遍：

这张表覆盖了我遇到的八成以上案例，剩下两成通常是框架版本相关的分配器行为，放到第 5 节讲。

4. 框架细节里的隐性持有者

核心训练路径干净之后，剩下的泄漏往往来自框架的边角料。这些地方平时根本不会去看，出问题时又特别难怀疑到它们头上。

4.1 DataLoader的worker与pin_memory

数据管道引起的显存增长有两种机制。第一种是pin_memory=True配合num_workers>0时，pin memory 线程会维护一个预取的缓冲区。正常情况下这个缓冲区大小是有限的，但如果主进程消费速度慢于生产速度，缓冲区会堆积。判断方法很简单：把num_workers设为 0、pin_memory设为 False，跑 500 步对比。

第二种更隐蔽：自定义的Dataset里如果持有 GPU 张量或者跨 worker 共享的对象，每个 worker 都会有自己的一份，而且生命周期和 worker 绑定。如果你用了persistent_workers=True，这些对象会一直活着。我之前遇到过一次，某份预计算的 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 如果没 detach，metric 又持有它

外层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()或者直接操作.data：

@torch.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(), alpha=1 - 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 40G，allocated 只有 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(f"step={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(f"out={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 杀掉进程——后者可能来不及保存，损失半天进度。

这套东西配齐之后，训练长跑基本就不会再有"跑到一半突然挂掉不知道为啥"的情况了，即使出问题，日志里也有足够信息让你在十分钟内定位到大致范围。