1. 显存泄漏是怎么回事先搞懂“越跑越高”的本质先聊个现象很多人在本地跑 PyTorch 训练脚本刚开始显存占用很正常比如 8G 的卡占用 6G结果跑着跑着变成 7.5G过一会儿直接 OOM进程被杀。重跑一遍还是老样子前几个 batch 正常后面必然爆炸。这种“训练越跑显存越高、最后爆掉”的问题十有八九就是显存泄漏。显存泄漏和普通的内存泄漏原理类似。内存泄漏是你 malloc 了空间但一直不释放进程占用的 RAM 只增不减显存泄漏则是你在 PyTorch 里分配了张量、计算图、中间激活值但这些对象在迭代结束后没有被正确释放导致 GPU 上的显存被“僵尸数据”越占越多。区别在于显存泄漏的后果更直接一旦达到显存上限CUDA 直接报CUDA out of memory训练中断前面的时间全白费。为什么 PyTorch 会泄漏显存核心原因是 PyTorch 的显存分配机制。PyTorch 默认使用缓存分配器CachingAllocator它会缓存已经释放的显存块以便后续重复利用而不是立刻还给 CUDA。这个机制本身没问题能减少分配开销、提升训练速度但问题在于如果训练逻辑里有 bug导致某些张量在loss.backward()之后仍然被持有引用或者把一些中间结果不断累积到 Python 列表里那么这些缓存块就永远不会被复用新的训练迭代还会继续申请新显存最终把显存塞满。我自己第一次遇到这个问题时处理方式很蠢直接换更大显存的卡或者调小 batch size。治标不治本甚至有时候显存占用没降多少反而因为 batch 变小训练变慢了。后来认真排查才发现90% 的显存泄漏都是代码层面的问题不是显卡问题。所以这篇文章我会从定位思路、排查工具、常见代码陷阱到修复方案给你一套可以直接照做的流程。不管你是刚入门 PyTorch 的新手还是已经在训练业务模型的老手这套方法都能让你少走几个星期弯路。2. 定位泄漏前的准备工作环境、工具和关键概念2.1 先确认是不是真的泄漏很多人一看到显存占用升高就慌其实要先区分“正常波动”和“真泄漏”。正常训练中显存占用会随着迭代进行小幅波动尤其是 batch 数据、随机采样、数据加载器的 worker 进程都会导致短时间占用变化。但要判断是否泄漏有一个简单粗暴的方法在训练循环的每个 epoch 结束或者每个固定迭代步数时打印一次当前显存占用。如果显存占用随着 epoch 单调上升几乎不回落那就是泄漏如果只是在一个 epoch 内部升高、epoch 结束又降下来那可能只是该 epoch 内累计了一些未释放的中间数据不算严重。我常用的验证脚本大概是这样的import torch import gc def print_gpu_mem(step, epoch): allocated torch.cuda.memory_allocated() / 1024**2 reserved torch.cuda.memory_reserved() / 1024**2 print(fStep {step} | Epoch {epoch} | Allocated {allocated:.1f} MB | Reserved {reserved:.1f} MB)注意memory_allocated()是实际张量占用的显存memory_reserved()是 PyTorch 从 CUDA 申请的缓存池大小前者才是我们排查泄漏的主要指标后者高一点是正常的因为缓存设计就是这样。如果allocated稳定升高那基本就是有张量没释放。另外还要确认环境因素有没有用 DataLoader 的num_workers如果设置了多个 worker每个 worker 会复制部分数据到 GPU显存占用也会有额外波动。这种不算泄漏但会让显存基准线变高排查时容易误判。建议先用num_workers0跑一个 demo排除 DataLoader 干扰确认是训练逻辑本身的问题。2.2 显存泄漏的类型张量表显存 vs 缓存分配器显存这里要说清楚PyTorch 显存泄漏分两种常见形态很多人混淆了导致排查方向错误。第一种是“真实泄漏”Python 对象中持有对 GPU 张量的引用导致这些张量在训练循环中不断累积。典型例子是把每个 batch 的 loss、输出或者损失项 append 到一个 list 里但第二天你没在意代码里一直没有清空这个 list。由于 GPU 张量被 list 持有引用PyTorch 的垃圾回收机制无法回收它们的显存。这种泄漏只要你把 list 清空显存立刻掉下来。第二种是“缓存分配器膨胀”PyTorch 的 CachingAllocator 为了性能会缓存不同大小的显存块。如果你的代码在每次迭代中创建了各种各样尺寸的临时张量比如动态 padding、不同长度的序列产生不同 shape 的中间结果分配器可能需要不断新开显存块来满足需求而旧块由于尺寸不匹配无法被复用导致 reserved 显存不断增长。这种“泄漏”其实不是真正持有引用而是缓存碎片化。我之前用torch.cuda.memory_reserved()看发现它涨得比allocated快很多就是这个原因。这两种类型的定位思路完全不同。真实泄漏要用gc模块、引用跟踪和objgraph之类的工具来找对象缓存膨胀则需要检查代码中是否存在动态 shape 的中间张量、是否频繁创建Tensor而完全不复用。后面我会分别给方案。2.3 必备的排查工具清单工欲善其事必先利其器。我推荐几个工具按需使用torch.cuda.memory_allocated()和torch.cuda.memory_reserved()最基本的两个 API用于快速看显存占用趋势。torch.cuda.max_memory_allocated()返回本次程序运行以来分配的最大显存结合reset_peak_memory_stats()重置可以定位是哪个阶段显存峰值飙升。torch.cuda.snapshot_memory()生成显存快照导出后可以用 Chrome 的chrome://tracing或者pytorch.org/memory_viz可视化看到详细的张量分配情况。这个是高级工具但排查复杂问题很有效。nvidia-smi监控整体显存占用但注意它显示的是进程级别的显存总量没法告诉你哪个 Python 对象占了显存只能做宏观参考。pympler和objgraph辅助定位 Python 对象持有引用的问题。gc.collect()配合gc.get_objects()检查垃圾回收行为有时候 PyTorch 的循环引用导致显存无法回收这一步能辅助判断。我个人习惯是先用memory_allocated每 100 步打一次点画出曲线判断形态然后再用 snapshot 做细粒度分析。不要一上来就上复杂工具先用最简单的指标定性。3. 最常见的五种泄漏代码模式你大概率踩过其中一个这部分我重点梳理实操中出现频率最高的几种泄漏代码写法每一条都是我在项目里踩过或帮别人定位过的问题。你可以对照检查自己的训练脚本。3.1 在 list 或 dict 里累积张量名气最大的泄漏模式也是最容易犯的。很多同学喜欢把 loss、输出特征、梯度等添加到 list 里用于后续统计或可视化但忘了清空。比如all_losses [] for epoch in range(epochs): for batch in dataloader: output model(batch) loss criterion(output, batch_target) all_losses.append(loss.item()) # 这里如果存的是 loss.item() 还好 # 但如果你写 all_losses.append(loss)就泄漏了如果 append 的是loss这个张量即使你后续只打印tensor.item()Python 的 list 会持有完整的过程张量引用。并且由于loss是在计算图里的反向传播结束后计算图理论上释放了但 list 持有 tensor导致显存无法释放。修复方法很简单只存标量用.item()或.detach().cpu().numpy()并且隔一段时间all_losses.clear()。如果你必须要保留输出特征做可视化使用 detach 后存入并控制存放数量上限比如用collections.deque(maxlen1000)。类似的还有把batch_idx遍历时的中间变量累积为全局变量、把每个 batch 的output.cpu()存起来等。记住一句话GPU 张量离开 GPU 用途后不要长时间保留如果要存就转成 CPU 的 numpy 或普通 Python 类型。3.2 在计算图中保存中间变量PyTorch 的自动求导机制是动态图。每个前向传播过程中如果某个变量被标记为需要梯度requires_gradTrue并且它被用于后续损失计算那么它的计算图节点和中间激活值都会被保存直到反向传播结束才会释放。这个机制是必需的但如果你在代码里隐式保留了对中间张量的引用比如hidden_states [] def forward_hook(module, input, output): hidden_states.append(output) # 这里保存了 GPU 张量且未 detach return output model.register_forward_hook(forward_hook)一旦 hook 里把 output 存在外部列表每个前向传播都会累积一个隐藏层的输出而且这些输出通常还带着 requires_gradTrue计算图不会释放显存就飞涨了。即使你只是想做特征提取也必须在 hook 里做output.detach().cpu()并且最好是每次都把列表清空或者直接在 hook 里把结果转成 numpy 存到文件中。另外还有一个容易忽略的点如果你自定义了nn.Module的forward函数里面写了out self.some_layer(x)然后又kept.append(out)哪怕你后面不用这个out它也会被 list 持有。所有带requires_gradTrue的保留引用都要警惕。3.3 动态 shape 导致缓存碎片化如果你做 NLP 或序列相关任务经常会用 pad_sequence、attention mask、不同长度的输入。这会导致同一个 batch 内的张量 shape 不固定或者不同 batch 之间的 shape 变化很大。PyTorch 缓存分配器对相同尺寸的内存块是有复用优先级的如果你每次迭代分配的都是不同大小的张量旧缓存块永远匹配不上新请求于是 reserved 越来越大。这种情况很像宿舍床位你入住时分配了一个床位退房后床位还留着但下一个人身高体重不同睡不下这个床宿管就再分配一个新床。一直这样搞宿舍楼满了。解决方法尽量把所有序列 padding 到同一个最大长度或者用paddingmax_length之类的固定策略。严格控制 batch 内的 padding 逻辑不要采用动态最大长度。如果实在避免不了动态 shape考虑手动调用torch.cuda.empty_cache()在 epoch 结束后清理缓存碎片但这只是缓解不能根治。这里多提一句torch.cuda.empty_cache()只是把缓存池中未使用的块还给 CUDA不会影响已分配的张量。它能解决 reserved 虚高的问题但解决不了 allocated 居高不下的引用问题。很多人以为定期调用它就能修复泄漏实际只能治标。3.4 循环中创建优化器、损失函数或移动模型到 GPU有这个坏习惯的人不多但一旦遇到就非常隐蔽。比如for epoch in range(epochs): optimizer torch.optim.Adam(model.parameters(), lrlr) # 每次循环重建优化器 ...如果优化器重建旧优化器中的状态动量、梯度缓存就会被丢弃但 PyTorch 的缓存分配器不一定立刻释放旧的显存块于是每轮 epoch 都会新分配一组优化器相关的显存。同样如果你在循环里反复model.to(cuda)并且每次都是同一个 model通常不会有大问题但如果创建新的nn.Module并to(cuda)旧模块释放和缓存复用之间可能会有延迟。更常见的坏习惯是循环内定义损失函数对象loss_fn nn.CrossEntropyLoss()或者每次迭代都x x.to(cuda)重新创建一个张量。单次不是问题但循环次数多了如果中间有对参数梯度的引用未释放显存就会缓慢增长。正确的做法是损失函数、优化器、模型都在训练循环之前创建遍历过程中只更新参数。这样显存分配是稳定的。3.5 梯度清零不彻底或反向传播后未置 Noneoptimizer.zero_grad()是最基本的操作但很多人只在训练循环的开头统一清零遇到特殊分支比如梯度累积就会漏。每次backward()时PyTorch 会计算所有叶子张量的梯度并累计到.grad字段。如果不清零.grad里的梯度张量会被继续保留而且每次反向传播都会更新它显存占用会随着梯度累积而增大。还有一种情况是反向传播后没有loss loss.detach()或者没有让计算图释放。虽然backward()之后计算图通常会被释放但如果你在前后有保留引用比如把 loss 存在一个 tuple 里那计算图还是占着显存。最稳妥的写法是optimizer.zero_grad() loss.backward() optimizer.step() # 如果不需要 loss 做反向传播之外的用途最好手动 detach loss loss.detach()对于梯度累积的场景积到足够步数后一定要optimizer.zero_grad()否则.grad越积越大虽然数值上你可能有意识做除法归一化但显存上梯度张量一直没释放。4. 逐一排查泄漏的实操流程从宏观指标到细粒度分析4.1 第一步确认泄漏类型并缩小范围建议先写一个非常简化的训练循环只包含一个 batch 的 forwardbackwardstep跑 100 个 iteration每 10 步打印一次显存看是否增长。如果简化循环不增长说明问题出在数据加载、batch 采样、后处理等外围逻辑如果增长再逐步缩小范围注释掉某些代码段看是否恢复。比如你可以跑一个基线脚本model MyModel().cuda() optimizer torch.optim.Adam(model.parameters(), lr1e-3) criterion nn.CrossEntropyLoss() for step in range(100): x torch.randn(16, 3, 224, 224).cuda() y torch.randint(0, 10, (16,)).cuda() optimizer.zero_grad() pred model(x) loss criterion(pred, y) loss.backward() optimizer.step() if step % 10 0: print(fstep {step}: allocated{torch.cuda.memory_allocated()/1024**2:.1f} MB, freserved{torch.cuda.memory_reserved()/1024**2:.1f} MB)如果这个基线脚本显存稳定那么问题肯定在你的数据或者训练循环额外逻辑里。如果这个脚本也增长可能是模型本身有什么奇怪的 state 或者动态图问题。这种二分法定位特别高效。4.2 第二步用 filename 定位每个张量的来源PyTorch 的Tensor都带有一个_file和_line属性记录了创建这个张量的代码位置。在 PyTorch 1.10 中有个方法可以直接查看所有 GPU 张量的位置import torch def trace_gpu_tensors(): for obj in gc.get_objects(): try: if torch.is_tensor(obj) and obj.is_cuda: print(fGPU tensor: shape{tuple(obj.shape)}, dtype{obj.dtype}, fdevice{obj.device}, file{obj._file}, line{obj._line}) except Exception: pass虽然遍历所有对象会有点慢但在排查中期很有效。你会看到很多重复的file和line那就是你的泄漏热点。如果大量张量都来自同一个数据加载函数或者同一个 for 循环体那基本就是定位了。4.3 第三步显存快照可视化如果上面两步还没定位就需要用 PyTorch 自带的内存快照功能。你可以在任意训练迭代中调用torch.cuda.memory_snapshot()这个函数返回一段 JSON 序列化数据包含当前所有已分配缓存块的地址、大小、分配栈、活跃状态等。导出来之后可以用官方内存可视化工具from torch.cuda import memory memory._dump_snapshot_to_file(snapshot.pkl) # 不同版本 API 有差异或者你也可以按官方文档生成一个 HTML 或 Chrome tracing 文件。分析时重点关注已经 allocated 的块列表找到占用显存最大的几个块的调用栈stack trace然后跳转到对应源码处。我自己遇到过一次非常隐蔽的泄漏一个第三方库的forward里为了做可视化每帧把特征图cat到全局 list不看快照根本想不到是那里。4.4 第四步使用 gc 模块检查 Python 对象的引用关系如果是真实泄漏torch.cuda.memory_allocated()会持续上升但显存快照里显示的是缓存池的块不一定直接告诉你哪个 Python 对象持有引用。这时用gc.get_objects()加上objgraph可以画出对象引用图pip install objgraph然后在你怀疑的泄漏处打印import objgraph objgraph.show_growth(limit20) # 显示新增对象的增长情况show_growth会告诉你哪些类型的对象在快速增长。如果大量新增Tensor、Variable对象再结合 stack trace 定位。这个方法适合排查 Web 服务这类跑很久才泄漏的场景训练任务也可以用但要注意gc.get_objects()开销较大不要在每步都调。4.5 第五步使用torch.profiler分析张量分配热点PyTorch 官方 profiler 可以记录每个操作分配和释放的显存。在训练循环外包裹from torch.profiler import profile, ProfilerActivity with profile(activities[ProfilerActivity.CUDA], profile_memoryTrue) as prof: run_training_steps() prof.export_chrome_trace(memory_trace.json)分析 trace 时看Memory Usage那一栏哪些操作上调用了大量cudaMalloc尤其是重复分配相同大小显存块的操作基本就是问题所在。不过这个工具对代码侵入性有点强通常我是在前几个步骤没定位出来才用它。5. 修复技巧与实战案例从代码层面解决 90% 的显存飙升5.1 修复“list 持有引用”的标准姿势假如你定位到了all_losses.append(loss)修复非常简单。但我想强调一个更全面的习惯所有需要保留的数据统一只留float或numpy标量所有需要在多个 epoch 之间共享的 CPU 数据存放在固定大小的队列中定期 pop 出队。一个安全的数据收集器可以这样写from collections import deque class MetricBuffer: def __init__(self, maxlen500): self.buffer deque(maxlenmaxlen) def append(self, tensor): # 只保存标量避免 GPU 张量引用 self.buffer.append(tensor.detach().cpu().item()) def get_average(self): return sum(self.buffer) / len(self.buffer) if self.buffer else 0.0用deque(maxlen500)的好处是即使你忘了清空它也会自动丢弃最老的元素上限可控不会无限增长。5.2 修复动态 shape 导致的缓存膨胀对于缓存碎片化我的建议是按“固定长度优先、长序列再切分”的原则重构数据预处理。以 NLP 为例与其每个 batch 都用动态最大长度不如统一设定max_seq_len512把超过的截断不足的 padding 到统一长度。这样每个 batch 的张量大小总是一致的PyTorch 分配器可以稳定复用缓存块。如果因为业务需求必须使用动态 shape你还可以在代码中手动控制缓存块的粒度比如预先分配一个大 block 并反复切用但这实现成本高不太建议入门者来做。落地一点的替代方案是每隔固定迭代次数执行一次torch.cuda.empty_cache()把碎片化的空闲块归还给 CUDA避免 reserved 持续膨胀。虽然不能根治但对训练稳定性已经有明显改善。我在一个多标签分类项目里用深度可分离卷积做特征提取每步预测时用不同尺寸的 RoI 特征显存 reserved 从 3G 一路涨到 7G加了empty_cache()后稳定在 4.5G 左右。治本方案还是统一了输出特征图尺寸让缓存复用率达到 90% 以上。5.3 正确使用torch.no_grad()和torch.inference_mode()如果你在训练阶段需要计算一些指标如准确率、IOU务必把这些计算放到torch.no_grad()下面。如果不这么做pred model(x)会产生新的计算图该计算图会一直持有中间张量直到你把它消耗掉比如调用.item()后可能释放但如果你存了 logits 的列表那就泄漏了。更严格的替代方案是torch.inference_mode()它比 no_grad 更进一步禁止自动求导跟踪显存占用更低、速度更快。举个例子在验证集上计算准确率时model.eval() with torch.no_grad(): for batch in val_dataloader: x, y batch pred model(x) # 只拿数值或 numpy 结果不要保留 pred 的 GPU 版本 acc (pred.argmax(dim1) y).float().mean().item()如果你在里面直接preds.append(pred)恭喜你踩坑了。正确做法是preds.append(pred.cpu().numpy())并且明确pred不再参与反向传播。5.4 梯度清零的陷阱与最佳实践梯度清零看似简单但在梯度累积场景下容易出错。比较安全的做法是使用model.zero_grad(set_to_noneTrue)而不是optimizer.zero_grad()。set_to_noneTrue会把.grad置为None而不是分配一个新的零张量这样可以避免每次迭代都重新分配梯度显存长期运行能显著降低显存峰值和抖动。代码示例optimizer.zero_grad(set_to_noneTrue)如果你的 PyTorch 版本较老不支持这个参数那就手动把每个参数的 grad 置为 Nonefor p in model.parameters(): p.grad None另外在使用loss.backward()之后如果这个 loss 生命周期已经结束可以显式del loss并gc.collect()但这里要注意gc.collect()本身有开销建议只在特定检查点比如每个 epoch 结束调用不要每个 step 都调。5.5 数据加载器对显存的影响num_workers与pin_memory数据加载器也可能导致显存异常升高。pin_memoryTrue会把数据放到锁页内存Page-Locked Memory虽然加速了数据拷贝到 GPU 的过程但如果你的collate_fn里对每个样本都调用了.cuda(non_blockingTrue)可能会提前将数据送入 GPU 并且不释放。排查方法是先设num_workers0, pin_memoryFalse测试显存基线再逐步开启。如果开启后显存多了几百 MB那属于正常范围因为 worker 进程缓存了部分样本但如果你发现有持续上升就要检查数据加载线程中的 GPU 操作。我建议在collate_fn中只做张量拼装不要做.cuda()转换。统一把.cuda()留在训练循环里等 batch 完整之后再转设备这样显存分配更可控。5.6 实战案例一个语义分割模型的显存泄漏修复过程分享一个我实际修过的例子帮你看完整流程。当时跑的是一个基于 Mask2Former 的语义分割模型batch size 设为 2显卡为 24G 的 3090。大概训练到 200 个 iteration 后显存从 18G 涨到 23G然后 OOM。第一次排查时我怀疑是 loss 累积但检查代码没有问题。后来用torch.cuda.memory_snapshot()做了分析发现大量显存被一个torch.nn.functional.interpolate操作分配的不同尺寸特征图占据。原来模型里的一个辅助解码器为了做多尺度特征融合每次前向都生成一个(batch, 2, H, W)的临时张量而 H 和 W 是动态变化的根据输入分辨率缩放导致缓存碎片化。修复方案很简单把所有辅助特征统一插值到固定尺度比如 32x32放在模型定义里做完插值再输出。改完后显存占用曲线几乎是一条直线峰值还从 23G 降到了 21G。这个案例说明动态 shape 的碎片化问题有时候比显式持引用更隐蔽需要靠快照工具才能发现。6. 进阶排查妙招代码审查之外还有这些“黑盒”手段6.1 在脚本中强制设置环境变量提前暴露问题很多 PyTorch 显存问题与 CUDA 缓存分配器的行为有关。在调试阶段你可以通过环境变量禁用缓存分配器让每次分配都直接调用 CUDA 的cudaMalloc这样虽然慢但显存上涨会立刻反映真实分配情况export PYTORCH_NO_CUDA_CACHING_ALLOCATOR1跑一遍训练脚本如果显存还是持续增长那就说明确实有张量没释放如果显存变得稳定了说明主要是缓存复用策略不佳。这个环境变量不需要一直用它只是帮你区分“真实泄漏”和“缓存碎片膨胀”非常有价值。6.2 使用tracemalloc监控 Python 内存分配虽然显存泄漏的是 GPU但 Python 侧的内存分配往往牵扯到 GPU 张量的释放。tracemalloc可以追踪 Python 对象的内存分配位置结合gc.get_objects()能判断哪个 Python 对象在持续增长间接定位显存泄漏。示例import tracemalloc tracemalloc.start() # 训练循环中每隔一段时间 snap1 tracemalloc.take_snapshot() # 跑一段时间后 snap2 tracemalloc.take_snapshot() top_stats snap2.compare_to(snap1, lineno) for stat in top_stats[:10]: print(stat)不过我没太常用这个因为 PyTorch 张量本身的 Python 对象是轻量的显存大头在 C 端tracemalloc 看不到 C 分配的显存。它就只是个辅助线索结合objgraph一起用。6.3 复现最小复现集方便向社区提问或自查如果你尝试了各种方法还是定位不到最好的办法是构造一个最小复现脚本。把模型换成两个线性层自定义一个简化版的数据集把可疑的代码段保留其他全部去掉。如果最小复现集还能复现显存泄漏那你就可以逐步删减直到剩下唯一一个可能触发泄漏的代码段。这个方法我推荐给所有入门同学在社区提问时简洁的最小复现集往往能让你快速获得有效帮助。我自己做最小复现集时会刻意保留以下部分一个固定随机种子的模型和数据精简的训练循环不包含任何日志、验证、检查点每 20 步打印显存用torch.cuda.empty_cache()和gc.collect()前后对照看能否回落如果精简后问题消失了说明是外围逻辑引起如果还在再层层嵌套加回其他模块。7. 防止显存泄漏的编码习惯从源头扼杀问题7.1 统一使用with torch.no_grad()包裹推理/指标计算在训练脚本里前向反向之外的任何网络前向计算比如验证集评估、计算 mAP、生成注意力图等全部放在torch.no_grad()或torch.inference_mode()上下文中。这个习惯能阻止所有意外计算图的累积。7.2 为训练循环增加显存护栏写一个每 100 步检查一次显存清理的函数如果显存没有回落就强制gc.collect()和torch.cuda.empty_cache()并输出警告。虽然这不是治本方案但能防止在长训练中突然 OOM给你留出人工干预的时间。def check_memory_safety(step, threshold_mb100): allocated torch.cuda.memory_allocated() / 1024**2 if allocated threshold_mb: print(fWarning: allocated {allocated:.1f} MB at step {step}, cleaning up.) gc.collect() torch.cuda.empty_cache()注意这个护栏不要设得太敏感否则频繁 empty_cache 会拖慢速度。比如你的模型 baseline 占用就 6G阈值可以设为 6.5G只在异常增长时触发清理。7.3 关注 PyTorch 版本差异不同 PyTorch 版本的缓存分配器行为不一致。有些旧版本在Tensor.data被修改后不会正确释放显存有些版本中no_grad覆盖不完全。如果你排查了很久仍无结果尝试升级到最新稳定版 PyTorch或者降级到你最初稳定运行的版本。我在 1.10 时代遇到过torch.cuda.empty_cache()偶尔无效的 bug升级到 1.13 后就好了。7.4 分布式训练中的显存泄漏注意点如果你在用DistributedDataParallelDDP或Horovod显存泄漏的原因更多样。常见的有DDP 在find_unused_parametersTrue时会有额外的梯度同步缓冲区每个 rank 进程如果都在加载同一个大数据文件到 GPU内存也会膨胀。排查 DDP 泄漏时要先用单卡跑同样的代码排除单卡本身就泄漏的可能再考虑分布式特有逻辑。使用torch.utils.checkpoint梯度检查点时由于前向传播会重新计算中间激活也可能导致前向过程中显存反复升降但不属于泄漏。如果你看到这类现象不要误判。8. 回归测试修复后如何验证显存问题真的解决了修复完泄漏很多人跑两步看到显存不涨就以为完事了结果训练一晚上又 OOM。回归测试要严谨一点建议做三件事。第一跑一个 500 步的小实验记录第 100/200/300/400/500 步的 allocated 显存画出一条趋势线。如果曲线斜率接近 0说明修复有效如果还是线性增长说明还有其他泄漏点没找到。第二测试torch.cuda.memory_reserved()和allocated的差距。如果 reserved 远大于 allocated并且持续增长说明缓存复用仍然不佳可以继续用empty_cache()缓解或者查动态 shape。第三设置一个长耗时训练任务比如 10 个 epoch每 epoch 打印显存并保存 checkpoint。如果第 10 个 epoch 的显存峰值与第 1 个 epoch 相差不到 200MB那基本可以断定泄漏已修复。另外还要留意多进程的情况。如果你的代码用了mp.spawn启动多进程每个子进程都会继承父进程的 CUDA context可能产生额外的显存开销这个不算泄漏但会造成看起来显存占用偏高。9. 写在最后一点个人经验总结排查显存泄漏的过程最忌讳的就是“头痛医头、脚痛医脚”——看到显存高就改小 batch size看到 OOM 就换更大显存的卡。这样不仅解决不了根本问题还会让你错过真正值得优化的方向。显存泄漏往往反映的是代码里有不合理的张量生命周期管理只要顺手清理干净训练速度和显存占用都会一并改善。我个人习惯是在写新训练脚本前就约定好一套规范loss 只保留标量、中间特征一律 detach 到 CPU、优化器和模型都在循环体外定义、定期用换行文本打印显存曲线。有了这套规范即使后来接手项目的人乱写代码也能很快定位问题。最后再送大家一个小技巧如果你是在 Jupyter Notebook 里跑训练频繁重新执行单元格也可能导致显存累积。每次重跑之前最好先再开一个新内核避免旧的 CUDA context 还占着显存。这一点虽然不算是代码 bug但实际中非常常见容易被忽略。希望这篇内容能帮你在下次遇到“训练越跑显存越高”时不再抓瞎。按着这套流程走一遍大部分问题都能在半小时内水落石出。
