1. 训练到第N个epoch突然OOM先别急着调batch size显存泄漏这个问题我敢说每个用PyTorch跑过稍大规模训练的人都遇到过。现象特别典型前几个epoch显存占用稳稳的跑到中途nvidia-smi一看显存曲线跟爬楼梯似的一个epoch比一个epoch高最后在某个随机的时间点直接甩给你一个CUDA out of memory。很多人的第一反应是把batch size砍一半或者上torch.cuda.empty_cache()结果发现要么治标不治本要么训练速度掉得没法看。这里要先厘清一个概念显存泄漏和显存峰值过高是两码事。峰值过高是模型本身或某个操作瞬时吃掉了大量显存比如attention矩阵在序列长度拉长时平方级膨胀而泄漏是显存占用随着训练步数单调递增即使输入尺寸完全不变。判断方法很简单——固定输入shape跑几百个step每隔几十步打印一次torch.cuda.memory_allocated()和torch.cuda.memory_reserved()如果allocated持续上涨那就是实打实的泄漏。这篇文章面向的是已经能跑通训练脚本、但被显存问题卡住的开发者。我会从计算图的持有关系讲起把常见的泄漏源头一个个拆开给出可复现的排查链路和修复代码。所有方案都在PyTorch 1.11到2.x上验证过不依赖任何特殊硬件。提示排查前先确认你的环境里torch.cuda.is_available()为True且驱动版本与CUDA版本匹配。环境本身有问题的话后面所有分析都是白费功夫。2. 计算图为什么会被扣住从autograd的引用链说起2.1 一张计算图的完整生命周期PyTorch的动态图机制决定了每次前向传播都会构建一张新的计算图。这张图由Function节点和Tensor边组成每个非叶子节点的Tensor都持有grad_fn而grad_fn又通过next_functions指向它的输入节点。正常情况下你调用loss.backward()之后PyTorch会释放掉这张图里保存的中间激活值只保留叶子节点的.grad。问题出在正常情况下这四个字。只要有任何一条Python引用链还指向计算图中的某个非叶子Tensor整张图就无法被释放。这就像你搬家时把钥匙落在了旧房子里只要钥匙还在搬家公司就没法把房子清空。常见的钥匙包括被append进list的loss、被存进dict的中间输出、被闭包捕获的Tensor、被全局变量引用的模型输出。2.2 一个最小可复现的泄漏demo先看一段能稳定复现泄漏的代码后面所有分析都围绕它展开import torch import torch.nn as nn class Net(nn.Module): def __init__(self): super().__init__() self.fc nn.Linear(1024, 1024) def forward(self, x): return self.fc(x) model Net().cuda() optimizer torch.optim.SGD(model.parameters(), lr0.01) losses [] # 罪魁祸首 for step in range(1000): x torch.randn(64, 1024).cuda() y model(x) loss y.sum() losses.append(loss) # 持有计算图引用 loss.backward() optimizer.step() optimizer.zero_grad() if step % 100 0: print(step, torch.cuda.memory_allocated() / 1024**2, MB)跑一下你会发现显存占用随着step线性增长。原因就是losses这个list把每个step的loss都存了下来而loss是计算图的根节点持有它就等于持有整张图。修复方式有两种存loss.item()转成Python float脱离计算图或者存loss.detach()共享存储但切断梯度链。2.3 detach和item的本质区别很多人分不清detach()和item()。detach()返回的是一个与原Tensor共享底层存储、但requires_gradFalse的新Tensor它仍然占用显存item()则把标量Tensor转成Python的float对象彻底脱离GPU。所以如果你只是想记录loss曲线用item()如果后续还要对这个值做GPU上的运算用detach()。注意with torch.no_grad()包裹的代码块里产生的Tensor默认不带梯度但如果你在块内把Tensor赋给了块外的变量它依然会存活。no_grad只影响新创建的Tensor不影响已有引用。3. 五类高频泄漏源从loss累积到自定义autograd3.1 累积loss、metric和中间输出这是最常见的一类。除了上面demo里的list还有几种变体把每个step的output存进list做可视化、把hidden_state缓存起来做时序分析、在验证循环里忘了切eval()模式导致BN层持续更新running stats并保留图。排查方法是在训练循环里加一行gc.get_referrers(loss)看看除了局部变量还有谁在引用它。修复原则很简单任何跨越step边界存活的对象都必须脱离计算图。记录指标用.item()缓存特征用.detach().cpu()需要保留梯度的场景比如梯度累积要显式管理生命周期。3.2 优化器状态与参数注册的隐性持有优化器本身会持有参数的引用这是正常的。但如果你在训练过程中动态地往model里注册新的buffer或parameter比如self.register_buffer(running_feat, feat)而这个feat是从计算图里来的那优化器或模型就会一直扣着这张图。还有一种情况是用了torch.nn.ParameterList或ModuleList在forward里动态append导致每次前向都新增参数显存自然只增不减。检查方法训练几个step后打印len(list(model.parameters()))和len(list(model.buffers()))如果数字在涨那就是动态注册的问题。3.3 自定义autograd.Function的ctx泄漏写自定义算子时forward里的ctx.save_for_backward()如果保存了不该保存的Tensor或者backward里返回的梯度引用了外部变量都会造成泄漏。更隐蔽的是在forward里把中间结果存到了self上class MyFunc(torch.autograd.Function): staticmethod def forward(ctx, x): ctx.save_for_backward(x) MyFunc.cache x * 2 # 类属性持有永久泄漏 return x * 2这种类属性或模块级变量持有Tensor的写法是很多跑着跑着就OOM的元凶。修复方式是改用ctx.save_for_backward并确保backward返回的梯度不引用任何外部Tensor。3.4 DataLoader的worker与pin_memoryDataLoader的num_workers 0时每个worker进程会独立加载数据。如果collate_fn里返回了CUDA Tensor或者dataset的__getitem__里做了GPU操作worker进程会持有显存。另外pin_memoryTrue会使用锁页内存虽然不直接吃显存但配合worker的预取队列prefetch_factor会放大内存占用。排查时把num_workers设为0跑一遍如果泄漏消失问题就在数据加载侧。修复方案是确保dataset只返回CPU Tensor把.cuda()操作放到训练循环里。3.5 全局缓存与闭包捕获最后这类最阴间在Jupyter Notebook里反复运行同一个cell旧的模型、优化器、Tensor都还挂在_变量或Out字典里或者用functools.lru_cache缓存了带Tensor的函数返回值再或者用闭包捕获了训练循环里的变量。这类泄漏在脚本模式下不容易出现但在交互式环境里极其常见。泄漏源典型症状快速验证方法loss累积显存随step线性涨注释掉append后重跑动态注册参数量递增打印parameters数量autograd.Function特定算子后开始涨替换为原生算子对比DataLoader多worker时更严重num_workers0对比全局缓存重启kernel后正常脚本模式重跑4. 定位工具链memory_summary、snapshot与引用追踪4.1 torch.cuda.memory_summary的读法PyTorch内置的torch.cuda.memory_summary()能打印出当前显存的详细分布包括allocated、reserved、active、inactive等指标。关键看两个数allocated是实际被Tensor占用的显存reserved是PyTorch向CUDA申请的总量。如果allocated稳定但reserved持续涨说明是碎片化问题不是泄漏如果allocated本身在涨那才是真泄漏。print(torch.cuda.memory_summary(deviceNone, abbreviatedFalse))输出里会按Tensor大小分桶统计能快速看出是大Tensor泄漏还是小Tensor堆积。4.2 用memory_snapshot做逐行归因torch.cuda.memory_snapshot()返回一个列表每个元素代表一个内存块包含size、allocated_size、frames等字段。frames里记录了分配这个块时的Python调用栈。你可以写个小脚本把snapshot按frames聚合找出哪个文件哪一行分配的显存最多import torch, collections snap torch.cuda.memory_snapshot() counter collections.Counter() for block in snap: if block[allocated_size] 0: for frame in block[frames]: counter[frame[filename] : str(frame[line])] block[allocated_size] for loc, size in counter.most_common(10): print(loc, size / 1024**2, MB)这个方法的威力在于它能直接告诉你泄漏发生在你的代码的哪一行而不是笼统地说某个Tensor没释放。4.3 gc与weakref追踪引用链当snapshot定位到具体行之后用gc.get_referrers()反查谁在引用这个对象import gc target some_leaked_tensor for ref in gc.get_referrers(target): print(type(ref), ref)如果引用者是list、dict或某个模块的属性基本就能锁定泄漏源。对于复杂场景可以用weakref.ref给可疑对象加弱引用回调在对象被回收时打印日志确认它到底有没有被释放。提示memory_snapshot在PyTorch 1.10以上才比较完善老版本建议升级。另外snapshot本身会消耗一定显存别在OOM边缘调用。5. 修复实战从empty_cache的误用到梯度累积的正确姿势5.1 empty_cache到底该不该用torch.cuda.empty_cache()的作用是把PyTorch reserved但未使用的显存还给CUDA驱动让其他进程能用。它不解决泄漏只是把已经泄漏的显存从reserved变成free但那些被引用的Tensor依然占着allocated。频繁调用还会导致后续分配变慢因为PyTorch得重新向驱动申请。正确用法在显存碎片化严重、且确认没有泄漏的前提下偶尔调用一次。比如在验证循环开始前或者模型切换时。把它当成整理房间而不是扔垃圾。5.2 梯度累积的显存陷阱梯度累积的常见写法是for i, (x, y) in enumerate(loader): loss model(x, y) / accum_steps loss.backward() if (i 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()这里如果loss被存进了list或者model(x, y)的输出被缓存泄漏照样发生。更隐蔽的是有些人为了节省显存在累积期间不调用zero_grad()结果梯度一直累加虽然不算泄漏但会导致梯度爆炸。梯度累积的正确姿势是每次backward后图自动释放只要不持有loss引用显存就是稳定的。5.3 用del和gc手动干预的时机在确认某个大Tensor不再需要时可以显式del tensor并调用gc.collect()。但要注意del只是删除引用如果还有其他引用存在对象依然存活。在训练循环里局部变量会在下一次迭代时被覆盖通常不需要手动del。真正需要手动干预的场景是在循环外创建了大Tensor循环内反复使用用完后想提前释放。big_cache torch.randn(10000, 10000).cuda() for step in range(1000): out model(big_cache) # ... 训练逻辑 del big_cache torch.cuda.empty_cache()5.4 一个完整的修复案例回到第2节的demo完整修复版本如下import torch import torch.nn as nn class Net(nn.Module): def __init__(self): super().__init__() self.fc nn.Linear(1024, 1024) def forward(self, x): return self.fc(x) model Net().cuda() optimizer torch.optim.SGD(model.parameters(), lr0.01) losses [] for step in range(1000): x torch.randn(64, 1024).cuda() y model(x) loss y.sum() losses.append(loss.item()) # 关键修复脱离计算图 loss.backward() optimizer.step() optimizer.zero_grad() if step % 100 0: print(step, torch.cuda.memory_allocated() / 1024**2, MB)改动只有一处但显存曲线会从爬楼梯变成一条水平线。这就是定位到根因之后修复的威力——不需要调batch size不需要加empty_cache一行代码解决。6. 几个反直觉的坑no_grad、inference_mode与BN的隐藏开销6.1 no_grad不等于不占显存torch.no_grad()只是让新创建的Tensor不带梯度但前向传播的中间激活值依然会占用显存只是backward时不需要它们所以可以在前向结束后立即释放。如果你在no_grad块里把中间结果存到了外部变量显存照样涨。验证循环里最常见的错误就是with torch.no_grad(): for x, y in val_loader: out model(x) all_outs.append(out) # 依然持有Tensor正确做法是all_outs.append(out.cpu())或者只记录指标。6.2 inference_mode比no_grad更彻底PyTorch 1.9引入了torch.inference_mode()它比no_grad更激进不仅不记录梯度还会禁用版本计数和autograd相关的元数据显存占用和速度都更优。纯推理场景建议优先用它。但注意inference_mode下创建的Tensor不能参与后续需要梯度的计算否则会报错。6.3 BatchNorm的running stats与显存BN层在训练模式下会更新running_mean和running_var这些是buffer不占多少显存。但如果你的模型在验证时忘了model.eval()BN会继续用batch统计量并更新running stats同时因为处于训练模式前向会构建计算图显存占用直接翻倍。这是一个非常隐蔽的坑尤其在自定义训练循环里容易漏掉。6.4 多卡训练下的显存不均用DataParallel时显存占用往往集中在主卡上因为loss和梯度汇总都在主卡。DistributedDataParallel则相对均衡。如果你发现只有某张卡OOM先确认是不是用了DataParallel。另外多卡场景下每张卡都会维护一份完整的模型副本和优化器状态显存需求是单卡的N倍N为卡数这个账要提前算清楚。坑点错误写法正确写法验证循环no_grad里append Tensorappend.cpu()或.item()推理模式用no_grad用inference_modeBN模式验证时忘eval验证前model.eval()多卡DataParallel优先DistributedDataParallel7. 把显存监控做成习惯训练脚本里的自检清单7.1 每个epoch打印一次显存指标在训练循环里加几行监控代码成本极低收益极高def log_memory(step): alloc torch.cuda.memory_allocated() / 1024**2 reserv torch.cuda.memory_reserved() / 1024**2 print(fstep {step}: allocated{alloc:.1f}MB reserved{reserv:.1f}MB)如果allocated在固定输入下持续上涨立刻停下来排查别等到OOM才动手。7.2 用tracemalloc配合CPU侧排查显存泄漏有时候是CPU侧对象泄漏的连带结果。tracemalloc能追踪Python对象的内存分配配合gc可以找出哪些Python对象在增长。虽然它不直接追踪GPU显存但能帮你发现list、dict的无界增长。7.3 训练前的显存预算跑之前先算一笔账模型参数显存 参数量 × 4字节fp32优化器状态以Adam为例 参数量 × 8字节一阶矩二阶矩激活值显存与batch size、序列长度、层数成正比。把这三部分加起来再留20%余量就是你的显存下限。如果预算超了先调模型结构或batch size别指望靠empty_cache硬撑。7.4 一个可复用的排查流程我把上面所有内容浓缩成一个排查流程遇到OOM时按顺序走固定输入shape打印allocated曲线确认是否泄漏。检查所有跨step存活的容器list、dict、全局变量确保存的是item或detach后的CPU Tensor。把num_workers设为0排除DataLoader侧问题。用memory_snapshot定位到具体代码行。用gc.get_referrers反查引用链。修复后重跑确认allocated曲线水平。这套流程我在多个项目上用过从简单的图像分类到复杂的多模态训练基本都能在半小时内定位到根因。显存泄漏不可怕可怕的是盲目调参和滥用empty_cache。把计算图的引用关系搞清楚问题就解决了一大半。