1、显存泄漏的根源

大家好，我是你们的GPU系统讲师。今天咱们聊聊显存泄漏——这个让无数AI工程师头疼的问题。

说实话，我见过太多团队在模型训练到一半时，突然报出"CUDA out of memory"的错误。那种感觉，就像你开车上了高速，突然发现油箱漏了。嗯，今天我们就来彻底搞明白，显存泄漏到底是怎么回事。

1.1 显存泄漏的定义

显存泄漏，说白了就是程序申请了显存空间，用完之后却不归还。你想想看，显存是有限的，比如一张A100只有80GB。如果每次迭代都泄漏一点点，累积下来，最终就会把显存撑爆。

我个人的习惯是，把显存泄漏看作一种"内存管理失忆症"——程序忘了自己曾经占过哪些地方，自然也就没法释放。

核心定义：显存泄漏是指已分配的GPU显存在使用完毕后，未被正确释放，导致可用显存逐渐减少的现象。

1.2 显存泄漏 vs 内存泄漏

很多同学会问：显存泄漏和内存泄漏有什么区别？

其实本质是一样的——都是资源没释放。但有两个关键差异：

对比维度 内存泄漏 显存泄漏
资源位置 CPU主存（RAM） GPU显存（VRAM）
影响范围 单个进程，系统可能OOM 整个GPU设备，影响所有进程
检测难度 相对容易，有valgrind等工具 较难，需要nvidia-smi等监控
恢复方式 进程退出后自动回收 进程退出后显存才释放

我记得有一次，团队里一个小伙子调试了三天，一直以为是内存泄漏，结果发现是显存泄漏。为什么？因为CPU内存看着很正常，但GPU显存却在悄悄上涨。所以，排查方向很重要。

1.3 显存泄漏的常见原因

根据我多年的实战经验，显存泄漏主要有四大元凶。咱们一个一个来看。

1.3.1 未释放张量

这是最常见的情况。你在训练循环里创建了张量，但忘记用del删除，或者没有让Python的垃圾回收机制及时处理。

# 错误示例：每次循环都创建新张量，但不释放
for i in range(1000):
    x = torch.randn(10000, 10000).cuda()  # 每次分配约400MB
    # 做一些计算...
    # 没有释放 x，显存持续增长

正确的做法是：

# 正确示例：及时释放
for i in range(1000):
    x = torch.randn(10000, 10000).cuda()
    # 做一些计算...
    del x  # 显式释放
    torch.cuda.empty_cache()  # 清空缓存

小技巧：我习惯在每次迭代末尾调用 torch.cuda.empty_cache()，虽然会有一点性能开销，但能有效防止显存泄漏。尤其是在调试阶段，这个习惯能帮你省下大量排查时间。

1.3.2 梯度累积

梯度累积本身是个好技术，能让小batch size模拟大batch size的效果。但很多人用错了。

为什么会泄漏？因为PyTorch默认会保留计算图，直到你调用backward()。如果你在累积梯度时没有正确清零，计算图就会一直保留在显存里。

# 错误示例：梯度累积导致计算图泄漏
optimizer.zero_grad()
for i, data in enumerate(dataloader):
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()  # 计算图保留
    if (i+1) % accumulation_steps == 0:
        optimizer.step()
        # 忘记清零梯度！

正确的做法是：

# 正确示例：每次迭代后清零梯度
optimizer.zero_grad()
for i, data in enumerate(dataloader):
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()
    if (i+1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()  # 关键：清零梯度

注意：我曾经在一个分布式训练项目中，因为梯度累积没清零，导致每个GPU的显存从12GB涨到了40GB，最后直接OOM。排查了整整两天才发现是这个问题。所以，梯度累积一定要配合 optimizer.zero_grad() 使用。

1.3.3 DataLoader句柄泄漏

这个比较隐蔽。DataLoader在加载数据时，会创建多个子进程来预取数据。如果这些子进程没有正确关闭，它们会持有CUDA上下文，导致显存无法释放。

# 错误示例：DataLoader未正确关闭
dataloader = DataLoader(dataset, num_workers=4, pin_memory=True)
for epoch in range(10):
    for data in dataloader:
        # 训练...
    # 每个epoch结束后，dataloader的子进程可能没有完全退出

解决方案：

# 正确示例：使用上下文管理器
with DataLoader(dataset, num_workers=4, pin_memory=True) as dataloader:
    for epoch in range(10):
        for data in dataloader:
            # 训练...
        # 退出with块时，自动清理子进程

避坑指南：我曾经在服务器上跑一个长时间训练任务，发现显存每过几个小时就涨一点。用 nvidia-smi 一看，发现有几个Python子进程一直挂着。后来才发现是DataLoader的句柄泄漏。从那以后，我所有DataLoader都改用 with 语句了。

1.3.4 CUDA上下文未释放

这是最容易被忽视的一点。每个使用CUDA的进程都会创建一个CUDA上下文，这个上下文会占用几百MB到几GB的显存。如果进程异常退出，或者没有正确销毁上下文，显存就泄漏了。

# 错误示例：多进程未正确管理CUDA上下文
import multiprocessing as mp

def worker():
    # 每个worker都会创建CUDA上下文
    model = MyModel().cuda()
    # 训练...
    # 退出时没有销毁上下文

# 启动多个worker
processes = [mp.Process(target=worker) for _ in range(4)]
for p in processes:
    p.start()
# 如果worker异常退出，CUDA上下文会残留

正确的做法是：

# 正确示例：使用torch.multiprocessing
import torch.multiprocessing as mp

def worker(rank):
    # 使用spawn方式启动，每个进程独立管理CUDA
    model = MyModel().cuda()
    # 训练...
    # 进程退出时，CUDA上下文自动销毁

mp.spawn(worker, nprocs=4, join=True)

重要提醒：使用 torch.multiprocessing 而不是原生的 multiprocessing，这是PyTorch官方推荐的做法。它能更好地管理CUDA上下文，避免泄漏。

知识体系总览

下面这张图，是我总结的显存泄漏知识体系。你可以把它当作排查时的路线图。

显存泄漏知识体系 显存泄漏根源 定义与区别 显存泄漏：GPU显存未释放 内存泄漏：CPU内存未释放 区别：资源位置、影响范围 四大常见原因 1. 未释放张量 创建后未del，计算图残留 2. 梯度累积 未正确清零，计算图累积 3. DataLoader句柄泄漏 子进程未关闭，持有CUDA上下文 4. CUDA上下文未释放 进程异常退出，上下文残留 排查方向：nvidia-smi + 代码审查

这张图把显存泄漏的核心知识点都串起来了。从定义到区别，再到四大原因，最后给出排查方向。你可以在实际工作中把它当作速查表。

好了，这一章的内容就到这里。记住，显存泄漏不是玄学，它是有迹可循的。只要你掌握了这些基础知识，排查起来就会事半功倍。

想要更多资料，联系可可

← 上一章 📖 返回目录 下一章 →