PyTorch显存管理机制：CUDA缓存分配器与碎片化实战

各位同学好，我是你们的GPU系统工程师。今天咱们聊聊PyTorch的显存管理机制。说实话，这块内容我当年踩过不少坑，有一次训练一个超大模型，明明显存还有几个GB，结果OOM了。排查了半天，发现是缓存分配器在搞鬼。嗯，咱们今天就把它彻底讲透。

1. CUDA缓存分配器：PyTorch的显存管家

PyTorch的显存管理，核心就是它的CUDA缓存分配器。说白了，它就像一个管家，负责向GPU申请显存，然后分配给各个张量使用。

我个人习惯把分配器的工作流程分成三步：

  • 申请：当你的代码创建张量时，分配器向GPU驱动申请显存
  • 缓存：张量释放后，显存不立即归还给驱动，而是留在缓存池里
  • 复用：后续有新的张量需要显存时，优先从缓存池里拿

为什么要搞这么复杂？直接申请、直接释放不行吗？

你想想看，每次向GPU驱动申请显存，都是一次系统调用，开销很大。我在项目中遇到过，频繁创建小张量时，光系统调用的时间就占了训练时间的15%。所以PyTorch搞了个缓存池，就是为了减少这种开销。

核心要点：PyTorch的显存分配器会预申请一大块显存，然后内部管理。你看到的显存占用，不全是实际数据占的，很大一部分是缓存。

2. 缓存池策略：分段管理与复用

PyTorch的缓存池，内部是怎么管理的？我画了一张图，大家一看就明白。

PyTorch CUDA缓存池结构图 大块缓存池（Large Pool） - 大于1MB的块 小块缓存池（Small Pool） - 小于等于1MB的块 CUDA缓存分配器 实际张量占用 缓存空闲块 注：缓存池中的空闲块不会立即归还给GPU驱动

PyTorch把缓存分成两个池子：

  • 大块缓存池：存放大于1MB的显存块。我一般用它来存模型参数和梯度
  • 小块缓存池：存放小于等于1MB的显存块。适合存中间激活值

为什么要分两个池子？

我在项目中遇到过，如果不分池，大块和小块混在一起，很容易产生碎片。比如你释放了一个2MB的块，但下一个请求要3MB，这个2MB的块就用不上，白白浪费了。

实战技巧：如果你发现显存占用很高但实际数据不多，可以用 torch.cuda.memory_summary() 查看缓存池的详细情况。我曾经靠这个命令定位到一个缓存泄漏的bug。

3. 显存碎片化：隐形杀手

显存碎片化，说白了就是显存被切成了很多小块，每个块都不大，但加起来总量不少。结果你要申请一个大块时，发现没有连续空间，OOM了。

为什么会这样？

我举个例子。假设你有一个100MB的连续显存，先分配了30MB给张量A，再分配40MB给张量B，最后分配30MB给张量C。然后你释放了张量A和张量C，显存变成：

| 空闲30MB | 占用40MB | 空闲30MB |

现在你想分配一个50MB的张量D，虽然总空闲有60MB，但没有连续50MB的空间，只能OOM。这就是碎片化。

注意：碎片化在长时间训练中尤其严重。我见过一个训练跑了3天，显存从12GB涨到15GB，但实际数据只占8GB，剩下的全是碎片。

怎么缓解碎片化？

  • 预分配大块：用 torch.cuda.empty_cache() 清空缓存，然后一次性申请大块显存
  • 控制张量生命周期：尽量让大小相近的张量同时创建、同时释放
  • 使用检查点：用激活检查点技术，减少中间张量的数量

4. PyTorch显存回收机制：你以为释放了，其实没有

很多同学以为，用 del tensor 或者 tensor = None 就能释放显存。嗯，这里要注意，事情没那么简单。

PyTorch的显存回收分两步：

  1. Python层面：引用计数归零，张量对象被销毁
  2. 缓存层面：显存块回到缓存池，但不归还给GPU驱动

也就是说，你 del 之后，显存占用不会立刻下降。它只是从「使用中」变成了「空闲缓存」。只有当你调用 torch.cuda.empty_cache() 时，才会把缓存池里的空闲块归还给驱动。

我曾经犯过一个错误：训练循环里每步都调用 empty_cache()，结果训练速度慢了3倍。因为每次清空缓存后，下一步又要重新申请，系统调用开销巨大。

最佳实践：不要频繁调用 empty_cache()。只在以下情况使用：

  • 训练结束后，释放显存给其他进程用
  • 显存碎片化严重，导致OOM时
  • 切换模型或数据集时

5. 实战：用代码看透显存管理

光说不练假把式。咱们写段代码，看看PyTorch的显存管理到底怎么工作的。

import torch

# 查看初始显存
print(f"初始显存占用: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
print(f"初始缓存大小: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")

# 创建张量
x = torch.randn(1000, 1000).cuda()
print(f"\n创建张量后:")
print(f"显存占用: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
print(f"缓存大小: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")

# 删除张量
del x
print(f"\n删除张量后:")
print(f"显存占用: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
print(f"缓存大小: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")

# 清空缓存
torch.cuda.empty_cache()
print(f"\n清空缓存后:")
print(f"显存占用: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
print(f"缓存大小: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")

运行这段代码，你会发现：

  • 删除张量后，memory_allocated 降为0，但 memory_reserved 不变
  • 调用 empty_cache() 后，memory_reserved 才下降

这就是PyTorch的缓存策略——宁可留着缓存，也不频繁申请释放。

调试利器：用 torch.cuda.memory_snapshot() 可以获取每个显存块的详细信息。我调试显存泄漏时，全靠这个函数定位问题。

6. 避坑指南：我踩过的那些坑

最后，分享几个我实际工作中遇到的坑：

  • 坑一：DataLoader的num_workers设置太大，每个worker都缓存了显存，导致OOM。解决方案：限制worker数量，或者用 torch.multiprocessing.set_sharing_strategy('file_system')
  • 坑二：在循环中创建临时张量，没有及时释放，导致缓存池膨胀。解决方案：用 with torch.no_grad() 包裹，或者手动 del
  • 坑三：混合精度训练时，master weights和fp16 weights同时存在，显存翻倍。解决方案：用 torch.cuda.amp 自动管理

好了，关于PyTorch显存管理机制，今天就讲到这里。记住一句话：显存占用 ≠ 实际数据大小，缓存池才是真正的显存消耗大户。下次遇到OOM，先看看缓存池里有多少空闲块，别急着加显卡。

本章总结：

  • CUDA缓存分配器：预申请、缓存、复用三步走
  • 缓存池策略：大块池和小块池分开管理
  • 显存碎片化：连续空间不足导致OOM
  • 回收机制：del不释放缓存，empty_cache才释放

想要更多资料，联系可可

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