3. PyTorch 显存管理机制：cudaMalloc、内存分配器缓存、torch.cuda.memory_summary()

搞深度学习的人，十个里有八个被OOM折磨过。模型一改大，batch一调高，啪，OutOfMemoryError。

很多人第一反应是「显存不够了，换卡」。其实很多时候，是你不了解PyTorch怎么管显存的。你想想看，每次你创建一个Tensor，底层发生了什么？申请一次显存真的要调用一次cudaMalloc吗？

我个人习惯是，先搞清楚分配器怎么工作的，再动手优化。这一章，我们就把它掰开揉碎了讲。

3.1 从 cudaMalloc 说起 —— 慢，是真的慢

CUDA的cudaMalloc是个底层API。它给GPU分配物理显存。但是这个调用非常慢。

我做过一个粗略测试：调用一万次cudaMalloc分配4KB的小块，总耗时接近几百毫秒。你可能觉得几百毫秒还好，但训练时每个step里可能触发几十次小分配，积少成多就卡住了。

更麻烦的是cudaFree。它不仅仅是释放，还会触发CUDA context的同步。这就像你还在高速上跑呢，突然踩了脚刹车。

为什么慢？因为cudaMalloc/cudaFree需要和GPU驱动通信，涉及用户态到内核态的切换，还有全局锁。这个开销远大于CPU上的malloc/free。

所以PyTorch不傻。它不直接拿cudaMalloc给小张量用。它用了另一套机制。

3.2 PyTorch 的内存分配器缓存 —— 先留着，别急着还

PyTorch的CUDA分配器，本质上就是一个大管家。它从GPU拿一大块显存，切成小块，谁要用就借，谁不用就收回库里，但不还给操作系统。

这种方式，专业术语叫「缓存式分配」(caching allocator)。

核心思想很简单：

  • 分配时：从已缓存的内存块里找一个大小合适的，直接复用。找不到，才调用cudaMalloc向GPU申请新块。
  • 释放时：不调用cudaFree，把内存块放回缓存池，标记为可用。

嗯，这里很关键。这样做的好处是，避免了频繁的cudaMalloc/cudaFree，大大加快了分配速度。我见过一些优化后的场景，分配速度能提升上百倍。

但副作用也明显——PyTorch占着的显存，不会释放给其他进程。你以为显存占用很高，其实里面有很多是「空闲但缓存」的块。

PyTorch的分配器还管理着两种缓存模式，我整理了一张表格：

模式 描述 适用场景
native 默认分配器，块大小按2的幂次对齐，合并相邻空闲块 通用训练/推理，大多数情况
cudaMallocAsync 使用CUDA 11.2+的流序分配，更细粒度 多流并发、重叠计算与通信的场景

我建议你，如果只跑单卡单流训练，就用默认的native分配器，别折腾。多流推理时，再考虑cudaMallocAsync。

注意，设置cudaMallocAsync需要两行代码：

torch.cuda.memory._set_allocator_settings('backend:cudaMallocAsync')
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'backend:cudaMallocAsync'

改了之后重启Python进程才生效。

3.3 核心 API 与配置控制

PyTorch提供了一组环境变量和控制接口，帮你调整分配器的行为。

先说最常用的PYTORCH_CUDA_ALLOC_CONF。格式是逗号分隔的key:value对：

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:64,expandable_segments:True"

几个重要参数：

  • max_split_size_mb：当一个缓存块太大时，是否允许分割。设置小一点，可以减少碎片，但可能增加allocator查找时间。默认是无限大，也就是不做分割限制。
  • expandable_segments：允许分配器从CUDA申请更大的虚拟地址区间，而不是每次只申请一个物理块。对减少地址空间碎片很有帮助。
  • roundup_power2_divisions：控制大小对齐策略。对防止小碎片有用。

我在项目中遇到过一个问题：多卡训练时，某张卡显存突然涨了1GB却查不出原因。后来用 expandable_segments:True 就稳定了。

注意！expandable_segments 需要 CUDA 11.4+ 和较新的 PyTorch 版本。老版本不兼容，会直接报错。

另外，你也可以在代码里通过 torch.cuda.memory 模块，手动控制一些行为：

# 清空缓存池中未使用的块，释放回CUDA（但不是给OS）
torch.cuda.empty_cache()

# 设置PyTorch分配器的split阈值
torch.cuda.memory.set_per_process_memory_fraction(0.8)

要我说，empty_cache() 这玩意儿很迷惑。你以为它会释放显存，其实它只是把缓存池里的空闲块还给cudaMalloc，并没有彻底消失。你跑nvidia-smi看，显存占用可能一点都不变。

3.4 使用 torch.cuda.memory_summary() 诊断显存

这是我最常用的诊断工具，没有之一。

torch.cuda.memory_summary() 会打印一份完整的显存使用报告，包含：

  • 已分配、已缓存、空闲块的数量和大小
  • 大块内存的详细地址信息
  • 分配器内部状态（如分割次数、合并次数）
  • 段内碎片化程度

调用方式极其简单：

print(torch.cuda.memory_summary())

输出内容长什么样？我截取一段典型示例：

| segment size | segment count | allocated size | allocated count |
|--------------|---------------|----------------|-----------------|
|  16384       | 12            | 8192           | 8               |
|  32768       | 8             | 16384          | 5               |
|  65536       | 4             | 32768          | 3               |
|  262144      | 2             | 131072         | 2               |
|  2097152     | 3             | 1048576        | 2               |

Allocated Memory: 1568000 bytes
Cache Hits: 1203
Cache Misses: 27
Segment Splits: 5
Segment Merges: 3

怎么看？

  • Cache Hits / Cache Misses：命中率高说明分配器工作良好，否则你可能需要调整分配器参数。
  • Segment Splits / Merges：分割太多说明请求的大小太乱，碎片化严重。
  • Allocated vs Cached：两者差距大，说明有很多空闲缓存没释放。但这不一定是坏事。

我记得有一次训练一个BERT-large模型，发现cache hits只有75%，大量misses导致频繁cudaMalloc。我加上 max_split_size_mb:64 之后，hits提升到98%，训练速度也快了5%。

3.5 实战：一张图看懂显存管理流程

下面这张图，把上面说的内容串起来了。你可以跟着它走一遍完整的分配与释放流程。

我个人建议你把它截图存下来，遇到OOM时对着排查。

PyTorch 显存管理流程与诊断 Python Tensor 创建 请求固定大小显存 缓存池中 是否有空闲块？ ✔ ✘ 直接复用缓存块 调用 cudaMalloc 新段加入缓存池 Tensor 参与计算 Tensor 引用归零 放回缓存池 不调 cudaFree memory_summary() 诊断 关键说明 1. 复用命中率由 PYTORCH_CUDA_ALLOC_CONF 中的参数控制。 2. empty_cache() 仅清空空闲缓存块，不释放给OS。 3. memory_summary() 可查看 splits/merges/cache hits/misses 等指标。 4. 碎片化严重时，考虑 expandable_segments=True。

3.6 避坑指南与实用建议

说了这么多，给你几个实用建议。

1. 不要滥用 empty_cache()

我曾经鬼迷心窍，在每个step后都调一次 torch.cuda.empty_cache()。结果训练慢了30%，因为腾出缓存后，下一个step又得重新分配。除非你想在换模型前释放一点显存，否则别碰它。

2. 碎片化是隐形杀手

如果你发现显存占用不高但OOM，十有八九是碎片化。解决方案：

  • 设置 max_split_size_mb:128 限制分割
  • 启用 expandable_segments:True
  • 尽量统一Tensor的大小，比如把batch size做成2的幂次

3. 多进程显存不共享

Distributed Data Parallel (DDP) 时每个进程独立分配显存，互不影响。但如果你开太多worker，单张卡上会叠加多个进程的分配器，显存占用可能翻倍。

4. 善用 memory_summary() 打日志

我习惯在每个epoch结束时输出一次 memory_summary()，记录到日志文件。这样一旦训练过程显存溢出，我可以回头排查是哪个epoch的分配行为异常。

小技巧：把 memory_summary 输出定向到文件，方便对比不同配置下的分配器表现。

with open('mem_log.txt', 'a') as f:
    f.write(torch.cuda.memory_summary())

3.7 总结

说到底，PyTorch的显存管理就三件事：

  1. cudaMalloc —— 慢，所以能不用就别用
  2. 缓存分配器 —— 用空间换时间，将空闲块暂存复用
  3. memory_summary —— 你的X光机，一眼看穿显存在干嘛

理解这个机制之后，你就能回答很多常见问题了：

  • 为什么显存没释放？—— 因为分配器缓存着。
  • 为什么内存占用稳步增长？—— 可能是碎片化，也可能是缓存块没重用到。
  • 为什么调了empty_cache没效果？—— 它不释放物理显存，只清空闲块。

我最后说一句：不要盲目崇拜「显存占用越低越好」。一个高效的分配器，占用略高但运行飞快，比一个低占用但频繁分配卡顿的系统，要好得多。你该关注的指标是 Cache Hits / Misses 比例和训练吞吐，而不是nvidia-smi上那个数字。

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