1. 训练显存曲线持续爬升先搞清楚它到底是不是泄漏显存泄漏这个词在深度学习圈子里被用得很随意很多人一看到nvidia-smi里的显存占用比刚启动时高就断定自己遇到了泄漏。实际上训练过程中显存曲线缓慢上升有相当一部分情况属于正常现象盲目去改代码反而会把训练逻辑搞乱。我先把判断标准讲清楚再谈怎么定位。1.1 三种显存变高的本质区别第一种是碎片化导致的假性增长。PyTorch 的 CUDA 缓存分配器CUDACachingAllocator会向驱动申请显存块并缓存起来即使张量被释放缓存也不一定立刻还给驱动。所以nvidia-smi看到的 reserved 显存往往高于实际 allocated 显存。这种情况在训练初期快速上升之后趋于平稳属于正常。第二种是真实的内存增长但非泄漏。比如你在训练循环里累积了 loss 列表、保存了每个 step 的输出用于后续分析、或者验证集评估时没有清空中间变量。这类增长有明确的上界跟训练步数成线性关系但不会无限膨胀。第三种才是真正的显存泄漏。特征是显存占用随 step 单调递增跑几百上千步后必然 OOM而且即使你把 batch size 调到很小最终还是会爆。这种泄漏的根源通常是计算图被意外持有、张量被全局变量引用、或者 CUDA 事件/流没有正确释放。判断方法很简单在训练循环里每隔 N 步打印一次import torch def log_mem(step): allocated torch.cuda.memory_allocated() / 1024**2 reserved torch.cuda.memory_reserved() / 1024**2 max_allocated torch.cuda.max_memory_allocated() / 1024**2 print(fstep {step}: allocated{allocated:.1f}MB, freserved{reserved:.1f}MB, max{max_allocated:.1f}MB)如果allocated持续单调上升且不收敛基本可以确认是泄漏如果只有reserved上升而allocated稳定那只是缓存策略问题不用慌。1.2 为什么越跑越高在 Transformer 类模型里尤其常见我做过不少 NLP 和视觉 Transformer 的训练调优发现这类模型特别容易出现显存爬升。原因在于它们的计算图结构复杂中间激活值多而且很多实现里会无意中保留对计算图的引用。比如你把loss存进一个 list 用于画曲线这个loss是带grad_fn的张量它持有整条反向图图里的所有中间激活都不会释放。跑几千步下来显存自然爆炸。另一个高频场景是带 KV Cache 的推理或训练。如果你在自回归生成时不断把新的 key/value 拼接到缓存里而缓存又是以 Python list 形式持有张量那每一步都会新增一份显存占用。这类问题在 mask2former、各类检测和分割模型的训练脚本里也时有出现因为它们的 loss 计算涉及多个辅助头和匹配过程中间变量多。提示判断泄漏时务必区分memory_allocated实际张量占用和memory_reserved缓存池占用。前者才是泄漏的信号后者可以通过torch.cuda.empty_cache()观察是否回落。2. 计算图被谁攥住了最常见的四类泄漏源头定位显存泄漏核心思路是找到谁还在引用那些本该释放的张量。Python 的引用计数和 GC 机制在 CUDA 张量上有个坑张量本身被释放了但它关联的 CUDA 内存要等计算图引用计数归零才回收。下面这四类源头覆盖了我排查过的八成以上案例。2.1 把带梯度的 loss 或输出存进容器这是最经典的泄漏。看这段代码losses [] for batch in dataloader: out model(batch) loss criterion(out, batch[label]) losses.append(loss) # 灾难 loss.backward() optimizer.step() optimizer.zero_grad()loss是一个 requires_grad 的张量它持有整个前向计算图。把它 append 进 list等于把每一步的计算图都钉在内存里。正确做法是存标量losses.append(loss.item()) # 或者 loss.detach()loss.item()会取出 Python float彻底断开与计算图的联系。loss.detach()也行但 detach 后的张量仍占一份显存只是不带图长期累积同样会涨。所以画曲线用.item()最稳妥。2.2 验证/评估阶段忘了 no_grad 或 inference_mode训练循环里插验证是常规操作但很多人写验证函数时忘了加torch.no_grad()def validate(model, loader): model.eval() total 0 for batch in loader: out model(batch) # 这里在构建计算图 total (out.argmax(1) batch[label]).sum() return total没有no_grad每次验证都会构建计算图而且这些图因为out被用于计算 accuracy 而无法释放。跑几个 epoch 后显存就上去了。正确写法torch.no_grad() def validate(model, loader): ...PyTorch 2.x 之后更推荐torch.inference_mode()它比no_grad更彻底会关闭版本计数等额外开销显存和速度都更优。但注意inference_mode下创建的张量不能再参与后续需要梯度的计算验证场景正好合适。2.3 全局缓存、日志钩子和回调持有张量很多训练框架自己写的或第三方的会注册 hook 来记录中间激活用于可视化或调试。如果 hook 里把激活值存进全局 dict 而从不清理那就是慢性泄漏。我见过一个案例某项目用register_forward_hook收集每层输出做注意力可视化结果 hook 把每步的激活都存进 list训练到 200 步就 OOM。排查这类问题可以用gc模块配合torch.cuda的统计import gc import torch # 训练若干步后 gc.collect() torch.cuda.empty_cache() print(torch.cuda.memory_allocated() / 1024**2)如果gc.collect()之后 allocated 明显下降说明有 Python 对象在循环引用被 GC 回收了如果没变化说明是 CUDA 层面的引用没断得往计算图方向查。2.4 DataLoader 的 worker 与 pin_memory 陷阱多进程 DataLoader 本身不直接导致显存泄漏但pin_memoryTrue配合某些自定义 collate 函数时如果 collate 里创建了 CUDA 张量这些张量会在 worker 进程里驻留。更隐蔽的是如果你在 Dataset 的__getitem__里做了 GPU 上的预处理比如用 GPU 做数据增强每个 worker 都会占一份显存而且 worker 重启时旧显存不一定及时释放。我的建议是数据预处理尽量放在 CPU 或独立的 GPU 预处理管线里Dataset 只返回 CPU 张量让主进程的pin_memory机制去搬运。这样显存占用可控也不会因为 worker 数量变化而波动。泄漏源头典型症状快速验证方法loss 存 listallocated 随 step 线性涨改成 .item() 后观察验证缺 no_grad每个 epoch 后跳一截加装饰器后对比hook 持有激活与 hook 注册数量相关禁用 hook 后对比worker 内 GPU 操作与 num_workers 相关设 num_workers0 对比3. 用工具把泄漏点钉死从 memory_summary 到 profiler光靠肉眼看代码效率太低。PyTorch 自带了几套工具能把显存分配追溯到具体代码行。我通常按先粗后细的顺序来先看整体分布再定位到具体张量最后用 profiler 抓分配栈。3.1 memory_summary 与 memory_snapshot 的实战用法torch.cuda.memory_summary()会打印一份详细的显存报告包括按张量大小分桶的分配情况。训练卡住时打印一次能快速看出是大张量堆积还是小张量碎片print(torch.cuda.memory_summary(deviceNone, abbreviatedFalse))更强大的是memory_snapshot它记录每个分配块的调用栈# 在怀疑泄漏的 step 前后各拍一次 torch.cuda.memory._record_memory_history(max_entries100000) # ... 训练若干步 ... snapshot torch.cuda.memory._snapshot() torch.cuda.memory._dump_snapshot(mem_snapshot.pickle) torch.cuda.memory._record_memory_history(enabledNone)拿到的 pickle 文件可以用 PyTorch 官方提供的可视化工具打开https://pytorch.org/memory_viz它会画出显存随时间的变化曲线并且能点开每个块看分配栈。我靠这个工具定位过一个很隐蔽的问题某个自定义 autograd Function 的 backward 里创建了临时张量却没释放导致每步泄漏几十 MB。3.2 用 profiler 抓谁在分配torch.profiler不仅能看时间还能看显存。开启profile_memoryTrue后它能告诉你每个算子分配和释放了多少显存from torch.profiler import profile, ProfilerActivity with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], profile_memoryTrue, record_shapesTrue, ) as prof: for i, batch in enumerate(dataloader): train_step(batch) if i 20: break print(prof.key_averages().table(sort_byself_cuda_memory_usage, row_limit20))输出里self_cuda_memory_usage这一列能直接告诉你哪个算子占的显存最多。如果某个算子的显存占用随 step 累积而不释放那它就是嫌疑对象。注意 profiler 本身有开销别在正式训练时一直开着定位阶段跑几十步就够。3.3 一个真实的排查链路复盘说个我去年遇到的案例。某检测模型训练batch size 设为 4 时能跑设为 8 时跑 300 步左右 OOM但显存曲线在前 100 步是平的之后突然开始爬。这种延迟泄漏很迷惑人。第一步我用memory_allocated打点确认确实是 allocated 在涨不是 reserved。第二步开memory_snapshot发现泄漏的张量形状是[8, 256, 56, 56]明显是某个特征图。第三步用 profiler 看分配栈定位到是模型里一个辅助损失分支的输出。第四步读代码发现这个分支在特定条件下正样本数超过阈值会走一条额外的计算路径而这条路径的输出被存进了一个模块级的 list 用于后续分析但那个 list 从来没被清理。修复很简单把那个 list 改成只存统计量或者用完即清。改完之后batch size 开到 16 都稳。这个案例的教训是延迟泄漏往往和数据的条件分支有关前 100 步没触发那条路径所以曲线是平的。注意memory_snapshot的max_entries别设太大否则记录本身会占大量内存。定位阶段设 10 万条足够正式训练时务必关掉。4. 修复之后怎么验证别让看起来好了骗了你找到泄漏点、改完代码不代表问题就解决了。我见过太多改完跑了一百步没爆就以为好了的情况结果正式训练跑到几千步又 OOM。验证环节必须做扎实否则前面白干。4.1 用长跑测试确认曲线收敛修复后的第一件事是跑一个足够长的测试。我的经验是至少跑 2000 步并且把memory_allocated每 50 步打一次点画成曲线。健康的曲线应该是前几十步快速上升到工作集大小之后在一个区间内小幅波动整体斜率接近零。如果曲线还在缓慢上升哪怕斜率很小也要警惕。因为训练动辄几万步小斜率累积起来照样爆。这时候可以算一下假设每步涨 0.1MB一万步就是 1GB对于 24GB 的卡来说可能就是压垮骆驼的最后一根稻草。import matplotlib.pyplot as plt steps, mems [], [] for step, batch in enumerate(dataloader): train_step(batch) if step % 50 0: steps.append(step) mems.append(torch.cuda.memory_allocated() / 1024**2) if step 2000: break plt.plot(steps, mems) plt.xlabel(step) plt.ylabel(allocated MB) plt.savefig(mem_curve.png)4.2 对比不同 batch size 和序列长度下的表现泄漏有时是条件触发的换个配置就不出现。所以验证时要把关键维度都扫一遍batch size 从 1 到最大能跑的序列长度或图像分辨率取几个档位。如果某个配置下曲线不收敛说明修复不彻底或者还有第二个泄漏点。我一般会做一个简单的矩阵测试配置batch1batch4batch8短序列跑 2000 步跑 2000 步跑 2000 步长序列跑 2000 步跑 2000 步跑 2000 步每个格子记录最终 allocated 和是否 OOM。全部通过才算修复完成。这个测试跑下来大概一两个小时但能省掉后面几天的反复调试。4.3 把显存监控加进训练脚本修复完成后我建议把显存监控作为常驻功能加进训练脚本而不是临时打点。这样以后换模型、换数据、换卡都能第一时间发现异常。实现很简单写个回调或者直接在训练循环里加class MemMonitor: def __init__(self, log_every100): self.log_every log_every self.baseline None def step(self, step): if step % self.log_every ! 0: return cur torch.cuda.memory_allocated() / 1024**2 if self.baseline is None: self.baseline cur growth cur - self.baseline if growth 500: # 超过 500MB 就告警 print(f[WARN] step {step}: mem grew {growth:.1f}MB since baseline)这个监控不占什么开销但能在泄漏刚冒头时就发出信号比等到 OOM 再回头查要省事得多。5. 几个容易忽略的边角场景与长期习惯前面讲的都是主干问题但实际项目里还有一些边角场景单独拎出来说说。这些坑不一定每次都遇到但遇到了很难查提前知道能省不少时间。5.1 混合精度训练里的 GradScaler 状态AMP自动混合精度训练时GradScaler会维护一些内部状态。正常情况下它不会泄漏但如果你在训练中途重置 scaler 或者手动操作它的状态可能留下悬空引用。另外autocast上下文里创建的张量如果被带出上下文长期持有也可能出问题。我的习惯是AMP 相关的变量都限制在训练步的局部作用域内不往外传。5.2 分布式训练中的显存不均衡用 DDP 或多卡训练时各卡显存占用可能不一致。如果某张卡先 OOM而其他卡还有余量问题往往出在数据分配不均或某张卡上的额外计算比如主卡的日志汇总。排查时先看各卡的memory_allocated找出异常的那张再单独分析它的计算路径。有时候是find_unused_parametersTrue导致的额外图保留关掉它能省不少显存。5.3 长期习惯把显存卫生写进代码规范最后分享几个我团队里推行的习惯能从根本上减少泄漏。第一训练循环里禁止直接 append 张量要存就存.item()或.detach().cpu()。第二所有验证和推理函数强制加torch.no_grad()或torch.inference_mode()用 lint 规则检查。第三任何模块级的 list/dict 缓存必须有明确的清理时机最好用weakref或者定期清空。第四新模型接入训练框架时先跑一个 500 步的显存稳定性测试通过了再上大规模训练。这些习惯看起来琐碎但坚持下来显存泄漏这类问题基本能在萌芽阶段被拦住。我在实际项目里最大的体会是显存问题从来不是调参能解决的它本质上是代码正确性问题得用工程手段去防而不是靠运气去躲。
1. 训练显存曲线持续爬升先搞清楚它到底是不是泄漏显存泄漏这个词在深度学习圈子里被用得很随意很多人一看到nvidia-smi里的显存占用比刚启动时高就断定自己遇到了泄漏。实际上训练过程中显存曲线缓慢上升有相当一部分情况属于正常现象盲目去改代码反而会把训练逻辑搞乱。我先把判断标准讲清楚再谈怎么定位。1.1 三种显存变高的本质区别第一种是碎片化导致的假性增长。PyTorch 的 CUDA 缓存分配器CUDACachingAllocator会向驱动申请显存块并缓存起来即使张量被释放缓存也不一定立刻还给驱动。所以nvidia-smi看到的 reserved 显存往往高于实际 allocated 显存。这种情况在训练初期快速上升之后趋于平稳属于正常。第二种是真实的内存增长但非泄漏。比如你在训练循环里累积了 loss 列表、保存了每个 step 的输出用于后续分析、或者验证集评估时没有清空中间变量。这类增长有明确的上界跟训练步数成线性关系但不会无限膨胀。第三种才是真正的显存泄漏。特征是显存占用随 step 单调递增跑几百上千步后必然 OOM而且即使你把 batch size 调到很小最终还是会爆。这种泄漏的根源通常是计算图被意外持有、张量被全局变量引用、或者 CUDA 事件/流没有正确释放。判断方法很简单在训练循环里每隔 N 步打印一次import torch def log_mem(step): allocated torch.cuda.memory_allocated() / 1024**2 reserved torch.cuda.memory_reserved() / 1024**2 max_allocated torch.cuda.max_memory_allocated() / 1024**2 print(fstep {step}: allocated{allocated:.1f}MB, freserved{reserved:.1f}MB, max{max_allocated:.1f}MB)如果allocated持续单调上升且不收敛基本可以确认是泄漏如果只有reserved上升而allocated稳定那只是缓存策略问题不用慌。1.2 为什么越跑越高在 Transformer 类模型里尤其常见我做过不少 NLP 和视觉 Transformer 的训练调优发现这类模型特别容易出现显存爬升。原因在于它们的计算图结构复杂中间激活值多而且很多实现里会无意中保留对计算图的引用。比如你把loss存进一个 list 用于画曲线这个loss是带grad_fn的张量它持有整条反向图图里的所有中间激活都不会释放。跑几千步下来显存自然爆炸。另一个高频场景是带 KV Cache 的推理或训练。如果你在自回归生成时不断把新的 key/value 拼接到缓存里而缓存又是以 Python list 形式持有张量那每一步都会新增一份显存占用。这类问题在 mask2former、各类检测和分割模型的训练脚本里也时有出现因为它们的 loss 计算涉及多个辅助头和匹配过程中间变量多。提示判断泄漏时务必区分memory_allocated实际张量占用和memory_reserved缓存池占用。前者才是泄漏的信号后者可以通过torch.cuda.empty_cache()观察是否回落。2. 计算图被谁攥住了最常见的四类泄漏源头定位显存泄漏核心思路是找到谁还在引用那些本该释放的张量。Python 的引用计数和 GC 机制在 CUDA 张量上有个坑张量本身被释放了但它关联的 CUDA 内存要等计算图引用计数归零才回收。下面这四类源头覆盖了我排查过的八成以上案例。2.1 把带梯度的 loss 或输出存进容器这是最经典的泄漏。看这段代码losses [] for batch in dataloader: out model(batch) loss criterion(out, batch[label]) losses.append(loss) # 灾难 loss.backward() optimizer.step() optimizer.zero_grad()loss是一个 requires_grad 的张量它持有整个前向计算图。把它 append 进 list等于把每一步的计算图都钉在内存里。正确做法是存标量losses.append(loss.item()) # 或者 loss.detach()loss.item()会取出 Python float彻底断开与计算图的联系。loss.detach()也行但 detach 后的张量仍占一份显存只是不带图长期累积同样会涨。所以画曲线用.item()最稳妥。2.2 验证/评估阶段忘了 no_grad 或 inference_mode训练循环里插验证是常规操作但很多人写验证函数时忘了加torch.no_grad()def validate(model, loader): model.eval() total 0 for batch in loader: out model(batch) # 这里在构建计算图 total (out.argmax(1) batch[label]).sum() return total没有no_grad每次验证都会构建计算图而且这些图因为out被用于计算 accuracy 而无法释放。跑几个 epoch 后显存就上去了。正确写法torch.no_grad() def validate(model, loader): ...PyTorch 2.x 之后更推荐torch.inference_mode()它比no_grad更彻底会关闭版本计数等额外开销显存和速度都更优。但注意inference_mode下创建的张量不能再参与后续需要梯度的计算验证场景正好合适。2.3 全局缓存、日志钩子和回调持有张量很多训练框架自己写的或第三方的会注册 hook 来记录中间激活用于可视化或调试。如果 hook 里把激活值存进全局 dict 而从不清理那就是慢性泄漏。我见过一个案例某项目用register_forward_hook收集每层输出做注意力可视化结果 hook 把每步的激活都存进 list训练到 200 步就 OOM。排查这类问题可以用gc模块配合torch.cuda的统计import gc import torch # 训练若干步后 gc.collect() torch.cuda.empty_cache() print(torch.cuda.memory_allocated() / 1024**2)如果gc.collect()之后 allocated 明显下降说明有 Python 对象在循环引用被 GC 回收了如果没变化说明是 CUDA 层面的引用没断得往计算图方向查。2.4 DataLoader 的 worker 与 pin_memory 陷阱多进程 DataLoader 本身不直接导致显存泄漏但pin_memoryTrue配合某些自定义 collate 函数时如果 collate 里创建了 CUDA 张量这些张量会在 worker 进程里驻留。更隐蔽的是如果你在 Dataset 的__getitem__里做了 GPU 上的预处理比如用 GPU 做数据增强每个 worker 都会占一份显存而且 worker 重启时旧显存不一定及时释放。我的建议是数据预处理尽量放在 CPU 或独立的 GPU 预处理管线里Dataset 只返回 CPU 张量让主进程的pin_memory机制去搬运。这样显存占用可控也不会因为 worker 数量变化而波动。泄漏源头典型症状快速验证方法loss 存 listallocated 随 step 线性涨改成 .item() 后观察验证缺 no_grad每个 epoch 后跳一截加装饰器后对比hook 持有激活与 hook 注册数量相关禁用 hook 后对比worker 内 GPU 操作与 num_workers 相关设 num_workers0 对比3. 用工具把泄漏点钉死从 memory_summary 到 profiler光靠肉眼看代码效率太低。PyTorch 自带了几套工具能把显存分配追溯到具体代码行。我通常按先粗后细的顺序来先看整体分布再定位到具体张量最后用 profiler 抓分配栈。3.1 memory_summary 与 memory_snapshot 的实战用法torch.cuda.memory_summary()会打印一份详细的显存报告包括按张量大小分桶的分配情况。训练卡住时打印一次能快速看出是大张量堆积还是小张量碎片print(torch.cuda.memory_summary(deviceNone, abbreviatedFalse))更强大的是memory_snapshot它记录每个分配块的调用栈# 在怀疑泄漏的 step 前后各拍一次 torch.cuda.memory._record_memory_history(max_entries100000) # ... 训练若干步 ... snapshot torch.cuda.memory._snapshot() torch.cuda.memory._dump_snapshot(mem_snapshot.pickle) torch.cuda.memory._record_memory_history(enabledNone)拿到的 pickle 文件可以用 PyTorch 官方提供的可视化工具打开https://pytorch.org/memory_viz它会画出显存随时间的变化曲线并且能点开每个块看分配栈。我靠这个工具定位过一个很隐蔽的问题某个自定义 autograd Function 的 backward 里创建了临时张量却没释放导致每步泄漏几十 MB。3.2 用 profiler 抓谁在分配torch.profiler不仅能看时间还能看显存。开启profile_memoryTrue后它能告诉你每个算子分配和释放了多少显存from torch.profiler import profile, ProfilerActivity with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], profile_memoryTrue, record_shapesTrue, ) as prof: for i, batch in enumerate(dataloader): train_step(batch) if i 20: break print(prof.key_averages().table(sort_byself_cuda_memory_usage, row_limit20))输出里self_cuda_memory_usage这一列能直接告诉你哪个算子占的显存最多。如果某个算子的显存占用随 step 累积而不释放那它就是嫌疑对象。注意 profiler 本身有开销别在正式训练时一直开着定位阶段跑几十步就够。3.3 一个真实的排查链路复盘说个我去年遇到的案例。某检测模型训练batch size 设为 4 时能跑设为 8 时跑 300 步左右 OOM但显存曲线在前 100 步是平的之后突然开始爬。这种延迟泄漏很迷惑人。第一步我用memory_allocated打点确认确实是 allocated 在涨不是 reserved。第二步开memory_snapshot发现泄漏的张量形状是[8, 256, 56, 56]明显是某个特征图。第三步用 profiler 看分配栈定位到是模型里一个辅助损失分支的输出。第四步读代码发现这个分支在特定条件下正样本数超过阈值会走一条额外的计算路径而这条路径的输出被存进了一个模块级的 list 用于后续分析但那个 list 从来没被清理。修复很简单把那个 list 改成只存统计量或者用完即清。改完之后batch size 开到 16 都稳。这个案例的教训是延迟泄漏往往和数据的条件分支有关前 100 步没触发那条路径所以曲线是平的。注意memory_snapshot的max_entries别设太大否则记录本身会占大量内存。定位阶段设 10 万条足够正式训练时务必关掉。4. 修复之后怎么验证别让看起来好了骗了你找到泄漏点、改完代码不代表问题就解决了。我见过太多改完跑了一百步没爆就以为好了的情况结果正式训练跑到几千步又 OOM。验证环节必须做扎实否则前面白干。4.1 用长跑测试确认曲线收敛修复后的第一件事是跑一个足够长的测试。我的经验是至少跑 2000 步并且把memory_allocated每 50 步打一次点画成曲线。健康的曲线应该是前几十步快速上升到工作集大小之后在一个区间内小幅波动整体斜率接近零。如果曲线还在缓慢上升哪怕斜率很小也要警惕。因为训练动辄几万步小斜率累积起来照样爆。这时候可以算一下假设每步涨 0.1MB一万步就是 1GB对于 24GB 的卡来说可能就是压垮骆驼的最后一根稻草。import matplotlib.pyplot as plt steps, mems [], [] for step, batch in enumerate(dataloader): train_step(batch) if step % 50 0: steps.append(step) mems.append(torch.cuda.memory_allocated() / 1024**2) if step 2000: break plt.plot(steps, mems) plt.xlabel(step) plt.ylabel(allocated MB) plt.savefig(mem_curve.png)4.2 对比不同 batch size 和序列长度下的表现泄漏有时是条件触发的换个配置就不出现。所以验证时要把关键维度都扫一遍batch size 从 1 到最大能跑的序列长度或图像分辨率取几个档位。如果某个配置下曲线不收敛说明修复不彻底或者还有第二个泄漏点。我一般会做一个简单的矩阵测试配置batch1batch4batch8短序列跑 2000 步跑 2000 步跑 2000 步长序列跑 2000 步跑 2000 步跑 2000 步每个格子记录最终 allocated 和是否 OOM。全部通过才算修复完成。这个测试跑下来大概一两个小时但能省掉后面几天的反复调试。4.3 把显存监控加进训练脚本修复完成后我建议把显存监控作为常驻功能加进训练脚本而不是临时打点。这样以后换模型、换数据、换卡都能第一时间发现异常。实现很简单写个回调或者直接在训练循环里加class MemMonitor: def __init__(self, log_every100): self.log_every log_every self.baseline None def step(self, step): if step % self.log_every ! 0: return cur torch.cuda.memory_allocated() / 1024**2 if self.baseline is None: self.baseline cur growth cur - self.baseline if growth 500: # 超过 500MB 就告警 print(f[WARN] step {step}: mem grew {growth:.1f}MB since baseline)这个监控不占什么开销但能在泄漏刚冒头时就发出信号比等到 OOM 再回头查要省事得多。5. 几个容易忽略的边角场景与长期习惯前面讲的都是主干问题但实际项目里还有一些边角场景单独拎出来说说。这些坑不一定每次都遇到但遇到了很难查提前知道能省不少时间。5.1 混合精度训练里的 GradScaler 状态AMP自动混合精度训练时GradScaler会维护一些内部状态。正常情况下它不会泄漏但如果你在训练中途重置 scaler 或者手动操作它的状态可能留下悬空引用。另外autocast上下文里创建的张量如果被带出上下文长期持有也可能出问题。我的习惯是AMP 相关的变量都限制在训练步的局部作用域内不往外传。5.2 分布式训练中的显存不均衡用 DDP 或多卡训练时各卡显存占用可能不一致。如果某张卡先 OOM而其他卡还有余量问题往往出在数据分配不均或某张卡上的额外计算比如主卡的日志汇总。排查时先看各卡的memory_allocated找出异常的那张再单独分析它的计算路径。有时候是find_unused_parametersTrue导致的额外图保留关掉它能省不少显存。5.3 长期习惯把显存卫生写进代码规范最后分享几个我团队里推行的习惯能从根本上减少泄漏。第一训练循环里禁止直接 append 张量要存就存.item()或.detach().cpu()。第二所有验证和推理函数强制加torch.no_grad()或torch.inference_mode()用 lint 规则检查。第三任何模块级的 list/dict 缓存必须有明确的清理时机最好用weakref或者定期清空。第四新模型接入训练框架时先跑一个 500 步的显存稳定性测试通过了再上大规模训练。这些习惯看起来琐碎但坚持下来显存泄漏这类问题基本能在萌芽阶段被拦住。我在实际项目里最大的体会是显存问题从来不是调参能解决的它本质上是代码正确性问题得用工程手段去防而不是靠运气去躲。
