跑PyTorch训练的工程师大多都经历过这种让人血压飙升的时刻早上看显存占用还是10GB中午变成16GB下午直接报CUDA out of memory训练戛然而止。更诡异的是只要你重启训练同一个脚本又能从头开始正常跑。这种“越跑显存越高”的现象就是典型的PyTorch显存泄漏问题。这个问题的麻烦在于它不像普通的程序崩溃那样给你一个明确的错误栈很多时候你只能从缓慢变慢的速度和一格一格涨上去的显存里感受到异常。我这些年排查过不少这类问题从最开始的瞎试empty_cache到后来用内存快照精确定位到具体代码行积累了一套可复用的方法论。这篇文章我会完整讲一遍怎么判断真假泄漏怎么给训练脚本装上显存监控为什么计算图会把显存变黑洞以及三条从浅到深的定位手段。最后还有一个真实案例的全过程复盘帮你把前面所有方法串起来。不管你是在调一个已经跑崩的模型还是想提前给训练脚本打好预防针都能从这篇文章里找到能直接用的东西。1. 显存曲线带着训练一起失控先分清“假性增长”和“真泄漏”1.1 一眼识别“假性增长”先别急着动代码先说个反直觉的结论并非所有显存“增长”都是泄漏。我见过很多同事一发现显存涨了就怀疑是泄漏然后去改模型结构改了大半天发现根本没作用因为问题压根不在模型。要理解这一点得先搞懂PyTorch自己的缓存分配器。GPU显存的申请和释放走的是cudaMalloc和cudaFree这俩操作相对昂贵。如果每个Tensor用完都立刻归还给驱动训练速度会被拖垮。所以PyTorch内部实现了一个caching allocator你申请一块显存等Tensor不再使用后它并不会马上把这块显存还回去而是留在进程自己的缓存池里供后续分配直接复用。于是你打开nvidia-smi看到的进程显存占用其实包含两部分一部分是真正被Tensor占用的另一部分是PyTorch“屯”在缓存池里的。对应到PyTorch的API分别是torch.cuda.memory_allocated()和torch.cuda.memory_reserved()。正常情况下reserved曲线的形态是阶梯跳变。比如训练第一次遇到大的batch、大的卷积核或矩阵乘法时缓存池发现现有的块不够用就重新向驱动申请一大块显存表现在nvidia-smi里就是显存一下子跳上一个台阶。之后如果又遇到更大的需求再跳一档。这种增长不是逐batch爬升的而是在某几个时间点跳到新平台后保持基本平坦。判断“假性增长”的关键就是看增长到底是一次性跳到高位后保持稳定还是每个batch、每个epoch都在持续爬升。前者是缓存扩展后者才可能是泄漏。还有个小技巧如果计算图本身没问题PyTorch后续会复用缓存块训练几十上百个epoch之后显存占用应该是平稳的。提示不要一看到nvidia-smi里显存高位就判定泄漏。先区分allocated和reserved再看多轮epoch的曲线如果allocated在多轮之后稳定不涨说明缓存池正常只是预留偏大。1.2 真泄漏的两种典型形态单调爬升型和阶梯跳变型根据我观察过的实际案例真正的显存泄漏基本可以归成两类形态。第一种是“单调爬升型”。显存每个batch或每个epoch都增长固定大小十几MB、几十MB都有。这种问题通常是一个Tensor或缓冲区被创建后又被某些对象持续引用导致它永远无法被垃圾回收。典型写法就是训练循环里把每个step的中间变量append进一个list而list又挂在模型或训练器上或者自定义loss函数为了可视化保留了一堆中间feature。增长量级相对固定和每次新增Tensor的总字节数成正比很好识别。第二种是“阶梯跳变型”。显存在某个阶段突然上一个大台阶之后回不到原来的水平。常见触发条件是开启了torch.backends.cudnn.benchmark True后模型对多种输入尺寸做auto-tune不断重新分配卷积workspace也可能是验证集里出现了一个尺寸异常大的batch缓存分配器为了喂饱它把reserved整体调高即使验证结束缓存也不会缩回去。这类问题只要不是每个epoch都涨危害相对小一些但如果把reserved误当成allocated来分析很容易冤枉好人。还有一种比较少见的“叠加型”模型在做参数剪枝、动态构图这类操作时每过N个epoch新增一批Tensor引用显存呈周期性台阶增长。这种需要配合日志时间戳和显存曲线做联合分析比前两类都难定位一些。1.3 一个五步自检法五分钟内判断要不要深入排查我的习惯是不管三七二十一先跑一个固定流程用最快速度判断到底是“要不要进入深水区”在训练脚本里每100个step打印torch.cuda.memory_allocated() / 1024**3观察它的走势。单独跑50个step记录起始值和结束值的差。如果差值超过1GB且不是由batch大小变化引起基本可以断定有问题。在验证循环前后各打一次allocated对比差值。如果每次验证结束都涨一块而且缩不回去说明验证环节有引用未释放。把所有自定义日志、可视化回调全部关掉再跑100个step。如果显存停止增长问题出在回调或日志代码里而不是模型主体。在每次backward()和optimizer.step()之后强制torch.cuda.empty_cache()看看有没有改善。注意这一步只是排除法不是修复手段。这五步走完你对问题类型和范围基本心里有数了再进入下一步就不会抓瞎。2. 第一套监控脚手架让显存占用在日志里“显形”2.1 为什么别只依赖nvidia-sminvidia-smi是大多数人排查显存问题的第一反应但我必须说它在这个场景下有天然死角它显示的是整个进程的总体显存占用而且默认几秒刷新一次你很难从它看到增长的Tensor到底是谁分配出来的。更麻烦的是它区分不了“PyTorch缓存池里的预留块”和“真正的活体Tensor”很容易误导判断。PyTorch自带的几个内存统计接口才是排查主力memory_allocated()、memory_reserved()、max_memory_allocated()以及两个深水区武器torch.cuda.memory_summary()和torch.cuda.memory_snapshot()。memory_summary会按分配段、活跃块、保留块给出一份汇总文本让你快速判断显存是被活体Tensor占着还是碎在缓存池里。memory_snapshot则能导出一份包含每个分配事件、地址、大小和调用栈的JSON配合分析脚本直接定位到具体代码行。所以我的建议是别一上来就盯nvidia-smi先基于PyTorch的API在训练脚本里埋好监控让显存增长自动在日志里现形。2.2 一套可复用的训练显存监控装饰器下面这个装饰器是我在多个训练项目里一直在用的模式。不侵入训练逻辑每隔N个step打印一次allocated和reserved顺带记一下耗时import time import torch def cuda_monitor(step_interval10): def decorator(fn): def wrapper(*args, **kwargs): torch.cuda.synchronize() start_time time.time() result fn(*args, **kwargs) torch.cuda.synchronize() elapsed time.time() - start_time if wrapper.counter % step_interval 0: allocated torch.cuda.memory_allocated() / 1024**3 reserved torch.cuda.memory_reserved() / 1024**3 peak torch.cuda.max_memory_allocated() / 1024**3 print(f[step {wrapper.counter:6d}] fallocated {allocated:.3f} GB, freserved {reserved:.3f} GB, fpeak {peak:.3f} GB, ftime {elapsed:.3f}s) torch.cuda.reset_peak_memory_stats() wrapper.counter 1 return result wrapper.counter 0 return wrapper return decorator很多用PyTorch Lightning这类框架的朋友可以把逻辑搬进on_train_batch_start和on_train_batch_end回调思路完全一样。这里有两个细节必须强调。第一打印前一定要torch.cuda.synchronize()。GPU上的kernel是异步执行的如果不做同步显存统计可能反映出“还有一堆kernel没跑完”的假象导致你看不清真实分配趋势。第二打印频率别太高十步一百步一次就行。在大模型训练里日志本身也有IO开销每一步都打印会把单步耗时拉长反过来干扰了你判断“显存泄漏导致速度变慢”的准确度。2.3 异常出现时先记录哪三组指标发现显存异常时不要慌我先记录三组指标指标含义判定泄漏的价值allocated实际占用GPU显存的Tensor字节数短期内快速增长基本就是Tensor泄漏reservedPyTorch向驱动申请的显存总量阶梯跳变说明是缓存池扩容或benchmarkmax_allocated从上次reset以来的峰值分配评估是否接近OOM阈值确认峰值异常举一个很常见的例子allocated从2.1GB涨到2.8GB但reserved一直是4.0GB。你从nvidia-smi看到的进程占用一直是4.0GB根本没变化于是误以为没泄漏。其实PyTorch在缓存池内部不断调整块的划分活体Tensor占用一直在涨只是还没触发新一轮cudaMalloc而已。反过来如果allocated平稳但reserved每隔几百步跳一次说明模型对不同输入shape产生了新的缓存需求这时要检查数据方面有没有batch size波动或padding变化。这三组指标配合时间戳输出到日志之后训练脚本就相当于自带了一台显存示波器后面定位问题会轻松太多。3. autograd的“账本”怎么变成显存黑洞理解计算图留存机制3.1 计算图的生命周期反向传播后到底释放了什么很多人以为显存泄漏就是“变量没删干净”其实不完全是。根源在于PyTorch的autograd机制。前向传播过程中PyTorch会记录一张计算图每个节点保存参与运算的输入、输出以及反向传播需要的梯度函数。训练模式下调用loss.backward()后PyTorch遍历整张图计算每个叶子节点的梯度并把结果写入param.grad。正常来说梯度计算完成后计算图会被释放保存中间激活值的那些Tensor就变成垃圾可以被回收。但有几个例外会延长计算图的寿命。第一你在前向过程中把一个Tensor挂到了模型或某个对象的属性上而这个属性在训练期间没有任何被清理的时机。它相当于给计算图“续了命”下一轮前向又生成新图但旧图的引用一直没断显存自然只进不出。第二自定义loss或回调为了做可视化、统计指标把中间结果保存下来却没有在合适的地方释放或做CPU拷贝同样导致引用累积。第三也是新人最容易踩的loss变量本身。如果loss是标量Tensor在后面某个日志评分函数里被存了下来或者错误地把Tensor连计算图直接丢给add_scalar之类的方法那这个loss Tensor会连同整个计算图一起被长期引用。3.2 empty_cache不是万能的很多人的第一反应是“显存涨了我就调torch.cuda.empty_cache()把它清掉”。我必须泼一盆冷水如果增长的Tensor还被计算图引用着empty_cache一点用都没有。因为它只把缓存池里那些“已释放但还没归还给驱动”的空闲块还给驱动并不会去清理正在使用的Tensor。用一句话概括这里的关系泄漏时增长的是allocated不是缓存池而empty_cache清的是缓存池。所以排查时正确思路是找“谁还在引用这个Tensor”而不是想“怎么把显存清掉”。只有当训练代码本身没有泄漏只是某些临时大块数据导致reserved过高时empty_cache才有压缩显存的效果但这也是治标。真正稳的做法是让分配的Tensor在生命周期结束后不再被任何持久对象引用。3.3 中间变量为什么是泄漏的高发区把原理落到代码上其实就是一个引用生命周期问题。正常的训练循环里局部变量的作用域在循环体内迭代结束就销毁Python自动回收。但如果中间变量被放进了list、被挂到self上、被存进全局字典它就跨过了迭代边界成为持久引用。一个典型的隐性累积例子你想可视化embedding分布写了self.cached_feature x。如果这个x是参与计算图的GPU Tensor每一轮都会被新值覆盖看起来不会涨但如果可视化代码写的是visualize_list.append(x)那就是另一回事了——每一个step的x都不会被覆盖每步新增一份GPU显存list越来越长显存当然一路攀升。处理中间变量的核心原则是不能让任何Tensor对象尤其是参与计算图的Tensor以“append”“缓存列表”“持久属性”的形式跨迭代存活。需要可视化或日志的数据必须显式做detach().cpu()并且最好只保留统计值而不是原始Tensor。4. 三招定位“显存刺客”从统计指标一路挖到Tensor的出生地4.1 第一招memory_summary看整体头部这是最朴素的定位法适合先看整体。在训练循环里插入print(torch.cuda.memory_summary(abbreviatedFalse))输出会告诉你当前显存里有多少个segments、多少个active blocks、多少个inactive blocks。OOM时还能帮你区分是不是碎片化问题。碎片化的典型特征是free块非常多但能匹配请求大块的连续空间太少这跟真正的Tensor泄漏不是一回事。用memory_summary能快速区分。如果文中的active blocks占了绝大部分说明“有人真实占用大块显存不放手”可以继续深挖了。如果active blocks不多但inactive blocks堆了一大堆碎片问题可能偏向缓存池策略或数据shape波动而不是单一的Tensor引用泄漏。4.2 第二招memory_snapshot生成全量分配记录memory_summary只给汇总到不了Tensor级。想看到底是哪个Tensor在增长就得memory_snapshotsnapshot torch.cuda.memory_snapshot() import json with open(mem_snapshot.json, w) as f: json.dump(snapshot, f, indent1)snapshot返回的是一个列表每个segment对象包含address、total_size、allocated_size、active_size、requested_size、stream等字段segment里还有blocks数组。最关键的是每个block往往带有一个frames字段记录了这块显存是哪一行代码、哪个函数分配的。知道这一点以后定位精确到代码行就变成了可能。我自己写过一个简单聚类脚本遍历snapshot里所有active blocks按file_name和line_number聚合求和size输出Top 10占用最大的分配点。下面是一个核心片段可以直接改来用from collections import defaultdict def analyze_snapshot(snapshot): totals defaultdict(lambda: [0, []]) for seg in snapshot: for blk in seg.get(blocks, []): if blk.get(allocated_size, 0) 0 or blk.get(active_size, 0) 0: size blk.get(requested_size, 0) or blk.get(size, 0) frames blk.get(frames, []) if frames: last frames[-1] key f{last[filename]}:{last[line_number]}:{last[function]} else: key unknown totals[key][0] size totals[key][1].append(blk) for key, (total, _) in sorted(totals.items(), keylambda x: -x[1][0])[:10]: print(f{total / 1024**3:.4f} GB {key})需要提醒一句不同PyTorch版本的snapshot字段名略有差异我写这段代码基于2.x你的版本如果不一致以实际输出为准。4.3 第三招二分注释排查法控制变量永远最可靠如果嫌自己写脚本麻烦社区里也有一些现成的可视化工具。PyTorch官方Sphere曾提供基于snapshot的内存可视化notebook第三方项目torch_mem_viz能把snapshot渲染成HTML按内存段和分配栈做交互式查看。这些工具的价值在于你能直观看到哪块内存长时间active、哪个call stack在反复分配。但我必须说实际工程里最快的方法往往不是花哨的可视化而是二分注释排查法。在训练循环里逐步注释掉被怀疑的代码块每注释一块就跑50个step看显存曲线。比如先关掉验证循环显存稳了问题就在验证再关掉自定义回调显存稳了问题就在回调再关掉某个复杂loss分支逐步缩小范围。这种“控制变量”法听起来土但定位速度通常比任何工具都快因为它直接排除了大量无关因素。5. 五种容易养出显存泄漏的写法每个都配了修复版5.1 累计训练指标时忘写.item()这是最经典的坑。很多人写训练循环时会累积loss做日志total_loss 0.0 for batch in train_loader: loss criterion(model(batch), target) total_loss loss # 错 loss.backward() optimizer.step()第一次跑可能没事因为Python的会创建新Tensor旧Tensor被覆盖后可能被回收。但在某些情况下loss还连着计算图累加操作会让计算图历史被保留着如果再叠加梯度累积、混合精度、延迟反向等场景多个step的计算图引用就容易同时存活着。显存轻微上涨在短训练里看不出来长跑几十个epoch就变成灾难。正确的做法是total_loss loss.item()loss.item()返回的是Python float和计算图、GPU显存完全无关。凡是要把loss、acc这类标量记录到日志或累积统计一律用.item()或loss.detach().cpu().item()。这是一条铁律。5.2 验证循环里漏掉with torch.no_grad()另一个高频问题很多人验证时只想着模型切到eval模式忘了包no_grad()。没有no_grad()的情况下验证前向也会创建计算图虽然不调用backward但中间激活Tensor的生命周期会被拉长。如果验证集还有几个大batch额外显存很可能直接触顶。model.eval() with torch.no_grad(): for batch in val_loader: output model(batch)no_grad()包住验证循环后前向就不会再构建计算图激活值用完即释放这是显存层面最基本的防线。只要no_grad()正确包住计算图的负担就基本消除了。5.3 自定义Dataset或collate_fn里持有模型Tensor数据加载器里的泄漏比较隐蔽因为报错往往不在DataLoader本身。比如__getitem__里用到了某个GPU Tensor作为索引或者collate_fn里有对GPU Tensor的处理逻辑这些Tensor的生命周期可能比预期长得多。还有一种常见情况Dataset里缓存了输入图像或语音波形体量一大系统内存先吃满间接拖累GPU分配。几条实测经验保证Dataset返回的样本都在CPU上GPU Tensor只在进入模型之前创建。DataLoader的num_workers不要无脑开很多尤其在worker内部使用了共享内存、pyarrow等库时容易出现句柄不释放诱发进程资源泄漏最终反映成训练环境整体变慢、显存异常。自定义collate_fn里严禁创建GPU Tensor并保留到了训练循环再统一做设备转移。5.4 模型forward里把中间量存成self属性这个案例我印象很深。有人为了让attention可视化脚本能读到某个中间特征在模型forward里写了self.mid_feat x。这个x是attention输出的Tensor参与了后续计算意味着挂到self上的同时也被计算图引用。虽然下一轮前向会重新赋值表面上看是覆盖但在混合精度、梯度裁剪、延迟反向等情况下旧值可能在某些分支路径上没有被覆盖导致计算图被多个step同时挂着。如果你确实需要中间特征做可视化或剪枝分析正确做法是维护一个可以随时清空的存根存detach().cpu()后的值并在每个训练step开始或epoch结束时主动清空class MyModel(nn.Module): def __init__(self): self._debug_feats {} def forward(self, x): feat self.attention(x) if self.training and self._capture_feat: # 存CPU拷贝避免GPU Tensor挂进计算图 self._debug_feats[attn] feat.detach().cpu() return self.classifier(feat)这样既保住了可视化需要的数据又不会把GPU Tensor变成持久引用。5.5 开启cudnn.benchmark后显存阶梯膨胀最后是环境层面的问题。torch.backends.cudnn.benchmark True在输入尺寸固定时会为每个卷积选择最优算法并缓存workspace。如果你的数据存在尺寸波动比如文本长度不一导致batch内padding后每次shape不同benchmark会频繁尝试不同算法每次重新分配workspace缓存。有些版本的cuDNN在释放旧workspace时机上比较懒导致reserved一路涨。症状是allocated不怎么变但reserved每隔几十步跳一档。这种情况的处理没有统一方案要看你的数据特性模型包含可变尺寸输入时要么关掉benchmark要么设置torch.backends.cudnn.deterministic True要么在数据层面把输入pad到统一尺寸并用mask控制。关掉benchmark后如果显存曲线平了基本就是它的锅。6. 一次真实泄漏的完整复盘从看到日志异常到改完代码验证稳定6.1 症状每个epoch显存固定多300MB那次是在训练一个多模态检索模型文本侧用BERT图像侧用ResNet最后做对比学习loss。训练跑到第4个epoch时我注意到日志里allocated从第一轮的3.8GB变成第二轮的4.1GB、第三轮的4.4GB。每个epoch增加约300MB模型总共才大约5亿参数这个增长明显不合理。更扎心的是跑到第5轮以后batch interval从每步0.9秒涨到1.3秒训练速度肉眼可见地退化。当时我第一反应也是查数据有没有问题但数据尺寸从头到尾是一致的。于是我把memory_summary()打印了出来发现active blocks数量非常多占用集中在几个segment。直觉告诉我这不是缓存池碎片是有Tensor在持续占用。6.2 定位snapshot里反复出现同一Call Stack接着我跑了一遍memory_snapshot()用前面那个聚类脚本统计Top 10分配点。结果一个函数稳稳占据了第一位xxx/multimodal/model.py: line 156: forward_embedding_loss这个函数里只有一个可疑行为它把每个batch的image_embedding和text_embedding都保留下来定期写入一个可视化文件用来观察两个模态的embedding空间是否对齐。代码大致是这样def forward_embedding_loss(image_emb, text_emb, step): self.vis_embeds.append((step, image_emb, text_emb)) # 可疑 ...image_emb是形状(batch_size, 768)的GPU Tensor每个batch保留一份。训练到第10个epoch时list里已经有几千个Tensor显存自然越涨越高。更隐蔽的是这段逻辑放在模型内部训练循环外又有一个回调负责在epoch结束时读取这个list画tSNE图所以它从来没有被清理的时机。6.3 根因为了可视化给模型挂了一个永不释放的缓存这个list里的Tensor每个只有几百KB因为batch是64乘以768个float也才几百KB。但几千个batch累积起来就不是小数目了。而且为了画出整个训练历史的轨迹代码在攒够之后也没有清空只一味追加典型的数据引用失控。修复其实很简单可视化完全不需要GPU Tensor也不需要保存全部历史。我改成每512个batch才采样一份CPU拷贝旧数据在画出图之后立即清空。同时把这个feature cache从模型内部挪到训练脚本外面作为回调参数传递训练step结束时显式clear()。这样模型加载了之后没有额外显存负担训练逻辑也更干净。6.4 修复与验证连续50个epoch的曲线终于平了修复后我用同一套配置从头重新训练了50个epoch。这次观察allocated曲线第一轮结束后3.8GB第50轮结束时仍然在3.8GB附近波动不超过±0.1GB。batch interval也稳定在0.9秒上下没有再出现越跑越慢的情况。这说明判断一个显存泄漏修复是否到位标准不是“某一时刻不涨”而是“连续跑多个epoch后allocated曲线整体持平”。为更保险我还做了一次压力测试临时把验证集batch加大一倍跑20个epoch。这时reserved确实出现了一次跳变但allocated没有同步上升说明缓存池扩容是正常的预期行为不会随着epoch推进演变成泄漏。这是我个人在实际排查中体会最深的一点不要一看到显存涨就以为代码有鬼也不要一看到缓存池扩容就急着改代码把allocated和reserved分开看很多问题会清晰很多。最后再分享一个建议每个训练项目启动时先花十分钟把显存监控埋进日志跑几十个epoch后回看曲线。这笔时间花得比任何排查工具都值因为它能让你在问题刚冒头时就发现它而不是等到OOM中断训练之后才匆匆忙忙去定位。 
长河君 · 2026/10/27 3:41:43 · 阅读 1.2k PyTorch显存泄漏排查指南：从监控到定位的完整方法论
跑PyTorch训练的工程师大多都经历过这种让人血压飙升的时刻早上看显存占用还是10GB中午变成16GB下午直接报CUDA out of memory训练戛然而止。更诡异的是只要你重启训练同一个脚本又能从头开始正常跑。这种“越跑显存越高”的现象就是典型的PyTorch显存泄漏问题。这个问题的麻烦在于它不像普通的程序崩溃那样给你一个明确的错误栈很多时候你只能从缓慢变慢的速度和一格一格涨上去的显存里感受到异常。我这些年排查过不少这类问题从最开始的瞎试empty_cache到后来用内存快照精确定位到具体代码行积累了一套可复用的方法论。这篇文章我会完整讲一遍怎么判断真假泄漏怎么给训练脚本装上显存监控为什么计算图会把显存变黑洞以及三条从浅到深的定位手段。最后还有一个真实案例的全过程复盘帮你把前面所有方法串起来。不管你是在调一个已经跑崩的模型还是想提前给训练脚本打好预防针都能从这篇文章里找到能直接用的东西。1. 显存曲线带着训练一起失控先分清“假性增长”和“真泄漏”1.1 一眼识别“假性增长”先别急着动代码先说个反直觉的结论并非所有显存“增长”都是泄漏。我见过很多同事一发现显存涨了就怀疑是泄漏然后去改模型结构改了大半天发现根本没作用因为问题压根不在模型。要理解这一点得先搞懂PyTorch自己的缓存分配器。GPU显存的申请和释放走的是cudaMalloc和cudaFree这俩操作相对昂贵。如果每个Tensor用完都立刻归还给驱动训练速度会被拖垮。所以PyTorch内部实现了一个caching allocator你申请一块显存等Tensor不再使用后它并不会马上把这块显存还回去而是留在进程自己的缓存池里供后续分配直接复用。于是你打开nvidia-smi看到的进程显存占用其实包含两部分一部分是真正被Tensor占用的另一部分是PyTorch“屯”在缓存池里的。对应到PyTorch的API分别是torch.cuda.memory_allocated()和torch.cuda.memory_reserved()。正常情况下reserved曲线的形态是阶梯跳变。比如训练第一次遇到大的batch、大的卷积核或矩阵乘法时缓存池发现现有的块不够用就重新向驱动申请一大块显存表现在nvidia-smi里就是显存一下子跳上一个台阶。之后如果又遇到更大的需求再跳一档。这种增长不是逐batch爬升的而是在某几个时间点跳到新平台后保持基本平坦。判断“假性增长”的关键就是看增长到底是一次性跳到高位后保持稳定还是每个batch、每个epoch都在持续爬升。前者是缓存扩展后者才可能是泄漏。还有个小技巧如果计算图本身没问题PyTorch后续会复用缓存块训练几十上百个epoch之后显存占用应该是平稳的。提示不要一看到nvidia-smi里显存高位就判定泄漏。先区分allocated和reserved再看多轮epoch的曲线如果allocated在多轮之后稳定不涨说明缓存池正常只是预留偏大。1.2 真泄漏的两种典型形态单调爬升型和阶梯跳变型根据我观察过的实际案例真正的显存泄漏基本可以归成两类形态。第一种是“单调爬升型”。显存每个batch或每个epoch都增长固定大小十几MB、几十MB都有。这种问题通常是一个Tensor或缓冲区被创建后又被某些对象持续引用导致它永远无法被垃圾回收。典型写法就是训练循环里把每个step的中间变量append进一个list而list又挂在模型或训练器上或者自定义loss函数为了可视化保留了一堆中间feature。增长量级相对固定和每次新增Tensor的总字节数成正比很好识别。第二种是“阶梯跳变型”。显存在某个阶段突然上一个大台阶之后回不到原来的水平。常见触发条件是开启了torch.backends.cudnn.benchmark True后模型对多种输入尺寸做auto-tune不断重新分配卷积workspace也可能是验证集里出现了一个尺寸异常大的batch缓存分配器为了喂饱它把reserved整体调高即使验证结束缓存也不会缩回去。这类问题只要不是每个epoch都涨危害相对小一些但如果把reserved误当成allocated来分析很容易冤枉好人。还有一种比较少见的“叠加型”模型在做参数剪枝、动态构图这类操作时每过N个epoch新增一批Tensor引用显存呈周期性台阶增长。这种需要配合日志时间戳和显存曲线做联合分析比前两类都难定位一些。1.3 一个五步自检法五分钟内判断要不要深入排查我的习惯是不管三七二十一先跑一个固定流程用最快速度判断到底是“要不要进入深水区”在训练脚本里每100个step打印torch.cuda.memory_allocated() / 1024**3观察它的走势。单独跑50个step记录起始值和结束值的差。如果差值超过1GB且不是由batch大小变化引起基本可以断定有问题。在验证循环前后各打一次allocated对比差值。如果每次验证结束都涨一块而且缩不回去说明验证环节有引用未释放。把所有自定义日志、可视化回调全部关掉再跑100个step。如果显存停止增长问题出在回调或日志代码里而不是模型主体。在每次backward()和optimizer.step()之后强制torch.cuda.empty_cache()看看有没有改善。注意这一步只是排除法不是修复手段。这五步走完你对问题类型和范围基本心里有数了再进入下一步就不会抓瞎。2. 第一套监控脚手架让显存占用在日志里“显形”2.1 为什么别只依赖nvidia-sminvidia-smi是大多数人排查显存问题的第一反应但我必须说它在这个场景下有天然死角它显示的是整个进程的总体显存占用而且默认几秒刷新一次你很难从它看到增长的Tensor到底是谁分配出来的。更麻烦的是它区分不了“PyTorch缓存池里的预留块”和“真正的活体Tensor”很容易误导判断。PyTorch自带的几个内存统计接口才是排查主力memory_allocated()、memory_reserved()、max_memory_allocated()以及两个深水区武器torch.cuda.memory_summary()和torch.cuda.memory_snapshot()。memory_summary会按分配段、活跃块、保留块给出一份汇总文本让你快速判断显存是被活体Tensor占着还是碎在缓存池里。memory_snapshot则能导出一份包含每个分配事件、地址、大小和调用栈的JSON配合分析脚本直接定位到具体代码行。所以我的建议是别一上来就盯nvidia-smi先基于PyTorch的API在训练脚本里埋好监控让显存增长自动在日志里现形。2.2 一套可复用的训练显存监控装饰器下面这个装饰器是我在多个训练项目里一直在用的模式。不侵入训练逻辑每隔N个step打印一次allocated和reserved顺带记一下耗时import time import torch def cuda_monitor(step_interval10): def decorator(fn): def wrapper(*args, **kwargs): torch.cuda.synchronize() start_time time.time() result fn(*args, **kwargs) torch.cuda.synchronize() elapsed time.time() - start_time if wrapper.counter % step_interval 0: allocated torch.cuda.memory_allocated() / 1024**3 reserved torch.cuda.memory_reserved() / 1024**3 peak torch.cuda.max_memory_allocated() / 1024**3 print(f[step {wrapper.counter:6d}] fallocated {allocated:.3f} GB, freserved {reserved:.3f} GB, fpeak {peak:.3f} GB, ftime {elapsed:.3f}s) torch.cuda.reset_peak_memory_stats() wrapper.counter 1 return result wrapper.counter 0 return wrapper return decorator很多用PyTorch Lightning这类框架的朋友可以把逻辑搬进on_train_batch_start和on_train_batch_end回调思路完全一样。这里有两个细节必须强调。第一打印前一定要torch.cuda.synchronize()。GPU上的kernel是异步执行的如果不做同步显存统计可能反映出“还有一堆kernel没跑完”的假象导致你看不清真实分配趋势。第二打印频率别太高十步一百步一次就行。在大模型训练里日志本身也有IO开销每一步都打印会把单步耗时拉长反过来干扰了你判断“显存泄漏导致速度变慢”的准确度。2.3 异常出现时先记录哪三组指标发现显存异常时不要慌我先记录三组指标指标含义判定泄漏的价值allocated实际占用GPU显存的Tensor字节数短期内快速增长基本就是Tensor泄漏reservedPyTorch向驱动申请的显存总量阶梯跳变说明是缓存池扩容或benchmarkmax_allocated从上次reset以来的峰值分配评估是否接近OOM阈值确认峰值异常举一个很常见的例子allocated从2.1GB涨到2.8GB但reserved一直是4.0GB。你从nvidia-smi看到的进程占用一直是4.0GB根本没变化于是误以为没泄漏。其实PyTorch在缓存池内部不断调整块的划分活体Tensor占用一直在涨只是还没触发新一轮cudaMalloc而已。反过来如果allocated平稳但reserved每隔几百步跳一次说明模型对不同输入shape产生了新的缓存需求这时要检查数据方面有没有batch size波动或padding变化。这三组指标配合时间戳输出到日志之后训练脚本就相当于自带了一台显存示波器后面定位问题会轻松太多。3. autograd的“账本”怎么变成显存黑洞理解计算图留存机制3.1 计算图的生命周期反向传播后到底释放了什么很多人以为显存泄漏就是“变量没删干净”其实不完全是。根源在于PyTorch的autograd机制。前向传播过程中PyTorch会记录一张计算图每个节点保存参与运算的输入、输出以及反向传播需要的梯度函数。训练模式下调用loss.backward()后PyTorch遍历整张图计算每个叶子节点的梯度并把结果写入param.grad。正常来说梯度计算完成后计算图会被释放保存中间激活值的那些Tensor就变成垃圾可以被回收。但有几个例外会延长计算图的寿命。第一你在前向过程中把一个Tensor挂到了模型或某个对象的属性上而这个属性在训练期间没有任何被清理的时机。它相当于给计算图“续了命”下一轮前向又生成新图但旧图的引用一直没断显存自然只进不出。第二自定义loss或回调为了做可视化、统计指标把中间结果保存下来却没有在合适的地方释放或做CPU拷贝同样导致引用累积。第三也是新人最容易踩的loss变量本身。如果loss是标量Tensor在后面某个日志评分函数里被存了下来或者错误地把Tensor连计算图直接丢给add_scalar之类的方法那这个loss Tensor会连同整个计算图一起被长期引用。3.2 empty_cache不是万能的很多人的第一反应是“显存涨了我就调torch.cuda.empty_cache()把它清掉”。我必须泼一盆冷水如果增长的Tensor还被计算图引用着empty_cache一点用都没有。因为它只把缓存池里那些“已释放但还没归还给驱动”的空闲块还给驱动并不会去清理正在使用的Tensor。用一句话概括这里的关系泄漏时增长的是allocated不是缓存池而empty_cache清的是缓存池。所以排查时正确思路是找“谁还在引用这个Tensor”而不是想“怎么把显存清掉”。只有当训练代码本身没有泄漏只是某些临时大块数据导致reserved过高时empty_cache才有压缩显存的效果但这也是治标。真正稳的做法是让分配的Tensor在生命周期结束后不再被任何持久对象引用。3.3 中间变量为什么是泄漏的高发区把原理落到代码上其实就是一个引用生命周期问题。正常的训练循环里局部变量的作用域在循环体内迭代结束就销毁Python自动回收。但如果中间变量被放进了list、被挂到self上、被存进全局字典它就跨过了迭代边界成为持久引用。一个典型的隐性累积例子你想可视化embedding分布写了self.cached_feature x。如果这个x是参与计算图的GPU Tensor每一轮都会被新值覆盖看起来不会涨但如果可视化代码写的是visualize_list.append(x)那就是另一回事了——每一个step的x都不会被覆盖每步新增一份GPU显存list越来越长显存当然一路攀升。处理中间变量的核心原则是不能让任何Tensor对象尤其是参与计算图的Tensor以“append”“缓存列表”“持久属性”的形式跨迭代存活。需要可视化或日志的数据必须显式做detach().cpu()并且最好只保留统计值而不是原始Tensor。4. 三招定位“显存刺客”从统计指标一路挖到Tensor的出生地4.1 第一招memory_summary看整体头部这是最朴素的定位法适合先看整体。在训练循环里插入print(torch.cuda.memory_summary(abbreviatedFalse))输出会告诉你当前显存里有多少个segments、多少个active blocks、多少个inactive blocks。OOM时还能帮你区分是不是碎片化问题。碎片化的典型特征是free块非常多但能匹配请求大块的连续空间太少这跟真正的Tensor泄漏不是一回事。用memory_summary能快速区分。如果文中的active blocks占了绝大部分说明“有人真实占用大块显存不放手”可以继续深挖了。如果active blocks不多但inactive blocks堆了一大堆碎片问题可能偏向缓存池策略或数据shape波动而不是单一的Tensor引用泄漏。4.2 第二招memory_snapshot生成全量分配记录memory_summary只给汇总到不了Tensor级。想看到底是哪个Tensor在增长就得memory_snapshotsnapshot torch.cuda.memory_snapshot() import json with open(mem_snapshot.json, w) as f: json.dump(snapshot, f, indent1)snapshot返回的是一个列表每个segment对象包含address、total_size、allocated_size、active_size、requested_size、stream等字段segment里还有blocks数组。最关键的是每个block往往带有一个frames字段记录了这块显存是哪一行代码、哪个函数分配的。知道这一点以后定位精确到代码行就变成了可能。我自己写过一个简单聚类脚本遍历snapshot里所有active blocks按file_name和line_number聚合求和size输出Top 10占用最大的分配点。下面是一个核心片段可以直接改来用from collections import defaultdict def analyze_snapshot(snapshot): totals defaultdict(lambda: [0, []]) for seg in snapshot: for blk in seg.get(blocks, []): if blk.get(allocated_size, 0) 0 or blk.get(active_size, 0) 0: size blk.get(requested_size, 0) or blk.get(size, 0) frames blk.get(frames, []) if frames: last frames[-1] key f{last[filename]}:{last[line_number]}:{last[function]} else: key unknown totals[key][0] size totals[key][1].append(blk) for key, (total, _) in sorted(totals.items(), keylambda x: -x[1][0])[:10]: print(f{total / 1024**3:.4f} GB {key})需要提醒一句不同PyTorch版本的snapshot字段名略有差异我写这段代码基于2.x你的版本如果不一致以实际输出为准。4.3 第三招二分注释排查法控制变量永远最可靠如果嫌自己写脚本麻烦社区里也有一些现成的可视化工具。PyTorch官方Sphere曾提供基于snapshot的内存可视化notebook第三方项目torch_mem_viz能把snapshot渲染成HTML按内存段和分配栈做交互式查看。这些工具的价值在于你能直观看到哪块内存长时间active、哪个call stack在反复分配。但我必须说实际工程里最快的方法往往不是花哨的可视化而是二分注释排查法。在训练循环里逐步注释掉被怀疑的代码块每注释一块就跑50个step看显存曲线。比如先关掉验证循环显存稳了问题就在验证再关掉自定义回调显存稳了问题就在回调再关掉某个复杂loss分支逐步缩小范围。这种“控制变量”法听起来土但定位速度通常比任何工具都快因为它直接排除了大量无关因素。5. 五种容易养出显存泄漏的写法每个都配了修复版5.1 累计训练指标时忘写.item()这是最经典的坑。很多人写训练循环时会累积loss做日志total_loss 0.0 for batch in train_loader: loss criterion(model(batch), target) total_loss loss # 错 loss.backward() optimizer.step()第一次跑可能没事因为Python的会创建新Tensor旧Tensor被覆盖后可能被回收。但在某些情况下loss还连着计算图累加操作会让计算图历史被保留着如果再叠加梯度累积、混合精度、延迟反向等场景多个step的计算图引用就容易同时存活着。显存轻微上涨在短训练里看不出来长跑几十个epoch就变成灾难。正确的做法是total_loss loss.item()loss.item()返回的是Python float和计算图、GPU显存完全无关。凡是要把loss、acc这类标量记录到日志或累积统计一律用.item()或loss.detach().cpu().item()。这是一条铁律。5.2 验证循环里漏掉with torch.no_grad()另一个高频问题很多人验证时只想着模型切到eval模式忘了包no_grad()。没有no_grad()的情况下验证前向也会创建计算图虽然不调用backward但中间激活Tensor的生命周期会被拉长。如果验证集还有几个大batch额外显存很可能直接触顶。model.eval() with torch.no_grad(): for batch in val_loader: output model(batch)no_grad()包住验证循环后前向就不会再构建计算图激活值用完即释放这是显存层面最基本的防线。只要no_grad()正确包住计算图的负担就基本消除了。5.3 自定义Dataset或collate_fn里持有模型Tensor数据加载器里的泄漏比较隐蔽因为报错往往不在DataLoader本身。比如__getitem__里用到了某个GPU Tensor作为索引或者collate_fn里有对GPU Tensor的处理逻辑这些Tensor的生命周期可能比预期长得多。还有一种常见情况Dataset里缓存了输入图像或语音波形体量一大系统内存先吃满间接拖累GPU分配。几条实测经验保证Dataset返回的样本都在CPU上GPU Tensor只在进入模型之前创建。DataLoader的num_workers不要无脑开很多尤其在worker内部使用了共享内存、pyarrow等库时容易出现句柄不释放诱发进程资源泄漏最终反映成训练环境整体变慢、显存异常。自定义collate_fn里严禁创建GPU Tensor并保留到了训练循环再统一做设备转移。5.4 模型forward里把中间量存成self属性这个案例我印象很深。有人为了让attention可视化脚本能读到某个中间特征在模型forward里写了self.mid_feat x。这个x是attention输出的Tensor参与了后续计算意味着挂到self上的同时也被计算图引用。虽然下一轮前向会重新赋值表面上看是覆盖但在混合精度、梯度裁剪、延迟反向等情况下旧值可能在某些分支路径上没有被覆盖导致计算图被多个step同时挂着。如果你确实需要中间特征做可视化或剪枝分析正确做法是维护一个可以随时清空的存根存detach().cpu()后的值并在每个训练step开始或epoch结束时主动清空class MyModel(nn.Module): def __init__(self): self._debug_feats {} def forward(self, x): feat self.attention(x) if self.training and self._capture_feat: # 存CPU拷贝避免GPU Tensor挂进计算图 self._debug_feats[attn] feat.detach().cpu() return self.classifier(feat)这样既保住了可视化需要的数据又不会把GPU Tensor变成持久引用。5.5 开启cudnn.benchmark后显存阶梯膨胀最后是环境层面的问题。torch.backends.cudnn.benchmark True在输入尺寸固定时会为每个卷积选择最优算法并缓存workspace。如果你的数据存在尺寸波动比如文本长度不一导致batch内padding后每次shape不同benchmark会频繁尝试不同算法每次重新分配workspace缓存。有些版本的cuDNN在释放旧workspace时机上比较懒导致reserved一路涨。症状是allocated不怎么变但reserved每隔几十步跳一档。这种情况的处理没有统一方案要看你的数据特性模型包含可变尺寸输入时要么关掉benchmark要么设置torch.backends.cudnn.deterministic True要么在数据层面把输入pad到统一尺寸并用mask控制。关掉benchmark后如果显存曲线平了基本就是它的锅。6. 一次真实泄漏的完整复盘从看到日志异常到改完代码验证稳定6.1 症状每个epoch显存固定多300MB那次是在训练一个多模态检索模型文本侧用BERT图像侧用ResNet最后做对比学习loss。训练跑到第4个epoch时我注意到日志里allocated从第一轮的3.8GB变成第二轮的4.1GB、第三轮的4.4GB。每个epoch增加约300MB模型总共才大约5亿参数这个增长明显不合理。更扎心的是跑到第5轮以后batch interval从每步0.9秒涨到1.3秒训练速度肉眼可见地退化。当时我第一反应也是查数据有没有问题但数据尺寸从头到尾是一致的。于是我把memory_summary()打印了出来发现active blocks数量非常多占用集中在几个segment。直觉告诉我这不是缓存池碎片是有Tensor在持续占用。6.2 定位snapshot里反复出现同一Call Stack接着我跑了一遍memory_snapshot()用前面那个聚类脚本统计Top 10分配点。结果一个函数稳稳占据了第一位xxx/multimodal/model.py: line 156: forward_embedding_loss这个函数里只有一个可疑行为它把每个batch的image_embedding和text_embedding都保留下来定期写入一个可视化文件用来观察两个模态的embedding空间是否对齐。代码大致是这样def forward_embedding_loss(image_emb, text_emb, step): self.vis_embeds.append((step, image_emb, text_emb)) # 可疑 ...image_emb是形状(batch_size, 768)的GPU Tensor每个batch保留一份。训练到第10个epoch时list里已经有几千个Tensor显存自然越涨越高。更隐蔽的是这段逻辑放在模型内部训练循环外又有一个回调负责在epoch结束时读取这个list画tSNE图所以它从来没有被清理的时机。6.3 根因为了可视化给模型挂了一个永不释放的缓存这个list里的Tensor每个只有几百KB因为batch是64乘以768个float也才几百KB。但几千个batch累积起来就不是小数目了。而且为了画出整个训练历史的轨迹代码在攒够之后也没有清空只一味追加典型的数据引用失控。修复其实很简单可视化完全不需要GPU Tensor也不需要保存全部历史。我改成每512个batch才采样一份CPU拷贝旧数据在画出图之后立即清空。同时把这个feature cache从模型内部挪到训练脚本外面作为回调参数传递训练step结束时显式clear()。这样模型加载了之后没有额外显存负担训练逻辑也更干净。6.4 修复与验证连续50个epoch的曲线终于平了修复后我用同一套配置从头重新训练了50个epoch。这次观察allocated曲线第一轮结束后3.8GB第50轮结束时仍然在3.8GB附近波动不超过±0.1GB。batch interval也稳定在0.9秒上下没有再出现越跑越慢的情况。这说明判断一个显存泄漏修复是否到位标准不是“某一时刻不涨”而是“连续跑多个epoch后allocated曲线整体持平”。为更保险我还做了一次压力测试临时把验证集batch加大一倍跑20个epoch。这时reserved确实出现了一次跳变但allocated没有同步上升说明缓存池扩容是正常的预期行为不会随着epoch推进演变成泄漏。这是我个人在实际排查中体会最深的一点不要一看到显存涨就以为代码有鬼也不要一看到缓存池扩容就急着改代码把allocated和reserved分开看很多问题会清晰很多。最后再分享一个建议每个训练项目启动时先花十分钟把显存监控埋进日志跑几十个epoch后回看曲线。这笔时间花得比任何排查工具都值因为它能让你在问题刚冒头时就发现它而不是等到OOM中断训练之后才匆匆忙忙去定位。