2、显存泄漏的检测工具：五把“手术刀”

显存泄漏这东西，就像程序里的“慢性失血”。

一开始没啥感觉，跑着跑着，突然OOM了。我见过太多团队，模型没崩，人先崩了。

所以，咱们得学会用工具。今天这五个工具，是我日常排查泄漏的“五把手术刀”。

它们各有侧重，组合起来用，基本能定位99%的问题。

2.1 nvidia-smi 实时监控：最直观的“血压计”

这是最基础的工具。说白了，就是看显存有没有“只涨不跌”。

我个人习惯，在训练脚本里加个定时任务，每5秒输出一次显存占用。

# 终端命令
watch -n 5 nvidia-smi

为什么用它？因为它快，零侵入。

但注意，它只能看到进程级别的总显存。如果泄漏发生在子模块里，它只能告诉你“出事了”，但说不出“哪里出事了”。

2.2 torch.cuda.memory_summary()：PyTorch 的“体检报告”

这个函数，我愿称之为“救命稻草”。

它会把当前GPU上所有张量的分配情况，打印得清清楚楚。

import torch

# 在怀疑泄漏的地方插入
print(torch.cuda.memory_summary())

输出里会告诉你：哪个张量占了多少显存，是什么类型，在哪个分配器里。

我曾经遇到过一个案例，模型跑着跑着显存涨了2GB。一查memory_summary，发现有个中间变量没释放，是个巨大的float32张量。嗯，改成float16就解决了。

2.3 torch.cuda.memory_snapshot()：给显存拍“CT”

如果说memory_summary是体检报告，那memory_snapshot就是CT扫描。

它能生成一个完整的分配历史，包括每个分配的时间、堆栈、大小。

import torch

# 获取快照
snapshot = torch.cuda.memory_snapshot()

# 分析快照（示例）
for seg in snapshot:
    if seg['segment_size'] > 100 * 1024 * 1024:  # 大于100MB
        print(f"大块分配: {seg['segment_size'] / 1024**2:.2f} MB")
        print(f"堆栈: {seg.get('history', [{}])[0].get('traceback', '无')}")

这个工具强在哪？它能告诉你“谁分配了这块显存”。

我一般这么用：在训练循环开始和结束时各拍一张快照，对比差异。如果结束时的快照里还有大量未释放的张量，那基本就是泄漏了。

2.4 py-spy 与 tracemalloc 联合使用：Python 层的“显微镜”

有时候，泄漏不在PyTorch层，而在Python层。

比如，你有个list不断append张量，或者有个缓存没清理。这时候，torch的工具就看不到了。

怎么办？用py-spy + tracemalloc。

# 安装
pip install py-spy

# 采样CPU和内存
py-spy record -o profile.svg --pid <你的进程ID>

配合tracemalloc，可以定位到具体的代码行：

import tracemalloc

tracemalloc.start()

# ... 你的训练代码 ...

snapshot = tracemalloc.take_snapshot()
top_stats = snapshot.statistics('lineno')

for stat in top_stats[:10]:
    print(stat)

我记得有一次，一个同事说显存泄漏查了三天没找到。我过去一看，他用py-spy一采样，发现有个__del__方法里引用了全局变量，导致对象永远无法释放。嗯，这种问题，光靠nvidia-smi根本看不出来。

2.5 自定义显存监控装饰器：你的“私人医生”

工具再好，不如自己写个趁手的。

我习惯写一个装饰器，自动监控每个函数的显存变化。

import torch
import functools

def monitor_gpu_memory(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # 调用前
        torch.cuda.synchronize()
        before = torch.cuda.memory_allocated()
        
        result = func(*args, **kwargs)
        
        # 调用后
        torch.cuda.synchronize()
        after = torch.cuda.memory_allocated()
        
        diff = after - before
        if diff > 100 * 1024 * 1024:  # 超过100MB
            print(f"[WARNING] {func.__name__} 显存增长: {diff / 1024**2:.2f} MB")
        
        return result
    return wrapper

# 使用示例
@monitor_gpu_memory
def train_step(batch):
    # 你的训练逻辑
    pass

这个装饰器，我几乎每个项目都会用。它不会告诉你泄漏原因，但能快速定位“哪个函数在吃显存”。

知识体系总览

这五个工具，从宏观到微观，覆盖了显存泄漏检测的完整链路：

你想想看，从nvidia-smi看到显存异常，到memory_summary定位张量，再到py-spy找到代码行，最后用装饰器自动化监控——这一套组合拳打下来，还有什么是查不出来的？

我个人建议，把这五个工具做成一个脚本，一键启动。这样每次训练前先跑一遍，心里踏实。