4. 显存泄漏排查

显存泄漏，说白了就是程序申请了显存不还。

你跑着跑着，nvidia-smi一看，显存占用像坐了火箭一样往上窜。最后啪一下，OOM Killer出来把进程干掉了。留下一堆僵尸进程，卡在那不动。

我遇到过好几次这种情况。有一次线上推理服务，跑了三天，显存从4G涨到了24G。用户反馈越来越慢，最后直接挂了。排查下来，就是某个算子每次调用都偷偷分配显存，但从不释放。

4.1 显存泄漏的典型症状

先说说怎么判断是不是显存泄漏。别一看到显存高就慌了，得看趋势。

症状	表现	说明
显存持续增长	nvidia-smi显示显存占用随时间线性或阶梯式上升	每次推理或训练迭代后，显存都比之前多一点点
OOM Killer触发	dmesg或/var/log/kern.log中出现"Out of memory: Kill process"	系统内存或显存耗尽，内核强制杀掉进程
进程残留	kill掉主进程后，子进程或GPU上下文仍然占用显存	nvidia-smi里还能看到进程，但ps aux里已经找不到了
性能下降	推理延迟增加，训练吞吐量下降	显存碎片化严重，分配新显存越来越慢

嗯，这里要注意一点：显存占用高不等于泄漏。比如你模型本身就大，batch size设得高，那显存占用高是正常的。关键看「增长」——是不是每次操作后都涨，而且不回落。

4.2 排查工具与手段

排查显存泄漏，我习惯用这几招。说白了就是「盯住分配，抓住释放」。

4.2.1 nvidia-smi 实时监控

最基础的手段。但别只看一眼，要持续看。

# 每1秒刷新一次，显示进程级显存占用
watch -n 1 nvidia-smi

# 或者用这个，只显示显存使用情况
nvidia-smi --query-gpu=index,memory.used,memory.total --format=csv -l 1

我个人习惯写个脚本，把显存占用记录到文件里，回头画个趋势图。这样一眼就能看出是不是在涨。

4.2.2 使用PyTorch的显存跟踪工具

如果你用PyTorch，那它的内置工具很好使。

import torch

# 打印当前显存分配情况
print(torch.cuda.memory_summary())

# 记录某个时刻的显存快照
torch.cuda.memory_snapshot()

# 开启显存统计
torch.cuda.memory_stats()

我记得有一次排查一个Transformer模型的泄漏，就是用memory_summary()发现了一个Embedding层的显存没被释放。每次前向传播都新建了一个临时tensor，但没显式删除。

4.2.3 使用cuda-memcheck

NVIDIA官方工具，专门查显存泄漏和越界访问。

# 检查显存泄漏
cuda-memcheck --leak-check full ./your_program

# 只检查泄漏，不检查其他错误
cuda-memcheck --tool memcheck --leak-check full ./your_program

这个工具会报告哪些分配没有对应的释放。不过它会让程序跑得慢很多，生产环境慎用。

4.2.4 手动标记与日志

有时候工具查不出来，就得靠人工埋点。

import torch
import gc

def log_memory(step_name):
    allocated = torch.cuda.memory_allocated() / 1024**2
    cached = torch.cuda.memory_reserved() / 1024**2
    print(f"[{step_name}] allocated: {allocated:.2f} MB, cached: {cached:.2f} MB")

# 在关键步骤前后打点
log_memory("before forward")
output = model(input_data)
log_memory("after forward")

loss = criterion(output, target)
log_memory("after loss")

loss.backward()
log_memory("after backward")

optimizer.step()
log_memory("after step")

# 手动清理
del output, loss
gc.collect()
torch.cuda.empty_cache()
log_memory("after cleanup")

你想想看，如果每次迭代后allocated都比上一次多，那肯定有泄漏。如果allocated不变但cached在涨，那可能是缓存策略的问题，不一定是泄漏。

4.3 常见泄漏原因与修复

我总结了几种最常见的显存泄漏场景。你排查的时候可以按这个顺序来。

4.3.1 未释放的中间变量

这是最常见的。比如在循环里创建了tensor，但没有显式删除。

# 错误写法：每次循环都创建新tensor，但不释放
for i in range(1000):
    x = torch.randn(1024, 1024).cuda()
    # 做一些操作...
    # 没有del x，x一直占着显存

# 正确写法：及时释放
for i in range(1000):
    x = torch.randn(1024, 1024).cuda()
    # 做一些操作...
    del x  # 显式释放
    # 或者用with作用域
    with torch.no_grad():
        y = torch.randn(1024, 1024).cuda()
        # y在with块结束后自动释放

4.3.2 DataLoader的pin_memory泄漏

这个坑我踩过。DataLoader设置pin_memory=True时，如果worker进程异常退出，pin_memory的显存可能不会被释放。

# 问题场景
dataloader = DataLoader(dataset, batch_size=32, pin_memory=True, num_workers=4)

# 解决方案
# 1. 减少num_workers，或者设为0
# 2. 使用persistent_workers=True，保持worker存活
# 3. 在每次epoch结束后手动清理
for epoch in range(num_epochs):
    for batch in dataloader:
        # 训练...
        pass
    # 清理
    torch.cuda.empty_cache()

4.3.3 模型或优化器的状态累积

有些人在训练循环里反复创建优化器或学习率调度器，导致状态不断累积。

# 错误写法
for epoch in range(100):
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    # 每次循环都创建新的优化器，旧的没释放

# 正确写法
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
for epoch in range(100):
    # 复用同一个优化器
    optimizer.zero_grad()
    # ... 训练 ...

4.3.4 多进程中的显存泄漏

用torch.multiprocessing时，子进程的显存可能不会随进程退出而释放。

# 问题：子进程结束后，显存仍然被占用
import torch.multiprocessing as mp

def worker(rank):
    model = MyModel().cuda(rank)
    # 训练...
    # 进程结束后，显存可能没释放

# 解决方案：在子进程结束时显式清理
def worker(rank):
    model = MyModel().cuda(rank)
    # 训练...
    del model
    torch.cuda.empty_cache()
    # 或者使用mp.spawn的join机制确保清理

4.4 避坑指南

4.5 知识体系图

下面这张图总结了显存泄漏排查的核心逻辑。你排查的时候可以照着这个流程走。