训练跑着跑着显存一路往上飙三个小时后直接 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, ascendingFalse) 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_noneTrue)。set_to_noneTrue会把梯度的.grad属性直接置为 None而不是调用.zero_()把张量清零。前者会真正释放掉梯度张量占用的显存后者只是把数值清零但空间还留着。在长训练场景下只要每个 step 都执行set_to_noneTrue梯度累积表剧无法形成梯度缓存池膨胀。optimizer.zero_grad(set_to_noneTrue)我实测过在 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_noneTrue)这个写法没问题因为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 的注意力矩阵都被保存在类实例上直到下一次覆盖才被释放。如果序列长度是 1000head 数是 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 显存问题排查速查表症状可能原因快速验证方法修复方向allocated 随 step 稳定净增循环内持有张量或计算图打印 memory_allocated 曲线看是否严格爬升检查 list/dict 中的 tensor改用 item() 或 detach().cpu()显存曲线锯齿状上涨缓存池碎片化观察 memory_reserved 是否大幅高过 allocated周期性调用torch.cuda.empty_cache()但不能根治重点减少尖峰张量生成验证阶段显存猛增验证循环收集 logits去掉收集列表只计算指标在循环内直接算 metric或 preds.detach().cpu()多卡后单步显存翻倍buffer 被 expand 后参与通信对比单卡与双卡的项数避免 buffer 参与 expand或改为普通 tensor混合精度训练后期 NaN 且显存涨梯度累积未正确 unscale_打印 grad 的 scale 和 norm在累积分支调用 unscale_重启后依然显存高非 PyTorch 进程占用用 nvidia-smi 查看所有进程检查其他进程和 CUDA context 释放我还要提醒一个容易被忽略的点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 标红搞得心力交瘁。
