PyTorch 显存泄漏排查与修复：从定位到根治的完整指南

训练跑着跑着，显存一路往上飙，三个小时后直接 OOM，连验证都跑不动。这个问题，做 PyTorch 训练的基本都撞上过。网上搜"显存泄漏"会跳出各种答案，但很多文章要么只讲概念不给排查路径，要么给一堆杀招看完还是不知道怎么定位到具体是哪一行代码出的问题。我前几年从单卡训练切到多卡，又切到长期跑实验，被这个事折磨了挺久，后来逐步把定位和修复的方法沉淀成了一套固定流程。这篇就把这套流程完整写出来，从判断是不是真泄漏、怎么用工具定位源头，到具体修哪些代码，再到几个真实案例复盘，希望能帮你少走点弯路。

如果你是刚开始用 PyTorch 训练模型、老板催着出实验结果的算法工程师，或者正在把训练任务从短迭代切成长跑任务的开发者，这篇东西都会有用。它解决的是"跑得越久越容易爆显存"这类非常具体的问题，不涉及分布式训练调参那一套，聚焦在显存这块。

1. 显存泄漏是怎么回事：先搞清楚"越跑越高"的真相

1.1 显存的三种"虚高"：缓存、碎片和真泄漏

很多人在这一步就误解了。先用 nvidia-smi 看到显存占用一直在涨，第一反应就是"完了，代码泄漏了"。但事实是，大部分情况下你看到的"持续上涨"根本不是真正的泄漏，而是 PyTorch 的显存管理机制在起作用。

PyTorch 为了加速训练，默认启用了一个叫 CUDA caching allocator 的机制。它跟操作系统管理内存的思路类似：申请一块显存后，用完不马上还给驱动，而是放进自己的缓存池里，下个 step 需要时直接复用。这样避免频繁调用 cudaMalloc，省掉了和驱动交互的开销。代价就是，从外部工具看，显存占用率会维持在一个高位，哪怕当前 step 真正用到的张量只有这个峰值的一半。这就像饭店备菜，午高峰过后备好的食材不会全部倒掉，而是留到晚上继续用。你在外面看到冰箱里塞满了菜，不代表今天的菜卖不完，只是存储策略的问题。

第二种是碎片化。训练过程会不断创建不同大小的张量，缓存池里的空闲块可能出现"大块不够用、小块用不上"的局面，这时候 allocator 会再向驱动申请新显存，整体占用就慢慢往上翘。这也是为什么有人发现，同一个模型同一个 batch size，训练到后期显存比刚开始高不少，但重启进程后又能正常跑。碎片不是泄漏，但表现得像泄漏。

真正的泄漏是第三种：某些张量或计算图对象被代码意外持有引用，导致它们的生命周期被无限延长。每一步迭代都会新建一批这样的对象，而旧的又没法被 Python 的垃圾回收机制释放，显存占用自然一条直线爬升。这就是需要重点排查的部分。所以定位的第一步，不是急着改代码，而是先区分你撞上的是缓存、碎片还是真泄漏。

1.2 训练循环里最常见的几个"引用坑"

结合我五年多的 PyTorch 排查经验，训练循环中出现真泄漏，9 成是引用持有问题。最典型的就是在训练循环里持续往某个 list 里塞东西。比如：

all_losses = [] for batch in train_loader: loss = model(batch) loss.backward() optimizer.step() all_losses.append(loss) # 罪魁祸首

这个loss在backward()之后其实已经不打算用了，按理说对应的计算图应该被释放。但你把它 append 到 list 里，它就带着整个反向图结构一起活了下来。每个 step 都会留下一个"尸体"，显存不涨才怪。

正确的做法是all_losses.append(loss.item())，因为.item()返回的是 Python 浮点数，和 GPU 上的张量彻底解绑。很多人知道这个，但实际写代码时图省事，顺手就 append 了带梯度的 loss，踩坑之后才明白。

还有一个隐蔽的坑，是保存模型的中间输出。比如为了做注意力可视化，把每层特征图都存到 dict 里，这个 dict 又在 step 循环外定义。注意力模块的输出一旦被持有多出几个 step，显存立刻翻倍。类似的还有在验证阶段把 logits 或 preds 保存下来用于后续计算指标，但忘了 detach 和转成 CPU。验证阶段虽然不求梯度，但 CUDA 张量只要还在 GPU 上且被引用，显存就不会释放。

我踩过的另一个坑是自定义 Dataset 或 DataLoader 的 worker 里，在__getitem__中不小心执行了某个 GPU 操作，把张量留在 GPU 上，再通过 collate 压回主进程。这种问题在单卡低 epoch 时根本看不出异常，一跑长训就现原形。

要区分"看着高"和"真的泄漏"，最直接的办法是把torch.cuda.memory_allocated()和torch.cuda.memory_reserved()打出来看。前者是当前真正使用的张量显存，后者是 PyTorch 预留的缓存池大小。如果 allocated 稳定在某个区间、reserved 偶尔高，那是缓存机制，不用管；如果 allocated 每个 step 都净增一大截，那就是真泄漏，直接进入定位环节。

2. 显存问题定位：先判断真假泄漏，再锁定源头

2.1 三个信号帮你判断"是不是真泄漏"

不用一上来就上 profiler，先用最便宜的三个信号做判断。

第一个信号：跑一会儿后重启进程再看，显存占用是否复位。如果重启后能回到初始占用，说明之前的高占用是进程在生命周期内人为持留的，大概率就是真泄漏。如果重启后依然高，那可能根本不是你的 PyTorch 代码问题，而是驱动、CUDA context，或者其他进程占用的显存。第二步要查是不是真的在持续增长。我建议把 monitoring 代码插在训练脚本里，每个 step 或每隔 N 个 step 打印torch.cuda.memory_allocated()，画一条时间序列曲线。真正的泄漏曲线特征是"严格单调爬升"，不回落。而缓存和碎片导致的上涨，会在显存不够时触发释放，出现一种锯齿状回落。看到锯齿，说明不是泄漏，只是 allocator 在极限腾挪。

第三个信号最实用：把训练代码里 optimizer 相关的部分不动，单独跑一个只有 forward 和 loss 计算、不调 backward 的版本，观察显存是否还涨。如果 forward 不回传梯度时显存稳定，那就说明问题出在反向传播或梯度保持环节；如果 forward 阶段都涨，那就要查模型本身的前向过程里是否有东西在累积，比如某些模块不断往 buffer 里写数据。

这三个信号配合起来，能快速把问题范围从"整个训练脚本"缩小到"前向"或"反向+优化器"这两个区域。很多人在这个问题上花了一整天，其实用这三个信号半小时就能锁定大概方向。

2.2 PyTorch 自带工具：memory_summary 和 memory_snapshot

在明确"真泄漏"之后，就该用 PyTorch 自带的显存工具做精细定位了。这里有两个我最常用的工具。

第一个是torch.cuda.memory_summary()，直接打印当前所有显存分配块的详细汇总。使用方法很简单：

print(torch.cuda.memory_summary())

它会输出类似这样的信息：每个分配段的地址范围、大小、是否 active、所属 stream、分配时的调用栈。注意，这个 summary 默认只显示当前进程的分配状态，而且因为缓存池的原因，可能显示很多已经释放但还在池里的块，所以别被大量"segment"吓到，重点看 active 的部分，那些才是真正存活着的张量。

如果你的显存分配块数量非常多，逐行看 summary 不现实。这时用torch.cuda.memory_snapshot()，它能拿到所有分配块的元数据，返回一个 dict 列表。把这些数据转成 DataFrame，按 size 排序，就能快速找出哪个"巨型张量"一直没释放：

import torch from pandas import DataFrame snap = torch.cuda.memory_snapshot() df = DataFrame(snap) df = df.sort_values("size", ascending=False) print(df[["size", "stream", "device_ptr", "active"]].head(20))

这个脚本我每次排查显存问题时都会跑。它解决了一个核心痛点：当显存条目多达上千个时，人眼无法从中挑出问题张量，但排序之后，最大的那些 active 块往往就是泄漏元凶。比如你看到一个 512MB 的块反复出现、一直 active，再结合分配时的调用栈（frames字段可以打印），基本就能定位到是哪一行代码创建了它。

不过有个细节要注意：memory_snapshot()拿到的调用栈需要torch.cuda.memory._snapshot()配合开启record_context才能看到。我通常在排查时加这样一行开启上下文记录：

torch.cuda.memory._record_memory_history(True)

在跑几十个 step 之后，调用torch.cuda.memory._dump_snapshot()把文件 dump 下来，用 Chrome tracing 打开，就能看到每个张量的完整生命周期。这是目前 PyTorch 官方推荐的方式，虽然有一定性能开销，但排查泄漏时这点开销完全值得。

2.3 用弱引用和垃圾回收辅助定位"持有者"

有时候快照也能看到某个大张量一直存活，但调用栈只显示到某个框架内部函数，看不出是你业务代码的哪一行持有它。这种情况，就要靠 Python 层的工具了。

Python 的gc模块配合weakref可以找出某个张量到底被谁引用着。大致的思路是：通过gc.get_objects()遍历 Python 堆上的所有对象，筛选出torch.Tensor类型，记录它们是否在 GPU 上、大小多少，然后结合显存快照，找出哪些张量是在预期中没有被释放的，再用gc.get_referrers()反向查出是谁引用了它。

import gc import torch objs = gc.get_objects() for obj in objs: if torch.is_tensor(obj) and obj.is_cuda: size = obj.numel() * obj.element_size() refs = gc.get_referrers(obj) print(size, obj.shape, len(refs), refs[:3])

注意，这个脚本不要直接跑全量，因为 Python 进程里的临时对象太多，会刷屏到崩溃。我的做法是配合 memory snapshot 先找到可疑的大块地址，再用类似方法在特定时机（比如 step 结束时）抓取所有 torch.Tensor，然后对比哪些是错误的续命者。

还有一个小技巧：在解释这些张量为何被持有之前，先检查是不是有类实例把张量存到了self上。比如自定义Trainer类，训练循环里反复赋值self.current_batch = batch，而batch每次都是新张量，旧的如果没被覆盖（或者覆盖频率跟不上），就会大量堆积。这类问题用 objgraph 之类的工具也能看，但我觉得最简单的是在可疑类里手动打印__dict__里的 Tensor 条目，一眼就能看到有没有不断增长的 list 或 dict。

3. 修复实操：把泄漏的代码改到不泄漏

3.1 梯度累积、optimizer.step 和 zero_grad 的正确姿势

先说梯度累积和zero_grad()这一对。很多人只知道optimizer.zero_grad()会清空梯度，但没意识到它有一个更省显存的选择：optimizer.zero_grad(set_to_none=True)。

set_to_none=True会把梯度的.grad属性直接置为 None，而不是调用.zero_()把张量清零。前者会真正释放掉梯度张量占用的显存，后者只是把数值清零但空间还留着。在长训练场景下，只要每个 step 都执行set_to_none=True，梯度累积表剧无法形成"梯度缓存池膨胀"。

optimizer.zero_grad(set_to_none=True)

我实测过，在 batch size 较大的模型上，这个参数的改动能让显存下降 1-2GB，而且完全不影响训练结果。官方文档里也推荐两边做法，但很多人还是习惯写成默认的zero_grad()。

接着说梯度累积。当你用梯度累积模拟大 batch 时，常见的错法是累积多个 step 的 loss 后才 step 一次，但如果中间某一步的 loss 没有 detach，计算图就会把多个 step 的反向路径串在一起。比如：

for i, batch in enumerate(loader): loss = model(batch) loss = loss / accum_steps loss.backward() if (i + 1) % accum_steps == 0: optimizer.step() optimizer.zero_grad(set_to_none=True)

这个写法没问题，因为loss.backward()反向之后，计算图默认会释放（保存在 loss 里的中间变量会被回收）。但如果你在backward()之前把 loss 存起来（比如self.last_loss = loss），那么累积到下一步时，上一步的计算图还没完全释放，显存占用就会变成累积倍数。我自己写累积逻辑时，坚持一个原则：除了缩放系数之外，不要在 backward 之后长期保存任何与计算图有关联的对象。

3.2 前向过程中的显存杀手与规避方式

前向阶段最容易造显存问题的，除了前面说的保存中间结果，还有一个是惯性调用了.detach()而忘了.clone()。举个例子，你想把某一层的输出拿去计算额外的 loss，为了不把传播路径串进主计算图，你可能写：

feat = layer(x).detach()

然后把这个feat存起来用于下一轮更新。这个 detach 只是切断了梯度关系，并没有阻止feat本身占用的显存。如果你存了很多份这样的feat，显存照样爆。正确的做法是.detach().clone()，或者直接把它移到 CPU 保存，让 GPU 侧的空间能被回收。

再一个前向杀手是 attention 的实现细节。很多人在自定义 transformer 模块时，把attention_weights存成类属性：

self.attn_weights = attn_weights

这是为了可视化需求，但它带来的后果是，每个 step 的注意力矩阵都被保存在类实例上，直到下一次覆盖才被释放。如果序列长度是 1000，head 数是 8，只要 batch 数一多，这一步就是显存大头。我的建议是，可视化功能不要直接改网络结构，而是单独写一个 hook 或者 forward 后临时获取，用完立即删除。

还有一类隐蔽问题：模型内部的 buffer。如果你在模块里写了register_buffer('running_feat', torch.zeros(...))，然后又不断用新 tensor 去赋值给它，虽然 buffer 数量不变，但每次赋值的 tensor 如果是从 CUDA 上下文计算出来的，旧的会被释放，新的会占新空间。如果赋值频率和显存分配策略不匹配，也会造成缓存池继续膨胀。排查时可以用model.state_dict()里 buffer 的数量和大小做对照，看看是不是某个自定义 buffer 被不当更新。

3.3 验证阶段与 DataLoader 的显存管理

很多显存泄漏，其实发生在验证阶段而不是训练阶段。原因是验证代码里为了计算各类指标或画图，常常把每个 batch 的预测结果完整地收集起来。比如：

for batch in val_loader: preds = model(batch) all_preds.append(preds) # 这句致命

preds是在 GPU 上的张量，append 到 list 之后，验证持续多久，显存就持有多久。修复方式有两个方向：一是及时转 CPU 再把 GPU 侧释放掉：

all_preds.append(preds.detach().cpu())

二是不收集所有 batch，改成在循环内计算指标、只保留标量。我倾向于推荐第二种，因为即使转移到 CPU，如果总数太大也会把内存打爆，何况还得做一次显存->内存的数据搬运。

DataLoader 这边也有说法。num_workers增加以后，主进程和 worker 进程之间的数据传递用的是共享内存，正常情况下不会重复占用显存。但如果你在自定义 collate 函数里执行了一些 GPU 操作，比如把 CPU 张量.cuda()后拼接，那么 collate 阶段的临时张量会落在 GPU 上，导致每个 step 都有额外显存开销，而且 worker 之间还可能重复分配。排查方法是把 collate 里的.cuda()去掉，把设备转移统一放到训练循环内。很多人忽略这个细节，其实是显存持续上涨的隐性来源。

4. 常见问题与排查技巧实录

4.1 三个真实案例复盘

案例一：某次跑 ViT 微调，显存以大概每 200 步 100MB 的速度稳定上升。我先用memory_reserved和allocated对照，发现 allocated 持续增长，说明是真泄漏。然后用 memory snapshot 看到一堆 224×224 的中间特征图一直存活。往下追调用栈，发现是train_logits列表在累积，目的是为了画准确率曲线，提前把每个 step 的 output 都存了下来。改成在循环内计算 acc 并只存标量后，显存曲线平滑如初。这个案例就是典型的"可视化代码污染训练代码"。

案例二：一个用混合精度训练的项目，torch.cuda.amp.GradScaler正常用着，但每到 1000 步左右，梯度会突然变成 inf，然后 loss 变成 NaN。我一开始以为是学习率问题，调了半天发现是梯度累积逻辑里scaler.scale(loss)之后，多步累积的梯度没有unscale_，导致 scale 因子在累积过程中不断叠加放大，数值溢出。虽然这不是直接的显存泄漏，但它会导致后续 tensors 变大、计算图保存异常，显存也会随之升高。修复方式是把unscale_和step放在累积完成的那个分支里。这个小教训说明，某些显存问题背后不是空间管理，而是数值稳定性引发的连锁反应。

案例三：多卡 DDP 训练，从 2 张卡扩展到 4 张卡后，单个 step 显存翻倍。查了很久，后来发现是模型中有一个很大的self.register_buffer('pos_emb', pos_emb)，但在每次 forward 时对pos_emb做了一次expand操作。在 DDP 下expand会创建张量副本并接触到 allreduce 通信，导致每个 rank 上都多了一份完整 embedding。解决方式是把pos_emb改为非 buffer 的普通张量，并在 forward 内部直接索引，不预展开。这个案例值得提醒：多卡场景下的显存泄漏经常不是"泄漏"，而是"重复分配"，排查思路要增加通信和 buffer 视角。

4.2 显存问题排查速查表

我还要提醒一个容易被忽略的点：torch.cuda.empty_cache()能回收缓存池里空闲的显存，但它不会释放正在被引用的张量，所以它对"真泄漏"几乎无效。很多人误以为训练到一半手动调一下空缓存就能解决问题，最多只是把碎片整理掉一些，后面还是会涨。正确的用法是在验证之前、或者即将结束训练前做一次显存整理，让峰值后遗留在池子里的大块空闲缓存还给驱动，给接下来的操作腾出空间。

提示：如果你用 jupyter notebook 或类似交互环境做实验，注意它可能为你隐式保存上一次 cell 输出的 CUDA 张量，也会造成显存一直高。这是环境层面的问题，和训练脚本无关，但很多人排查半天之后才发现是 notebook 缓存了中间输出。

根据我的经验，显存泄漏这件事，大部分时候不是玄学，而是"你不经意间保存了不该保存的东西"这么简单。严格的代码审查习惯、每个 step 至少打印一次memory_allocated的监测脚本、再加上几次 memory snapshot 的排查，基本能把问题控制住。我在实际项目中养成了一个习惯：每次改动训练脚本后，都先跑 200 步看显存曲线是否平稳，再继续长训。这个习惯帮我拦下了好几次重大事故，非常奏效。

最后再分享一个实用的小技巧：在训练脚本里周期性地记录torch.cuda.max_memory_allocated()和当前 step 的torch.cuda.memory_allocated()，把差值打印出来。这个差值如果不断增大，说明你"曾经用到的高水位"和"当前活跃量"之间的差距在扩大，其实反映了某个周期性的操作在创建大张量后没及时释放。观察几次，你就能总结出它是跟forward、backward还是eval绑定出现的，然后顺藤摸瓜找到元凶。祝各位实验跑得稳，不用再被 OOM 标红搞得心力交瘁。