第4章 CUDA Caching Allocator：碎片化战争

"Yet another caching allocator for CUDA device allocations. ... With this allocator, allocations and frees should logically be considered usages of the memory segment associated with streams, just like kernel launches."

—— c10/cuda/CUDACachingAllocator.cpp:71 的设计注释开头

4.1 为什么不直接 cudaMalloc

按最朴素的实现，PyTorch 每次创建 GPU 张量时调一次 cudaMalloc，析构时调一次 cudaFree。但这条路 走不通。

打开 NVIDIA 的 CUDA 文档，能读到这样两句话：

Managing memory allocations using cudaMalloc and cudaFree causes the GPU to synchronize across all executing CUDA streams.

—— CUDA Programming Guide「Stream-Ordered Memory Allocator」一节

For all other pointers, this API may perform implicit synchronization.

—— CUDA Runtime API 手册 cudaFree 条目（"other pointers" 指不是由 cudaMallocAsync / cudaMallocFromPoolAsync 分配的指针）

也就是说 —— cudaMalloc / cudaFree 会打断 stream 之间的并发，并可能让调用线程等到 GPU 上已排队的工作跑完才返回（官方措辞是 "may"，何时真的同步由驱动决定）。PyTorch 源码自己也把它当昂贵操作：alloc_block() 调 cudaMalloc 前会临时放开 allocator 的锁，注释直接称之为 "cudaMalloc sync call"（c10/cuda/CUDACachingAllocator.cpp:3269），说它持锁调用代价很高。

对一个朴素的 ML 框架，这个代价很难承受：

# 朴素实现下，每次创建 / 释放张量都要调一次 cudaMalloc / cudaFree
a = torch.randn(1024, 1024, device='cuda')   # cudaMalloc → 可能隐式同步
b = torch.randn(1024, 1024, device='cuda')   # cudaMalloc → 可能隐式同步
c = a + b                                     # 还有这个 op 自身的临时显存分配
del a                                         # cudaFree → 可能隐式同步
del b                                         # cudaFree → 可能隐式同步

一次 forward + backward 可能创建上千个临时张量。如果每次分配 / 释放都可能引入同步 —— CPU 没法提前把 kernel 排进队列，GPU 会反复空转等待，利用率随之下降。

4.1.1 解法：用户态的"显存池"

既然 cudaMalloc 太重，PyTorch 的解法是 批发再零售：

1. 一次性 cudaMalloc 一大块（例如 2 MiB / 20 MiB）
2. 把这一大块切成小段，用户每次"分配"实际上是从池子里割一段给他
3. 用户"释放"时，不调 cudaFree，把这段还给池子
4. 下次有新的分配请求，先在池子里找合适的段，找不到再 cudaMalloc 新的大块

这就是 CUDACachingAllocator —— PyTorch 在 GPU 显存管理上做的最重要的工程优化之一。源码在 c10/cuda/CUDACachingAllocator.cpp，v2.11.0 下全文 4667 行。这一章我们把它逐层拆开。

4.1.2 这份代价到底贵在哪

cudaMalloc 的开销有多大？关键不在"分配字节"这个动作本身（那只是几条簿记指令），而在它可能附带的隐式同步 —— 调用可能一直阻塞到 GPU 上已排队的工作跑完。所以它的代价不是一个固定数字，而是"当前 GPU 队列里还剩多少工作"：空载时只有 driver 调用开销，队列排满时就是队列的排空时间。这也解释了为什么它在训练 hot loop 里格外致命 —— 恰恰是 GPU 最忙的时候，它最慢。

Caching allocator 的目标就是把这份开销摊销掉：训练初期付几次 cudaMalloc 的同步代价，之后命中池子的"分配"都是 O(log n) 的簿记，不调用 cudaMalloc。

4.2 核心数据结构：Block 与 BlockPool

打开 c10/cuda/CUDACachingAllocator.cpp:189，看 Block 结构：

struct Block {
    c10::DeviceIndex device;       // 哪张 GPU
    cudaStream_t stream;           // 分配时的 stream
    stream_set stream_uses;        // 这块被哪些 stream 用过
    size_t size;                   // 字节数（pool 里的大小）
    size_t requested_size;         // 用户实际请求的字节数
    BlockPool* pool;               // 所属池
    void* ptr;                     // 显存地址
    bool allocated;                // 是否被持有
    bool mapped;                   // 虚拟地址是否被物理页面 backed (expandable_segment)
    Block* prev;                   // split 链表前驱
    Block* next;                   // split 链表后继
    int event_count;               // 还有多少个未完成的 CUDA event
    int64_t gc_count_base;
    ExpandableSegment* expandable_segment_;
    ...
};

注意几个关键字段：

• stream_uses：记录分配流之外还有哪些 stream 用过这块（recordStream 会忽略分配流本身）。这是流间安全复用的核心
• prev / next：每次"切分"产生的子块用双向链表串起来。这让 释放时合并相邻空闲块（coalescing）成为 O(1) 操作
• event_count：只有 event_count = 0（所有相关 GPU 工作完成）时，块才能被回收

BlockPool 是 Block 的集合：

struct BlockPool {
    std::set<Block*, Comparison> blocks;        // 空闲（已缓存）块，按 (stream, size, ptr) 排序
    std::set<Block*, Comparison> unmapped;      // expandable_segment 下未映射的虚拟段，按 (stream, ptr) 排序
    const bool is_small;                        // 区分 small / large 池
    PrivatePool* owner_PrivatePool;             // CUDA Graph 私有池
    ...
};

blocks 是一棵有序的红黑树（std::set），同一 stream 内按 size 排序 —— 分配时用 lower_bound 找"刚好够"的块，O(log n)。这是 best-fit 分配策略的标准实现。

4.2.1 best-fit vs first-fit 之争

经典分配器有两条路：

• first-fit：找第一个 ≥ size 的块。实现简单，但容易在池子前部留下大量"几乎合适但不够"的小碎片
• best-fit：找"恰好够"的最小块。每次分配都精准吃掉最小块，剩下的大块保持完整 —— 但搜索代价稍高，且如果"恰好够"留下极小残块也是浪费

PyTorch 用的是 best-fit：文件顶部的设计注释写的就是 "attempts to find the smallest cached block that will fit the requested size"。在有序集合上做 best-fit，代价只是一次 lower_bound。

而 BlockComparatorSize 不是简单的 size < —— 它先比 stream、再比 size、最后比 ptr 地址（c10/cuda/CUDACachingAllocator.cpp:867）。stream 排在第一位，于是同一 stream 的空闲块在树里连成一段；get_free_block 用 lower_bound 定位后，若命中的块 stream 不同就直接判"没找到"。这对应设计注释的第一条：块释放后只能在同一个 stream 上再分配，不会跨流复用。

4.2.2 unmapped 集合：为 expandable_segments 服务

注意 BlockPool 还有一个 unmapped 集合，按地址而非 size 排序。这个集合只在 expandable_segments 模式下使用 —— 它存储那些虚拟地址已分配但物理页未映射的"幽灵块"。后面 §4.8 我们会看到它的用法。

4.3 两池子设计：small vs large

CUDACachingAllocator 把所有分配请求按大小分流：

flowchart TD
    Req["allocate(size)"] --> Decide{size 多大?}
    Decide -->|≤ 1 MiB| SP["small_blocks 池<br/>每个块 2 MiB<br/>按 512 字节倍数 round"]
    Decide -->|> 1 MiB| LP["large_blocks 池<br/>段大小见下表"]
    SP --> CudaMalloc1["池中无合适块: cudaMalloc 2 MiB 段"]
    LP --> CudaMalloc2["池中无合适块: cudaMalloc 20 MiB 段（< 10 MiB）<br/>或 round(size, 2 MiB)（≥ 10 MiB）"]
    CudaMalloc1 --> Slice["从大块切 size 段返回"]
    CudaMalloc2 --> Slice
    style SP fill:#dbeafe,stroke:#3b82f6
    style LP fill:#fef3c7,stroke:#f59e0b

关键常量都在 c10/core/AllocatorConfig.h:16-24：

constexpr size_t kSmallBuffer = 2097152;     // 2 MiB
constexpr size_t kMinBlockSize = 512;        // 512 字节最小对齐
constexpr size_t kSmallSize = 1048576;       // 1 MiB
constexpr size_t kMinLargeAlloc = 10485760;  // 10 MiB
constexpr size_t kRoundLarge = 2097152;      // 2 MiB

为什么要分两池子？因为如果不分：

• 一个 100 字节的请求和一个 1 GB 的请求都进同一个池子
• 100 字节的请求把某段切走，剩下大段被切成小碎片
• 后续 1 GB 请求找不到足够大的连续段，触发新的 cudaMalloc
• 大量小请求和少量大请求混合，碎片化以最快速度发生

把 small 与 large 隔离让"小张量大量出现"不会污染"大张量的连续性"。这是经典的 size-class allocator 设计，与 jemalloc / tcmalloc 的思想一致。

4.3.1 大小分流的边界值

PyTorch 用三个阈值定义大小分流：

源码把这条规则写在 get_allocation_size() 里（c10/cuda/CUDACachingAllocator.cpp:3130）：

static size_t get_allocation_size(size_t size) {
  if (size <= kSmallSize) {                    // ≤ 1 MiB
    return kSmallBuffer;                       //   → 去 cudaMalloc 一个 2 MiB 段
  } else if (size < kMinLargeAlloc) {          // 1 MiB ~ 10 MiB
    return AcceleratorAllocatorConfig::large_segment_size();  // → 一个 20 MiB 段
  } else {                                     // ≥ 10 MiB
    return kRoundLarge * ((size + kRoundLarge - 1) / kRoundLarge);  // → 向上取整到 2 MiB 倍数
  }
}

范围	池	新段大小
≤ 1 MiB	small_blocks	2 MiB（kSmallBuffer），切片给用户
1 MiB - 10 MiB	large_blocks	20 MiB（large_segment_size_，可配），切片给用户
≥ 10 MiB	large_blocks	向上取整到 2 MiB 倍数

注意中间那一段（1-10 MiB）不是"按需 cudaMalloc"，而是去切一个 20 MiB 的大段。文件顶部的设计注释把理由写得很直白："To reduce fragmentation, requests between 1MB and 10MB will allocate and split a 20MB block, if no free block of sufficient size is available." —— 这个尺寸段的请求单次不小、频次也不低，与其一个个 cudaMalloc，不如批一个 20 MiB 段来切。

4.4 分配流程：best-fit + split

完整的分配路径：

sequenceDiagram
    autonumber
    participant U as 调用方（如创建张量）
    participant A as Allocator
    participant P as BlockPool
    participant K as CUDA Driver
    U->>A: alloc(size, stream)
    A->>A: round size to 512 倍数
    A->>P: blocks.lower_bound(size)
    alt 找到合适块
        P-->>A: 返回 block (size >= req)
        opt block 比请求大很多
            A->>A: split: 切出 req 大小, 剩余还回池子
            A->>A: 用 prev/next 链起来
        end
    else 池里无合适块
        A->>K: cudaMalloc(round_size)
        K-->>A: 新 block
        A->>P: 加入池子
        A->>A: split (如果新块比请求大)
    end
    A->>A: block.allocated = true
    A->>A: block.stream = stream
    A-->>U: 返回 block.ptr

简化的 C++ 代码（c10/cuda/CUDACachingAllocator.cpp 中 malloc 实现）：

// 极简版 (实际代码 200+ 行处理各种边界)
Block* malloc(size_t size, cudaStream_t stream) {
    size = round_size(size);                    // 至少 512 字节倍数
    auto& pool = (size <= kSmallSize) ? small_blocks : large_blocks;

    // 1. 在 pool 里找"≥ size 的最小块"
    auto it = pool.blocks.lower_bound(make_search_key(size, stream));
    if (it != pool.blocks.end()) {
        Block* block = *it;
        pool.blocks.erase(it);
        if (should_split(block, size)) {
            block = split(block, size);          // 切出 size, 剩余回池
        }
        block->allocated = true;
        return block;
    }

    // 2. 池里没找到，cudaMalloc 新块
    void* ptr;
    size_t alloc_size = get_allocation_size(size);
    cudaMalloc(&ptr, alloc_size);
    Block* block = new Block(device, stream, alloc_size, &pool, ptr);
    if (should_split(block, size)) {
        block = split(block, size);
    }
    block->allocated = true;
    return block;
}

源码里并没有叫 split 的函数，切分在 alloc_found_block()（c10/cuda/CUDACachingAllocator.cpp:1641）里就地完成：把一个大块分成两个，前一个返回给用户，后一个放回池里，两者通过 prev/next 双向链表保持邻接关系。

4.4.1 split 的 trade-off

should_split(block, size) 不会无限制地切。它有几条规则：

// c10/cuda/CUDACachingAllocator.cpp:3110
bool should_split(const Block* block, size_t size,
                  bool is_expandable_segments_active) {
  if (no_split_pools.find(block->pool->owner_MempoolId()) != no_split_pools.end())
    return false;                       // 这个 mempool 被标记为"不许切"
  size_t remaining = block->size - size;
  if (block->pool->is_small || is_expandable_segments_active) {
    return remaining >= kMinBlockSize;  // small 池：剩下 ≥ 512 字节就切
  } else {
    return (size < AcceleratorAllocatorConfig::max_split_size()) &&
           (remaining > kSmallSize);    // large 池：请求 < max_split_size 且剩下 > 1 MiB 才切
  }
}

• small 池（以及开了 expandable segments 时）门槛极低：剩余 ≥ kMinBlockSize（512 字节）就切
• large 池要同时满足两条：本次请求小于用户配的 max_split_size_mb，且切完剩余大于 kSmallSize（1 MiB）
• 被标记进 no_split_pools 的 mempool 一律不切

这就是 PyTorch 在"减少 cudaMalloc 调用"和"控制碎片化"之间的权衡：切得越细，下次小请求越容易命中已有块；但切得太细，大请求又找不到连续大块。

4.4.2 一个具体例子：split 怎么发生

假设池子里有一个 64 MB 的空闲块（地址 0x1000），用户请求 4 MB：

分配前: [______64MB free______]  (0x1000 - 0x4001000)

split 后:
  [4MB used][_______60MB free_______]
  0x1000     0x401000      0x4001000
  ↑ 返回    ↑ 留在池子里, 通过 prev/next 与第一段链起来

第一段被标 allocated=true 返回给用户；第二段 size=60MB，allocated=false，仍在 pool.blocks 里、按 size 排序。它的 prev 指向第一段，第一段的 next 指向它。

后续用户释放第一段时：

// 第一段 free 时
block->allocated = false;
// 检查 next 邻居 (60MB) 是否空闲 → 是
// 合并!
block->size = 4MB + 60MB = 64MB;
delete block->next;
// 现在 block 又是一个 64MB 的空闲块

这就是 coalescing 的全过程。理想情况下经过反复分配/释放，池子最终会回到"一个 64MB 整块"的状态，可以被 empty_cache 还给 CUDA 驱动。但只要有任何一个子段还被持有 —— 整个链就无法合并回原始大小。

4.5 free 流程：合并相邻空闲块

释放比分配复杂，因为要处理 stream 同步：

sequenceDiagram
    autonumber
    participant U as 调用方（张量存储析构）
    participant A as Allocator
    participant P as BlockPool
    participant E as cuda::Event
    U->>A: free(block)
    A->>A: block.allocated = false
    A->>A: 检查 stream_uses
    alt 用过多个流
        A->>E: 在每个其他流上 cudaEventRecord
        A->>A: block.event_count = N
        Note over A: 块进入"等待回收"状态<br/>不立即放回池子
    else 单流
        A->>A: 立即可回收
    end
    Note over A: 下次 malloc 开头轮询事件（无后台线程）
    A->>E: cudaEventQuery(...)
    E-->>A: 全部 done
    A->>A: 合并相邻空闲块 (coalesce)
    A->>P: 放回 pool.blocks

合并的逻辑藏在 try_merge_blocks 里（简化版）：

void free_block(Block* block) {
    block->allocated = false;
    // 尝试合并 prev 邻居
    if (block->prev && !block->prev->allocated) {
        block->size += block->prev->size;
        block->ptr = block->prev->ptr;
        delete block->prev;     // 邻居被吸收
    }
    // 尝试合并 next 邻居
    if (block->next && !block->next->allocated) {
        block->size += block->next->size;
        delete block->next;
    }
    block->pool->blocks.insert(block);   // 合并后回池
}

coalescing 是池化分配器对抗碎片化的核心机制。理想情况下：每次释放都能向 prev / next 各扩张一次，最终把整个 cudaMalloc 来的大块"还原"成一个大空闲块。但这只有在所有切分出来的子块同时空闲时才能成立 —— 这正是后面要讲的 empty_cache 失效的根源。

4.6 stream-aware：跨流安全复用

考虑这段代码：

s = torch.cuda.Stream()
a = torch.randn(1024, device='cuda')          # 在 default stream 分配
with torch.cuda.stream(s):
    b = a + 1                                  # 在 s 流上读 a
del a                                          # default stream 释放 a
c = torch.randn(1024, device='cuda')          # 新分配 — 可能复用 a 的内存？

allocator 只知道 a 是在默认流上分配的；del a 时它会立刻把这块内存还回池子，c 在同一默认流上分配就可能拿到同一地址。可 s 流上的 a + 1 未必已经跑完 —— 数据竞争！ b 读到的可能是 c 刚写入的新值，而不是 a 的旧值。（这段示例还省略了 s.wait_stream(torch.cuda.current_stream())，s 流读 a 之前本该先等默认流把 a 生成出来。）

PyTorch 的解法：recordStream 与 stream_uses。当一个张量被另一个流使用时，调用 tensor.record_stream(other_stream)，allocator 记录这个块的 stream_uses。释放时如果 stream_uses 非空（分配流本身不会被记进去），allocator 在其中每个流上插入 cudaEvent，把块标记为 "等待事件" 状态，event 全部完成后才真正回池。

这是为什么 PyTorch 鼓励手动 tensor.record_stream(s) —— 自动跟踪做不到（没法在 Python 端拦截每次跨流读写），用户必须显式声明意图。

第 16 章讲 stream 与 collective 的协作时会回到这个机制。

4.6.1 一个真实场景：自己开 stream 做 H2D prefetch

先澄清一个常见误解：DataLoader 的 pin_memory=True 不涉及任何 CUDA stream。看 torch/utils/data/_utils/pin_memory.py 就知道，_pin_memory_loop 是主进程里的一个普通 Python 线程，它对每个 batch 调的是 data.pin_memory() —— 一次 CPU 侧的"把内存钉进页锁定区"操作，H2D 拷贝仍然由你自己的 x.to('cuda', non_blocking=True) 在当前 stream 上发起。pin_memory 让这次拷贝可以异步，但它本身不搬数据。

真正需要 record_stream 的是你自己开 stream 做 prefetch 的时候：

copy_stream = torch.cuda.Stream()
with torch.cuda.stream(copy_stream):
    gpu_batch = cpu_batch.to('cuda', non_blocking=True)   # 在 copy_stream 上分配 + 拷贝
torch.cuda.current_stream().wait_stream(copy_stream)      # 计算流等拷贝流
gpu_batch.record_stream(torch.cuda.current_stream())      # ← 少了这句就危险

gpu_batch 是在 copy_stream 上分配的，allocator 只知道它属于 copy_stream。当它在 Python 端被释放时，如果没有 record_stream 告诉 allocator "默认流也用过它"，这块内存会被立刻还给池子、并可能马上分给下一个在默认流上的张量 —— 而此时默认流上读它的 kernel 可能还没跑完。

忘了 record_stream 的症状是"输出有时对、有时错"。诊断的标准方法：在怀疑的地方加 torch.cuda.synchronize()，如果错误消失，几乎肯定是 stream race。这种隐藏 bug 在自己手写通信/计算重叠的分布式代码里特别常见。

4.6.2 event_count 与延迟回收

Block::event_count 字段和 cudaEvent_t 配合实现"延迟回收"：

// 简化的多流释放路径
void process_events() {
    for (auto& [event, block] : pending_events) {
        if (cudaEventQuery(event) == cudaSuccess) {
            block->event_count -= 1;
            if (block->event_count == 0) {
                free_block(block);             // 真正回池（合并邻居后插回 pool）
            }
        }
    }
}

每次 malloc 开头（不在 CUDA Graph capture 期间时），allocator 调 process_events() 按流轮询 pending event（c10/cuda/CUDACachingAllocator.cpp:3754）。完成的 event 让对应 block 的 event_count 减 1；event_count 归零时块才真正被合并并放回池子。

这个机制让 PyTorch 在跨流场景下既正确又高效：正确（不释放还在用的内存），又高效（不强制 device 全停）。代价是池子里有一部分块处在"等待事件"状态，暂时不可用 —— 多流用得越密、这部分占比越高，是 caching allocator 设计中"用空间换时间"的经典权衡。想看它有多大，可以用 torch.cuda.memory_stats() 里的 active_bytes.all.current 减 allocated_bytes.all.current：free() 时 allocated 立即扣减，active 要等 free_block() 真正回池才扣减，两者之差正是还在等事件的字节数（按源码判断、未实跑）。

4.7 torch.cuda.empty_cache() 为什么常常不如预期

终于到本章最具实战价值的话题。打开 c10/cuda/CUDACachingAllocator.cpp:3443 的 release_cached_blocks：

bool release_cached_blocks(...) {
    synchronize_and_free_events(context);   // 等所有 event 完成

    // 关键：只释放"非 split"的 block
    release_blocks(large_blocks, context);
    release_blocks(small_blocks, context);
    ...
}

而 release_blocks 的判断是 prev == nullptr && next == nullptr（即 !block->is_split()）才 cudaFree；expandable segment 里的块另走 unmap_block()，把其中完全空闲的物理页 unmap（c10/cuda/CUDACachingAllocator.cpp:3622）。

为什么有这个限制？因为 cudaFree 必须配对 cudaMalloc —— 你不能 cudaFree 一个 cudaMalloc 大块的中间一段。注意 "未切分" 不等于 "从没被切过"：§4.5 的 coalescing 会在子块全部空闲时把它们合并回一整块，这时 prev/next 都是空，照样能被释放。真正还不回去的，是段里还有任何子块被占用的那些段（还在等跨流事件的子块也算）—— 其中空闲的部分只能留在池里。

让我们看一个例子：

# 训练初期：
torch.cuda.empty_cache()                           # 假设清空
# 第一个 forward：
a = torch.randn(1, device='cuda')                  # 池子从 cudaMalloc 拿 2 MiB 段, 切出 512 字节（最小块）给 a
b = torch.randn(1, device='cuda')                  # 再切走 512 字节
del a                                              # a 的子块还给池, 但 b 仍占着同一段
torch.cuda.empty_cache()                           # 这 2 MiB 段里还有 b → 不能 cudaFree
del b                                              # 子块全部空闲, 合并回完整的 2 MiB 段
torch.cuda.empty_cache()                           # 现在这一段可以还给驱动了

第一次 empty_cache() 之后，memory_reserved() 仍是 2 MiB，而 memory_allocated() 只有 512 字节 —— 一个 512 字节的张量钉住了整整一段。放大到训练规模，这就是 "张量已经删了、显存却没降下来" 这个常见困惑的源头（按源码判断、未实跑）。

4.7.1 那 empty_cache 什么时候才有用

只有 池子里有整段空闲的缓存段 时它才能还回显存。torch/cuda/memory.py 里的 docstring 说得很克制：它把 "all unoccupied cached memory" 释放给其他 GPU 应用、让 nvidia-smi 看得见，但 "doesn't increase the amount of GPU memory available for PyTorch"。具体场景：

• 训练前期峰值过后，某些大型临时张量（如 attention 的中间结果）被释放，且这些张量是"独占一整个 cudaMalloc 块" 时
• 切换 batch size 后，原 batch size 留下的整块被腾出
• 跨阶段切换（如 train → eval）

但训练 hot loop 里的小张量分配/释放，empty_cache 几乎帮不上忙。调用它本身的开销反而成了拖累：它要先 cudaEventSynchronize 等完所有未完成的跨流事件，再对每个释放的段调 cudaFree，下一轮还得重新 cudaMalloc 回来。

比较合理的调用时机大致两类：

1. 你刚做完一个会留下大量空闲缓存的阶段，准备进入新阶段（如 train → val）
2. 你需要给其他进程腾显存（如 nvidia-smi 看起来你占了 16 GB，但训练只用 8 GB）

4.7.2 一个反复被踩的坑

社区 issue 里反复出现的一个误用：

# 错误：在每个 batch 后调 empty_cache
for batch in dataloader:
    loss = model(batch).sum()
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    torch.cuda.empty_cache()    # ← 几乎肯定让训练变慢

调用 empty_cache 本身要等待所有未完成的跨流事件（synchronize_and_free_events），并逐段 cudaFree —— 按前面引的 CUDA 文档，后者可能隐式同步；下一个 batch 开始时这些段又得重新 cudaMalloc，同样可能同步。每个 batch 来一遍，CPU 提前排队的余量就被反复打掉。这种代码看起来"很负责"，实际上很可能拖慢训练。

社区里还能见到一种更夸张的写法：在每个 forward 后调一次 gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize()。这三连发把"CPU 提前排队、GPU 埋头跑"的流水线彻底打断 —— 每一步都要等 GPU 排空，GPU 因此在每次同步点后重新爬坡，用户却常常百思不得其解为什么"清理资源"反而变慢。这种"加清理一定更稳"的反直觉是 caching allocator 设计的代价 —— 它要求用户克制，相信 PyTorch 自己管得很好。

所以在训练循环里不要例行调用 empty_cache，让 caching allocator 自己管理。如果真的 OOM，调 expandable_segments:True 或者 max_split_size_mb 等配置 —— 何况 allocator 在 cudaMalloc 失败时本就会自己释放缓存段再重试（§4.7.4 的 cudaMalloc retries 数的就是这件事）。

需要和别的进程、或者 PyTorch 池子之外的显存使用者（例如 NCCL 自己申请的通信缓冲）共用一张卡时，empty_cache 是有意义的 —— 它把整段空闲的缓存段归还驱动，让它们能用上。但它解决不了碎片，也不提升 hot loop 性能。

4.7.3 PyTorch 的内存统计 API

PyTorch 提供一组内省 API 让你看清池子状态：

API	含义
torch.cuda.memory_allocated()	所有当前张量持有的字节数（"用户视角"）
torch.cuda.memory_reserved()	PyTorch 池子里的总字节数（已 cudaMalloc 的）
torch.cuda.max_memory_allocated()	历史峰值的 allocated
torch.cuda.max_memory_reserved()	历史峰值的 reserved
torch.cuda.memory_stats()	详尽 dict，包含 small/large 池各自的分配/释放/碎片化统计
torch.cuda.memory_summary()	人类可读的池子状态文字报告

判断"是否碎片化严重"的简单办法：

allocated = torch.cuda.memory_allocated()
reserved = torch.cuda.memory_reserved()
fragmentation_ratio = (reserved - allocated) / reserved
# 以下阈值是经验值, 且这个比值里也包含可复用的正常缓存:
# 0.0 ~ 0.1: 正常
# 0.2 ~ 0.4: 轻度碎片化, 可考虑 expandable_segments
# > 0.5: 严重碎片化, 应该排查工作流

工业级训练会把这些数字推到 Prometheus / Grafana 监控，一眼就能看出"训练吞吐下降是不是显存碎片化导致的"。

4.7.4 memory_summary() 输出解读

调用 torch.cuda.memory_summary() 会得到类似这样的输出（节选）：

|===========================================================================|
|                  PyTorch CUDA memory summary, device ID 0                 |
|---------------------------------------------------------------------------|
|            CUDA OOMs: 0            |        cudaMalloc retries: 12        |
|===========================================================================|
|        Metric         | Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed  |
|---------------------------------------------------------------------------|
| Allocated memory      |   8203 MB  |  12502 MB  | 145210 MB  | 137007 MB  |
| Active memory         |   8203 MB  |  12502 MB  | 145210 MB  | 137007 MB  |
| Requested memory      |   8200 MB  |  12498 MB  | 145202 MB  | 137001 MB  |
| GPU reserved memory   |  16384 MB  |  16384 MB  |   16384 MB |        0 B |
| Non-releasable memory |   8181 MB  |  12378 MB  |  121432 MB |  113251 MB |
| Allocations           |     12345  |     45678  |    856321  |    843976  |

几个关键指标：

• cudaMalloc retries（num_alloc_retries）：cudaMalloc 失败后、先释放全部可释放的缓存段再重试的次数（c10/cuda/CUDACachingAllocator.cpp:3285）。每一次 retry 都要同步等待事件并逐段 cudaFree，是一次明显的性能损失
• Allocated vs GPU reserved：差值是"被池子持有但未给用户"的字节数 —— 其中既有可复用的缓存，也有碎片
• Non-releasable memory（inactive_split_bytes）：被切分的段里那些空闲、但因同段还有占用而无法 cudaFree 的字节数 —— 这是 empty_cache 解决不了的

读懂这份 summary，是诊断 OOM 的基本功（§4.13 有完整的诊断案例）。

4.8 expandable_segments:True：v2.1+ 的杀手级配置

PyTorch v2.1（2023 年下半年发布）引入了一个新模式：expandable_segments（引入它的是 PR #96995 "Expandable blocks in allocator"；到 v2.11 官方文档仍标注为 experimental，默认关闭）。它彻底改变了 caching allocator 处理碎片化的方式。

启用方式：

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python train.py

或在代码里：

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

4.8.1 它做了什么

传统模式下，每个段都是一次 cudaMalloc 拿来的固定大小的地址区间，段与段之间无法合并。源码里的 Note [Expandable Segments]（c10/cuda/CUDACachingAllocator.cpp:275 起）举的例子是 batch size 从 N 变成 N+1：一部分 (N+1) 尺寸的张量还能塞进旧的 N 尺寸段，但塞不满，段尾留下用不上的空闲"薄片"；另一部分只能 cudaMalloc 新段。显存不够时，这些散落在各段段尾的空闲薄片没法拼起来用 —— 空闲总量够，却 OOM。

expandable_segments 用 CUDA Virtual Memory Management API（cuMemMap / cuMemUnmap / cuMemCreate，CUDA 10.2+）做：

1. 一次性预留一大段虚拟地址空间（cuMemAddressReserve，大小足以覆盖整卡显存，不占物理显存）
2. 按需把物理显存页"映射进来"（大池按 20 MiB、小池按 2 MiB 为一个 chunk）
3. 需要回收时把完全空闲的物理页"取消映射"，但虚拟地址段保留
4. 每个 stream 的分配都在同一个不断增长的段里平铺，虚拟地址连续，碎片少得多

graph TB
    subgraph Traditional["传统 cudaMalloc 模式"]
        T1["batch N: 大张量各自 cudaMalloc 一个段"]
        T2["batch N+1: 部分张量塞进旧段, 但放不满, 段尾留下空闲碎片"]
        T3["其余张量只能再 cudaMalloc 新段"]
        T4["显存余量不足 → 各段段尾的碎片彼此不相邻, 拼不起来 → OOM"]
    end
    subgraph Expandable["expandable_segments 模式"]
        E1["每个 stream 一个段: 预留足以覆盖整卡显存的虚拟地址 (不占物理)"]
        E2["需要时 cuMemCreate + cuMemMap 物理页, 追加到段尾"]
        E3["batch N+1 的张量在同一段里依次平铺"]
        E4["OOM 前或 empty_cache 时, 完全空闲的物理页可 unmap 还给 CUDA"]
        E5["碎片少得多 → 更容易凑出所需空间"]
    end
    T4 -.->|fails| OOM[OOM]
    E5 -.->|succeeds| OK[OK]
    style OOM fill:#fee2e2,stroke:#ef4444
    style OK fill:#dcfce7,stroke:#22c55e

收益的形状很好理解：碎片化的本质是"空闲字节总量够、但没有一段连续的够大"，而 expandable segments 把"连续"这件事挪到了虚拟地址层面 —— 物理页可以东一块西一块，用户看到的地址仍然连续。所以它对变长序列、动态 shape、峰值起伏大的训练收益最明显，对形状恒定、每步分配模式完全一样的训练几乎没差别。判断自己属于哪一类，看 memory_summary() 里 Non-releasable memory 与 cudaMalloc retries 两个数字（§4.7.4）。

4.8.2 expandable_segments 的内部机制

打开源码 c10/cuda/CUDACachingAllocator.cpp:368 的 ExpandableSegment（这一份被 PYTORCH_C10_DRIVER_API_SUPPORTED 包着；:783 还有一份"不支持时直接报错"的桩），末尾的字段是：

struct ExpandableSegment {
  // ...
  c10::DeviceIndex device_;
  std::optional<cudaStream_t> stream_;
  CUdeviceptr ptr_{};                 // 虚拟地址段起点
  size_t segment_size_;               // 每个 chunk 的粒度
  size_t mapped_size_;                // 已映射的物理字节数
  size_t max_handles_;                // 虚拟段能容纳多少个 chunk
  struct Handle {
    CUmemGenericAllocationHandle handle;
    std::optional<std::variant<int, CUmemFabricHandle>> shareable_handle;
  };
  std::vector<std::optional<Handle>> handles_;   // 每个 chunk 一个句柄，nullopt = 未映射
  std::vector<c10::DeviceIndex> peers_;          // 还要映射到哪些 peer device
};

handles_ 这个 vector<optional<Handle>> 就是"哪些 chunk 已经 backed 了物理页"的位图 —— nullopt 表示这一格只有虚拟地址、没有物理页。关键操作是两个：

• map(SegmentRange)（:399）：cuMemCreate 出物理句柄，再 cuMemMap 到虚拟段的对应偏移
• unmap(SegmentRange)（:500）：cuMemUnmap + 释放句柄，但不释放虚拟地址

普通的 free 只把块放回池子；等到 release_cached_blocks（empty_cache() 或 cudaMalloc 失败后的重试路径）时，release_blocks 才对 expandable segment 里的空闲块调 unmap_block()，把其中完全空闲的 chunk unmap、物理页归还 CUDA —— 但虚拟地址仍然挂在池子里，等待下次复用；只有整个段都空了，才连段带虚拟地址一起释放（release_expandable_segment）。这就是 expandable 模式能让"看起来连续的虚拟空间"在物理上有"洞"的根本机制。

4.8.3 何时不要开 expandable_segments

虽然这个模式很香，但它有几个真实限制需要警惕：

• 走的是 CUDA Driver API，不是所有构建都有 —— 源码里这一整块被 #if !defined(USE_ROCM) && defined(PYTORCH_C10_DRIVER_API_SUPPORTED) 包着，条件不满足时编到的是 :783 那份直接 TORCH_INTERNAL_ASSERT(false, "expandable segment not supported") 的桩。 不过在没有 PYTORCH_C10_DRIVER_API_SUPPORTED 的构建上，CUDAAllocatorConfig::expandable_segments() 会先打一次 "expandable_segments not supported on this platform" 警告并返回 false（c10/cuda/CUDAAllocatorConfig.h:34），配置等于没开
• 跨进程共享是另一条路径 —— 普通 cudaMalloc 的内存靠 cudaIpcGetMemHandle 共享，而 expandable segment 得走 ExpandableSegment::share() / fromShared() 这套自己的 handle 传递（源码里的 ShareHeader / shareable_handle 就是干这个的）。所以凡是依赖 CUDA IPC 的路径，都要确认这条链路在你的版本上通 —— v2.11 的 getBaseAllocation() 里仍有一条 TORCH_CHECK，报错信息就是 "Tensors allocated with expandable_segments:True cannot be shared between processes"，并建议在数据加载 worker 里临时关掉它
• Profiler 的内存 trace 解释复杂度上升 —— "虚拟地址有了但物理页没 map"这种中间状态，在快照可视化里需要额外理解

实战建议：先在单机单卡训练上开启 expandable 看效果，确定有收益再推到分布式训练。第 21 章 §21.41 的调优流程图里也有"碎片化 → 开 expandable_segments"这一分支。

4.9 CUDA Graph 集成：PrivatePool 隔离

CUDA Graph 是一种把 GPU 操作记录成静态 graph 然后整体重放的机制。它能极大降低 CPU launch 开销，但有一个硬约束：graph 里固化了所有显存地址，重放时这些地址必须仍然指向"正确的"内存。

如果 caching allocator 在 graph 期间分配/释放显存，下次重放时这些地址可能已经被分给别的张量 —— 数据被覆盖。

PyTorch 的解法：手工 torch.cuda.graph 捕获时，每个 CUDA graph 默认关联一个 PrivatePool（torch.compile 的 cudagraph_trees 则让同一设备上的所有 graph 共用一个池，见第 15 章）（c10/cuda/CUDACachingAllocator.cpp:970 起的 struct PrivatePool）。源码注释（:129-132）里把这套交互写成四个回调 notifyCaptureBegin / notifyCaptureAboutToEnd / notifyCaptureEnded / notifyCaptureDestroy，但这是早期的命名；v2.11 里 aten/src/ATen/cuda/CUDAGraph.cpp 实际调用的是三个函数：capture_begin 调 beginAllocateToPool，capture 开始后该流上的分配都走 private pool；capture_end 调 endAllocateToPool，切回主池；graph 析构时调 releasePool。Private pool 要等 releasePool 之后才可能被释放，保证地址稳定。

源码注释 Note [Interaction with CUDA graph capture]（c10/cuda/CUDACachingAllocator.cpp:103 起）专门解释了这套机制：

Because capture bakes in memory addresses, the memory used during capture must be available for the graph to use during replay. ... DeviceAllocator satisfies allocations from a graph-private memory pool during capture, and doesn't begin cudaFreeing those addresses until the graph is destroyed.

第 15 章 torch.compile 章会拆 CUDA Graph 与 Inductor 的协同；那里你会看到 private pool 是怎么和编译后的代码联动的。

4.9.1 PrivatePool 的生命周期

PrivatePool 不是简单的"另开一个池子"，它有完整的生命周期回调：

stateDiagram-v2
    [*] --> Idle
    Idle --> Capturing: beginAllocateToPool<br/>(capture_begin)
    Capturing --> Captured: endAllocateToPool<br/>(capture_end)
    Captured --> Replaying: graph.replay()
    Replaying --> Replaying: 多次 replay 复用同样地址
    Captured --> Destroyed: releasePool<br/>(graph 析构)
    Destroyed --> [*]: 下次 release_cached_blocks 时<br/>释放已空闲的段

注意一个细节：graph 销毁前，private pool 永远不释放。即使你的 graph 只 replay 了一次，私有池里那些被 graph 固化的地址也要等 CUDAGraph 对象真正析构后才能释放；而且析构也只是让池子的 use_count 减到 0、进入"可释放"名单（releasePool，c10/cuda/CUDACachingAllocator.cpp:2685）—— capture 期间分配的输出张量若还被你持有，对应的段照样还不回去，要等下一次 release_cached_blocks 发现它们空闲了才 cudaFree。这是为什么频繁创建 / 销毁 CUDA graph 容易导致显存压力 —— 每个 graph 短暂的"私有池"都要占据它捕获时的 high-water mark。

torch.compile 的 mode 'reduce-overhead' 大量用 CUDA Graph，所以它在小模型上加速明显，但在显存紧张的大模型场景可能更容易 OOM。注意 mode='max-autotune' 同样会开 CUDA Graph；不想要 CUDA Graph 时用 'max-autotune-no-cudagraphs' 或默认模式，见第 15 章 §15.6 的模式表。

4.10 配置 PYTORCH_CUDA_ALLOC_CONF：大模型训练的救星

PYTORCH_CUDA_ALLOC_CONF 是大模型训练遇到 OOM 时的第一个救命武器。（v2.11 里通用的名字是 PYTORCH_ALLOC_CONF，PYTORCH_CUDA_ALLOC_CONF 作为兼容别名继续有效，见 c10/core/AllocatorConfig.cpp:14 起；本章沿用更常见的旧名。）常用配置：

配置	含义	何时用
expandable_segments:True	用 VM API 减少碎片	训练长序列、动态 shape 模型
max_split_size_mb:512	≥ 512 MB 的块不切分（默认不限）	想保护几个大 buffer 不被碎片化
garbage_collection_threshold:0.6	占用超过允许上限的 60% 时主动回收旧缓存段（默认 0，即关闭；须配合 per-process memory fraction，见 §4.10.2）	长跑训练显存爬高时主动收紧
roundup_power2_divisions:[256:1,1024:2,>:32]	按大小区间设不同的对齐	调优特定 workload 的碎片率
backend:cudaMallocAsync	改用 CUDA 自带的异步分配器（PyTorch 要求 CUDA 11.4+，默认 native）	部分场景比默认快

一个常见的 OOM 救急组合是：

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:512

这条命令在某些场景能让原本 OOM 的训练直接跑通。不过要知道两者叠加时的实际效果：开了 expandable segments 后，should_split()（c10/cuda/CUDACachingAllocator.cpp:3110）一律按"剩余 ≥ 512 字节就切"处理，max_split_size_mb 不再限制切分；只剩下 get_free_block()（同文件 :3140）里"小请求不拿超大块"那条约束还起作用。怎么读 torch.cuda.memory._dump_snapshot() 产出的可视化，见第 21 章 §21.18；用它对比不同配置下的显存曲线，再决定取舍。

4.10.1 max_split_size_mb 的实战价值

max_split_size_mb 是最容易被低估的一条配置。它告诉 caching allocator："任何大于这个值的块，都不要切分"。

为什么这是有用的？考虑大模型训练：weights 张量一旦分配几乎一直存活，KV cache 也长期占用。这些"长寿张量"往往很大（几百 MB 到几 GB）。如果 allocator 把它们所在的大块切分给短寿小张量用，长寿张量的释放就被绑死在那些小张量上——只要任何一个短寿张量还在，整个大块都不能 cudaFree。

设 max_split_size_mb:512 让大于 512 MB 的块独享，不会被切分给小张量"借用"。代价是这些大块结束后短寿张量找不到合适块时要走新 cudaMalloc。但对长跑训练，避免长寿张量与短寿张量纠缠 比省几次 cudaMalloc 重要得多。

4.10.2 garbage_collection_threshold 的工作机制

这个配置告诉 allocator：当它从 CUDA 拿到的显存总量超过"允许上限 × 阈值"（如 0.6）时，主动回收旧的缓存段，而不是等 cudaMalloc 失败才 release_cached_blocks。

源码里有一个容易踩的前提："允许上限"指的是 allowed_memory_maximum，只有调用过 torch.cuda.set_per_process_memory_fraction()（或配 per_process_memory_fraction）且比例小于 1 时才有值。malloc() 和 get_free_block() 都先判断 allowed_memory_maximum.has_value()，没设比例时这条 GC 路径根本不会走（c10/cuda/CUDACachingAllocator.cpp:1450 起、:2107）。触发时，garbage_collect_cached_blocks() 只在 large pool 里挑未切分、非 expandable、"年龄"（多少次 get_free_block 调用没被复用）不低于平均值的缓存段 cudaFree，直到回到阈值以下（:3209）。

为什么需要主动 GC？因为 PyTorch 的"free 路径"是惰性的 —— 块进入 unallocated 状态后会留在池子里等待下次复用，直到 cudaMalloc 失败才被强制 reclaim。官方文档的说法是，它用来避免代价高昂的 "sync-and-reclaim-all"（release_cached_blocks），对延迟敏感的应用尤其有用；而且它不做同步，开销比 release_cached_blocks 小。

实战上，长跑训练显存爬高时可以试 garbage_collection_threshold:0.6，但记得同时设 per-process memory fraction，否则它不生效。

4.10.3 backend:cudaMallocAsync 是新选择

CUDA 11.2 起提供了 cudaMallocAsync API —— 分配和释放按 stream 排序，不像 cudaMalloc / cudaFree 那样隐式同步。PyTorch 从 v2.0 起提供对应的 backend（PR #82682），v2.11 要求 CUDA 11.4 或更新（c10/cuda/CUDAAllocatorConfig.cpp:51 起）。

启用：PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync

它的优势：

• 把 caching 工作交给 CUDA driver，PyTorch 池子简化
• 在某些 NVIDIA 驱动版本上比 PyTorch 自家 caching allocator 略快

代价：

• 不能用 PyTorch 的 memory_snapshot 工具诊断：CUDAMallocAsyncAllocator::snapshot() 直接报错 "Calling snapshot with backend:cudaMallocAsync is not meaningful"（c10/cuda/CUDAMallocAsyncAllocator.cpp:786）
• 在某些工作负载上比默认 allocator 略慢（特别是大量小张量场景）
• expandable_segments、max_split_size_mb、roundup_power2_divisions、garbage_collection_threshold 都是 native allocator 的机制，对它不起作用（后三者同时出现时会有 "backend:cudaMallocAsync ignores ..." 警告）

默认 backend 仍是 native，也就是本章讲的这套 caching allocator。

实战建议：在你的工作负载里实测两个 backend 的吞吐差异，而不是凭网上的"听说"做选择。一行环境变量就能切换，对照实验成本很低。

4.10.4 各配置的优先排查顺序

在我的经验里，遇到 OOM 的排查顺序是：

1. 先看 memory_summary()，确认 reserved/allocated 比 + cudaMalloc retries 数字
2. 比值明显偏高（经验上 > 1.3）或 retries 很多 → 优先开 expandable_segments:True
3. 仍 OOM → 加 max_split_size_mb:256（或更小，根据模型权重大小定；下限是 large_segment_size_mb，默认 20）
4. 长跑训练再加 garbage_collection_threshold:0.6（配合 per-process memory fraction）
5. 仍 OOM → 上 snapshot viewer 找具体张量

按这个顺序走，相当一部分由碎片引起的 OOM 不需要改算法或缩 batch size 就能缓解；若 allocated 本身就逼近显存上限，那就不是 allocator 配置能解决的了。

4.11 显存快照：可视化诊断的利器

torch.cuda.memory._record_memory_history() 从 v1.13 起就有，配套的 _dump_snapshot() 和下面示例里 max_entries= / enabled=None 这套写法是 v2.1 起的（PR #97406）。它们记录每一次 alloc/free 的时间序列与调用栈，dump 成 .pickle 文件后可以用官方 viewer（pytorch.org/memory_viz）可视化：

import torch

torch.cuda.memory._record_memory_history(max_entries=100000)

# 你的训练代码
train_one_step()

torch.cuda.memory._dump_snapshot('snapshot.pickle')
torch.cuda.memory._record_memory_history(enabled=None)

把 snapshot.pickle 拖进 viewer，你能看到：

• 每一段显存的"出生"时刻、调用栈、大小
• 哪些段被复用、哪些段一直存活
• 总显存占用随时间的曲线
• 碎片化随时间的演变

这是诊断 OOM、内存泄漏、碎片化的事实标准工具。第 21 章会专门有一节讲它。

4.11.1 一个可视化的真实示例

snapshot viewer 里你会看到一张"显存时间线"：

• 横轴是时间，纵轴是显存用量
• 每段彩色矩形是一个 alloc，颜色随源代码位置变化
• 矩形左右边界是 alloc / free 的时间点
• 矩形高度是 size

健康的训练 loop 看起来是 大量同一颜色的矩形规律地起伏（每个 batch 都有相似的 alloc/free 模式）。碎片化严重的训练看起来像 瑞士奶酪 —— 大量小矩形把空间钉得满满当当但中间布满"洞"。

snapshot viewer 还能定位"持续生长"的 alloc：如果某个特定 stack trace 对应的 alloc 数量随时间一直增加，那就是显存泄漏 —— 通常是某段 Python 代码持有引用没释放。这种泄漏在 PyTorch 训练里出乎意料地常见：autograd 的反向图被某个对象长期持有、callback 闭包捕获了张量、缓存的中间结果忘了清理 —— 都会让 snapshot 出现"持续生长"。

可视化工具的最大价值不是看绝对数字，而是 看模式。一旦把模式认清，OOM 的根因就能用经验快速定位。

4.11.2 PyTorch 与 OS swap：一个有趣的不交集

CPU 内存可以 swap 到磁盘 —— 不够用时 OS 把不活跃的页换出去。但 GPU 显存不能 swap：CUDA 没有"页换出去"的标准机制，物理显存不够就是不够。

这听起来像 PyTorch 的劣势，但它实际是设计前提：caching allocator 可以安全地把空闲块留在池子里，不必担心 OS 把它们偷走。CPU 端的内存池要小心 OS 的 page-out 行为，GPU 端的池子完全独享显存生命周期。

例外：CUDA 的 Unified Memory（配合 cudaMemAdvise 的 cudaMemAdviseSetPreferredLocation 等提示）理论上能让 GPU"swap"到 CPU 内存。但 PyTorch caching allocator 默认不用 unified memory —— 页迁移可能让 kernel 突然 stall，性能难以预测。第 18 章 FSDP 会讲一种"用 CPU offload 模拟 swap"的策略，那是受控的、按算法明确触发的。

4.12 跨专栏关联

• 《vLLM 推理内核深度解析》专栏第 4 章 PagedAttention：vLLM 的 KV Cache 也是池化分配，但是用 fixed-size block + PageTable，不切分。设计权衡完全不同 —— 你可以对照看"caching allocator vs paged allocator"两条路
• jemalloc / tcmalloc / Linux SLUB：这些通用分配器的 size-class 分流思想，是 PyTorch 两池子设计的"远房亲戚"；把 jemalloc 的 arena/bin 结构和本章的 small/large BlockPool 对照读，会发现解决的是同一类问题

4.13 一个真实案例：训练突然 OOM 怎么诊断

讲一个我从社区 issue 见过的典型场景：一个训练任务跑了 10 小时后突然 OOM，前面 9.5 小时显存稳定。怎么诊断？

第一步：看 cudaMalloc retries。memory_summary() 里如果 retries 数字很大（经验上 >100），说明 allocator 反复遇到 cudaMalloc 失败、只能清空缓存再重试 —— 池子里的空闲块派不上用场，往往意味着碎片化严重。

第二步：看 reserved vs allocated 的差。如果 reserved 比 allocated 大很多（比如 reserved=24GB / allocated=12GB），50% 是碎片或被池占着的空闲块。

第三步：开 _record_memory_history 跑一段时间后 dump snapshot。在可视化工具里搜 "long-lived but small" 的张量 —— 它们就是切分大块的"钉子"。常见嫌疑犯：

• 优化器的 momentum / variance（每个 weight 一份）
• LayerNorm 的中间统计量
• DDP 的 gradient bucket

第四步：上对应的 fix。常见解法：

• 加 expandable_segments:True 大幅减少碎片
• 用 torch.cuda.set_per_process_memory_fraction(0.95) 给 PyTorch 的池子设上限，为池子之外的显存使用者（CUDA context、NCCL 缓冲等）留出余量；它也是 garbage_collection_threshold 生效的前提（§4.10.2）
• 重新设计训练流程把"长寿小张量"集中分配（比如统一用 fused 优化器）

第 21 章 §21.41 的调优流程图把这条逻辑收成了其中的 OOM 分支。

4.14 caching allocator 的设计哲学一句话

整章看完，PyTorch 的 caching allocator 设计哲学可以用一句话概括：

把"昂贵的硬件操作（cudaMalloc）"换成"廉价的用户态簿记（块管理 + 链表合并）"，用空间（池子里的空闲块）和复杂度（split / coalesce / stream events）换时间（不触发 device sync）。

这与操作系统内核里 SLUB 分配器、jemalloc 等通用分配器的思想一脉相承。区别在于 PyTorch 多了 stream-aware 与 CUDA Graph 支持这两个深度学习特有维度。

4.15 几条记账意识

关于"显存什么时候真正被释放"，整理一组心智模型让你训练时少踩坑：

• del tensor：Python 端 PyObject refcount -1，到 0 时 C++ TensorImpl 析构，调 StorageImpl 的 deleter，deleter 通常只是把 block 还给 caching allocator 池子，不调 cudaFree
• tensor = None：等价于 del 当前引用
• 离开 Python 作用域：函数返回时局部变量自动 del
• tensor.fill_(0)：不释放内存，只清零内容
• torch.cuda.empty_cache()：尝试把 PyTorch 池子里的整段空闲的缓存段还给 CUDA 驱动，不影响用户持有的张量
• gc.collect()：强制 Python GC，捕获循环引用 —— 这能释放某些被 autograd 反向图保留住的张量
• 进程退出：进程的 CUDA 上下文销毁，全部显存交还驱动，不管 PyTorch 池子状态

最常被混淆的两条：

1. del tensor 不一定降 nvidia-smi 的数字，因为 caching allocator 还把内存留着
2. empty_cache 不释放任何"用户还在持有的张量"，它只动 PyTorch 池子里的空闲部分

把这组心智模型记牢，你就能精确地推理"这一步显存应该是多少"，而不是事后看 nvidia-smi 困惑。

4.16 一个练习：手动制造碎片化

import torch

torch.cuda.empty_cache()
print(f'Initial: allocated={torch.cuda.memory_allocated()/1e6:.1f}MB '
      f'reserved={torch.cuda.memory_reserved()/1e6:.1f}MB')

# 制造碎片
tensors = []
for i in range(1000):
    tensors.append(torch.randn(256, device='cuda'))  # 1 KB 张量

# 删除间隔元素，留下"瑞士奶酪"
del tensors[::2]

print(f'After del even: allocated={torch.cuda.memory_allocated()/1e6:.1f}MB '
      f'reserved={torch.cuda.memory_reserved()/1e6:.1f}MB')

# 试着调 empty_cache，观察 reserved 几乎不变
torch.cuda.empty_cache()
print(f'After empty_cache: reserved={torch.cuda.memory_reserved()/1e6:.1f}MB')

# 真要释放，必须把所有引用都删了
del tensors
torch.cuda.empty_cache()
print(f'After del all + empty_cache: reserved={torch.cuda.memory_reserved()/1e6:.1f}MB')

跑一次，亲眼看到 empty_cache 在碎片化场景里的"无能"：1000 个 1 KB 张量都落在同一个 2 MiB 的 small 段里，删掉一半后这个段仍被另一半钉着，reserved 不会降；全部删掉后才能还回去（按源码判断、未实跑）。有兴趣可以再用 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 跑一遍对照：small 池的物理页粒度也是 2 MiB，这个例子里半满的那一页照样 unmap 不掉。

4.17 横向看：其他 ML 框架的 GPU 内存管理

caching allocator 不是 PyTorch 独家。把主流框架放在一起对比能看出设计思想的演进：

框架	内存管理策略	关键特点
PyTorch	Caching allocator + 两池子 + expandable_segments	灵活、可调、用户态控制
TensorFlow 2.x	BFC (Best-Fit with Coalescing) Allocator	与 PyTorch 类似的 best-fit + coalesce，但没有两池子分流
JAX	启动时预分配大块显存 + XLA 编译期规划 buffer	默认在首个 JAX 操作时预分配 75% 显存（JAX 官方文档 GPU memory allocation），之后在这块池子里分配；XLA 编译期为每个计算规划 buffer 复用
vLLM	PagedAttention BlockTable + 动态 block 池	固定 block size, 显式 page 管理（受操作系统虚拟内存启发）
TensorRT-LLM	静态规划的 Workspace + Tensor 内存复用	编译期规划好每个 tensor 的内存位置，几乎零碎片

各方案的取舍：

• PyTorch / TF：动态、灵活，付碎片化代价
• JAX / TensorRT：编译期规划，碎片少，但动态形状的灵活性差一些
• vLLM：在动态批处理场景下用 page 抽象解决"变长 KV cache"的碎片问题

理解这些差异，你就能在选型时做出有依据的决定 —— 训练 + 多变 batch / 序列长度 选 PyTorch；推理 + 固定 shape 选 TRT；研究 / 高吞吐推理选 vLLM。

值得多说一句的是 vLLM 的 PagedAttention 思路。它把 KV Cache 切成固定大小的 page（默认 16 个 token），用一张 BlockTable 把逻辑序列映射到物理 page —— 完全模仿了 OS 的虚拟内存机制。这种设计在变长序列推理里几乎消除碎片，但代价是每次 attention 计算都要走一次 indirect addressing。PyTorch 的 caching allocator 与 vLLM 的 PagedAttention 是两套不同的工具，分别为"训练时各种 shape 的张量"和"推理时变长 KV cache"设计。它们在 PyTorch 里可以共存 —— vLLM 的 KV block 仍然是用 caching allocator 一次性分配的，PagedAttention 只是在那块大内存内部做 logical paging。

下一章拆 dispatcher：PyTorch 多分派的核心引擎，第 1 章 §1.4 与第 3 章 §3.5 提过的 DispatchKey 现在终于要派上用场。