2、显存分配机制：CUDA显存分配器原理、cudaMalloc与cudaFree的底层行为、PyTorch缓存分配器（caching allocator）的工作机制

显存分配，听起来好像就是「申请一块内存，用完释放」这么简单。但实际做深度学习训练时，你会发现事情远没那么单纯。我见过不少团队，模型结构没问题，数据加载也没错，就是训练到一半突然报 CUDA out of memory。查来查去，发现是显存碎片化在作祟。

要理解碎片化，得先搞清楚显存是怎么分配的。今天我们就从底层往上，一层层拆开来看。

2.1 CUDA 显存分配器：底层的那位「管家」

CUDA 驱动层维护着一个显存分配器。它的职责很简单：管理 GPU 上的物理显存页。

当你调用 cudaMalloc 时，驱动会做两件事：

  1. 从空闲页列表中找一块足够大的连续物理显存
  2. 如果没有合适的，就触发一次「回收」或「整理」

这里有个关键点：CUDA 驱动分配的最小单位是「页」，通常是 64KB 或 2MB（取决于 GPU 架构和分配大小）。

核心概念：CUDA 驱动分配显存时，是按「页」对齐的。你申请 1 字节，它可能给你 64KB。这本身就会造成内部碎片。

我个人习惯把 CUDA 驱动层的分配器想象成一个「粗放型管家」。它只管大块分配，不太关心小块内存的复用。你申请 100 次 1MB，它可能给你 100 个不同的 2MB 页，中间留下大量空洞。

2.2 cudaMalloc 与 cudaFree 的底层行为

咱们直接看代码。这是最基础的用法：

float* d_ptr;
cudaMalloc(&d_ptr, 1024 * sizeof(float));  // 申请 4KB
// ... 使用 ...
cudaFree(d_ptr);                            // 释放

看起来简单吧？但底层发生的事情，比你想象的多：

  • cudaMalloc 调用时：驱动会检查当前进程的显存上下文。如果空闲列表里有合适大小的连续块，直接返回。如果没有，驱动会尝试从系统预留的显存池中切一块出来。如果还不够，就会触发一次「同步等待」，等之前提交的 kernel 执行完，回收它们占用的显存。
  • cudaFree 调用时：驱动不会立即把显存还给系统。它会把这块内存标记为「空闲」，放回空闲列表。这样做是为了下次 cudaMalloc 能快速复用。

注意：频繁的 cudaMalloc 和 cudaFree 是性能杀手。每次调用都有内核态切换的开销，而且容易导致显存碎片化。我曾经在一个项目中，每步训练都动态分配和释放中间结果张量，结果训练速度直接掉了 30%。

为什么会这样？因为 cudaMalloc 和 cudaFree 是同步操作。它们会阻塞 CPU 线程，直到 GPU 完成所有未完成的工作。你想想看，每次分配都要等 GPU 同步，这得多慢。

2.3 PyTorch 缓存分配器（Caching Allocator）

PyTorch 的设计者当然知道 cudaMalloc 的痛点。所以他们搞了一个「缓存分配器」，也叫 Caching Allocator。说白了，就是在 CUDA 驱动之上又加了一层「智能管家」。

它的工作机制，我总结为三步：

  1. 预分配：当你第一次在 GPU 上创建张量时，PyTorch 会调用 cudaMalloc 申请一大块显存（比如 2GB）。这块显存被缓存分配器管理，不会轻易还给驱动。
  2. 内部复用：后续你创建张量时，缓存分配器会从自己的「空闲块列表」中找一块大小合适的显存。如果找不到，它会尝试合并相邻的空闲块，或者从预分配的大块中切出一块。
  3. 延迟释放：当你删除张量（调用 del 或超出作用域）时，缓存分配器不会调用 cudaFree，而是把这块显存标记为「空闲」，放回自己的池子里。

这个机制的好处很明显：减少了 cudaMalloc 和 cudaFree 的调用次数，避免了频繁的内核态切换。同时，通过复用显存块，也降低了碎片化的概率。

小技巧：你可以通过 torch.cuda.memory_summary() 查看当前缓存分配器的状态。它会显示已分配块、空闲块、以及碎片化程度。我每次调优显存时，第一件事就是跑这个命令。

但缓存分配器也不是万能的。它有一个「阿喀琉斯之踵」：当分配大小不一时，容易产生外部碎片。

举个例子：你分配了 100MB、200MB、300MB 三个张量，然后释放了中间的 200MB。这时空闲列表里有一个 200MB 的空洞。如果你接下来要分配一个 250MB 的张量，缓存分配器发现没有合适的连续块，就会触发一次「整理」——尝试合并相邻的空闲块。但如果 200MB 空洞两边的块都还在使用，那就合并不了。最终，缓存分配器只能调用 cudaMalloc 向驱动申请新的显存。

这就是碎片化的典型场景。

2.4 一张图看懂显存分配流程

下面我用一张 SVG 图来展示整个流程，从 PyTorch 张量创建到物理显存分配：

PyTorch 显存分配流程 PyTorch 张量 torch.tensor(...).cuda() 缓存分配器 查找空闲块 / 合并 / 切分 命中率 ≈ 90%+ 命中（直接复用） 未命中 cudaMalloc 内核态调用 CUDA 驱动 页分配 / 回收 GPU 物理显存（页式管理） 张量 A (100MB) 空洞 张量 B (120MB) 空洞 张量 C (150MB) 空闲大块 碎片化：空洞无法合并，导致大块分配失败 返回显存指针

2.5 避坑指南

基于上面的原理，我给你几个实战建议：

  • 尽量复用张量：不要频繁创建和销毁中间结果。用 torch.zeros_like 或原地操作（inplace=True）来复用显存。
  • 注意分配大小：尽量让张量大小对齐到 512 字节或 1KB。这能减少内部碎片。
  • 监控碎片化程度：用 torch.cuda.memory_stats() 查看 allocated_bytes.all.peak 和 reserved_bytes.all.peak 的比值。如果比值接近 1，说明碎片化严重。

我曾经踩过的坑：在一个 NLP 项目中，我用了动态 batch 大小，每个 step 的序列长度都不一样。结果显存碎片化严重，训练到第 1000 步就 OOM 了。后来我改成固定序列长度 + padding，碎片化问题直接消失。你想想看，固定大小的张量分配，缓存分配器几乎 100% 命中，效率能不高吗？

嗯，到这里你应该对显存分配机制有了清晰的认识。从 CUDA 驱动的页式管理，到 cudaMalloc 的同步开销，再到 PyTorch 缓存分配器的复用策略，每一层都有自己的脾气。理解了这些，你就能在实战中避开那些「莫名其妙」的显存错误了。

一句话总结：显存分配不是简单的 malloc/free，而是一场「复用 vs 碎片」的博弈。PyTorch 的缓存分配器帮你赢了大半，但剩下的那部分，得靠你自己来优化。

想要更多资料，联系可可

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