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

“Memory allocation is one of the few things in deep learning that, when it goes wrong, takes down the whole training run.”

—— PyTorch core team note in c10/cuda/CUDACachingAllocator.cpp 顶部注释

4.1 为什么不直接 cudaMalloc

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

打开 NVIDIA CUDA 文档你会读到一句让人惊掉下巴的话：

cudaMalloc and cudaFree are synchronous with respect to the device. They cause an implicit synchronization of all streams on the device.

也就是说 —— 每次 cudaMalloc 都要等当前所有 GPU 工作完成，然后才能返回。这是为了保证新分配的内存不会被任何还在跑的 kernel 写到（CUDA 没有异步分配器的内存别名保证）。

对一个朴素的 ML 框架，这个设计完全致命：

# 每个张量创建都是一次 GPU 全停同步
a = torch.randn(1024, 1024, device='cuda')   # cudaMalloc → device sync
b = torch.randn(1024, 1024, device='cuda')   # cudaMalloc → device sync
c = a + b                                     # 还有这个 op 自身的临时显存分配
del a                                         # cudaFree → device sync
del b                                         # cudaFree → device sync

一次 forward + backward 可能创建上千个临时张量。如果每次都触发 device sync —— GPU 利用率会从 95% 直接掉到 30% 以下。

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

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

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

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

4.1.2 一个量化的代价

cudaMalloc 同步开销有多大？粗略数据：

• 在空载 GPU 上：单次 cudaMalloc 大约 50-200 微秒
• 在繁忙 GPU 上（kernel 队列里有任务）：单次可达 1-10 毫秒（要等队列完成）
• 而真正的”分配字节”动作只需要几十纳秒

也就是说 —— cudaMalloc 的 99% 时间花在同步上，只有 1% 时间在真正分配。Caching allocator 的目标就是把这 99% 的开销摊销掉：训练初期付几次 cudaMalloc 的 sync 代价，之后所有”分配”都从池子里 O(log n) 拿，完全不触发 GPU 同步。

这是为什么用 PyTorch caching allocator 的训练吞吐通常比”裸调 cudaMalloc”的实现快 2-3 倍 —— 不是因为算子算得快了，而是因为分配/释放的隐式同步消失了。

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 用过（不只是分配 stream）。这是流间安全复用的核心
• prev / next：每次”切分”产生的子块用双向链表串起来。这让 释放时合并相邻空闲块（coalescing）成为 O(1) 操作
• event_count：只有 event_count = 0（所有相关 GPU 工作完成）时，块才能被回收

BlockPool 是 Block 的集合：

struct BlockPool {
    std::set<Block*, Comparison> blocks;        // 所有块，按 size 排序
    std::set<Block*, Comparison> unmapped;      // expandable_segment 下未映射的虚拟段
    const bool is_small;                        // 区分 small / large 池
    PrivatePool* owner_PrivatePool;             // CUDA Graph 私有池
    ...
};

blocks 是按 size 排序的红黑树 —— 分配时用 lower_bound 找”刚好够”的块，O(log n)。这是 best-fit 分配策略的标准实现。

4.2.1 best-fit vs first-fit 之争

经典分配器有两条路：

• first-fit：找第一个 ≥ size 的块。快（O(n) 平均 < O(log n) 平摊），但容易在池子前部留下大量”几乎合适但不够”的小碎片
• best-fit：找”恰好够”的最小块。每次分配都精准吃掉最小块，剩下的大块保持完整 —— 但搜索代价稍高，且如果”恰好够”留下极小残块也是浪费

PyTorch 选 best-fit 的原因是 GPU 显存太宝贵，不能用 first-fit 的”大概够就行”策略浪费整页。代价是每次分配多走几次红黑树节点。

而 BlockComparatorSize 不是简单的 size < —— 它先比 size、再比 stream（让同一 stream 的块优先复用，避免跨流同步成本）、再比 ptr 地址。这种”多键比较”让池子的内部顺序天然有利于”流亲和”复用模式。

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/>按 2 MiB 倍数 round"]

    SP --> CudaMalloc1["首次: cudaMalloc 2 MiB 大块"]
    LP --> CudaMalloc2["首次: cudaMalloc round(size, 2MiB)"]

    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 用三个阈值定义大小分流：

注意中间那一段 (1-10 MiB) 不做”打包”，而是按需 cudaMalloc。原因是：1-10 MiB 的请求频次不高、单次足够大，打包反而增加碎片。这是 PyTorch 对真实工作负载分布的经验总结。

4.4 分配流程：best-fit + split

完整的分配路径：

sequenceDiagram
    autonumber
    participant U as 用户调用 cudaMalloc(size)
    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 函数是切分操作的实现：把一个大块分成两个，前一个返回给用户，后一个放回池里，两者通过 prev/next 双向链表保持邻接关系。

4.4.1 split 的 trade-off

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

• 如果剩余部分小于阈值（默认 small 池 0.5 MiB、large 池 1 MiB），不切（避免产生过多小碎片）
• 如果当前 size 已经接近 max_split_size_mb（用户可配），也不切

这就是 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 还给 OS。但只要有任何一个子段还被持有 —— 整个链就无法合并回原始大小。

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

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

sequenceDiagram
    autonumber
    participant U as 用户 cudaFree(ptr)
    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: 后台周期性轮询事件
    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 的内存？

如果 PyTorch 立刻把 a 的内存还给池子让 c 用，但 s 流上的 a + 1 还没跑完 —— 数据竞争！ c 读到的是没被覆盖的 a 的旧值，但 b 读到的可能是 c 刚写入的新值。

PyTorch 的解法：recordStream 与 stream_uses。当一个张量被另一个流使用时，调用 tensor.record_stream(other_stream)，allocator 记录这个块的 stream_uses。释放时如果 stream_uses 包含多流，allocator 在每个使用过的流上插入 cudaEvent，把块标记为 “等待事件” 状态，event 全部完成后才真正回池。

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

第 17 章 DDP 与第 21 章 profiler 都会回到这个机制。

4.6.1 一个真实场景：DataLoader pin_memory 与跨流复用

DataLoader 的 pin_memory=True 模式会启动一个专门的”pinned memory worker stream”，把数据从 pinned CPU 内存异步搬到 GPU 默认流。这意味着：每个 batch 的输入张量都被两个流碰过（pinned worker stream + default stream）。

如果 PyTorch 没有 stream-aware 复用，下一个 batch 的输入张量分配可能拿到上一个 batch 还没用完的内存，直接覆盖训练数据。但因为 caching allocator 的 stream_uses 机制 + DataLoader 内部对张量的 record_stream 调用，这种 race 在用户代码层面”似乎从未发生过” —— 整个机制是透明的。

但如果你写自定义的多流代码（比如自己开 stream 做 prefetch），忘了 record_stream 就会触发非常隐蔽的”输出有时对、有时错”的 bug。诊断这种 bug 的标准方法：在怀疑的地方加 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_internal(block);    // 真正回池
            }
        }
    }
}

每隔一段时间（在某些 caching allocator 操作前会触发），PyTorch 轮询所有 pending event。完成的 event 让对应 block 的 event_count 减 1；event_count 归零时块才真正被合并并放回池子。

这个机制让 PyTorch 在跨流场景下既正确又高效：正确（不释放还在用的内存），又高效（不强制 device 全停）。代价是池子里有一部分块处在”等待事件”状态，暂时不可用。在多流密集场景这部分内存占比可能 5-10%，是 caching allocator 设计中”用空间换时间”的经典权衡。

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 的判断是 block.is_split() == false 才能释放（即 prev == nullptr && next == nullptr）。

为什么有这个限制？因为 cudaFree 必须配对 cudaMalloc —— 你不能 cudaFree 一个 cudaMalloc 大块的中间一段。所以 empty_cache() 只能把”还没被切过的整块”还给 OS。任何被切分过的池子，里面的子块即使全部空闲，也无法返还。

让我们看一个反例：

# 训练初期：
torch.cuda.empty_cache()                           # 假设清空
# 第一个 forward：
a = torch.randn(1, device='cuda')                  # 池子从 cudaMalloc 拿 2 MiB 大块, 切出 4 字节给 a
b = torch.randn(1, device='cuda')                  # 切走 4 字节
del a, b                                           # 子块还给池, 但 prev/next 还在
torch.cuda.empty_cache()                           # 想还回 OS, 但这 2 MiB 已被切过 → 不能 free

nvidia-smi 显示 PyTorch 进程仍然占着 2 MiB —— 即使所有子块都空闲。这就是 “显存看起来没释放但 PyTorch 已经放手” 这个常见困惑的源头。

4.7.1 那 empty_cache 什么时候才有用

只有 池子里有”完整未切分”的整块 时它才有效。具体场景：

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

但训练 hot loop 里的小张量分配/释放，empty_cache 几乎帮不上忙。调用它本身的开销（synchronize 整设备）反而成了拖累。

PyTorch 团队的官方建议：只在两种场景调用 empty_cache：

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

4.7.1 一个反复被踩的坑

社区 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 本身会强制做一次 device sync（等所有 pending kernel 完成才能查 free 块状态）。每个 batch 都同步 GPU 让 batch 之间的并行性彻底消失。这种代码看起来”很负责”，实际上是性能杀手。

PyTorch GitHub issue 里反复出现一种更夸张的反模式：在每个 forward 后调一次 gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize()。这三连发让训练吞吐直接掉到 1/5，但用户常常百思不得其解为什么”清理资源”反而变慢。这种”加清理一定更稳”的反直觉是 caching allocator 设计的代价 —— 它要求用户克制，相信 PyTorch 自己管得很好。

正确做法是 完全不调用 empty_cache，让 caching allocator 自己管理。如果真的 OOM，调 expandable_segments:True 或者 max_split_size_mb 等配置；只有在进程会被外部观察显存（如多任务共享 GPU）时才偶尔调一次 empty_cache 让 nvidia-smi 显示低一些。

第二种场景下 empty_cache 是有意义的 —— 它把”完整大块”归还系统，让其他进程能用。但对 hot loop 性能没帮助。

4.7.2 PyTorch 的内存统计 API

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

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

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.3 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：池子触发了多少次”再 cudaMalloc 一次”。每一次 retry 都是一次性能损失
• Allocated vs GPU reserved：差值就是”被池子持有但未给用户”的字节数 —— 即碎片
• Non-releasable memory：被切分过、即使 free 也无法 cudaFree 还给 OS 的字节数 —— 这是 empty_cache 解决不了的

读懂这份 summary，是诊断 OOM 的基本功。第 21 章会深入。

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

PyTorch v2.1（2023 末）引入了一个新模式：expandable_segments。它彻底改变了 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 一次性分配固定大小的连续物理 + 虚拟地址段。如果你需要 32 GB 但只有 28 GB 连续可用 —— OOM。

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

1. 一次性预留一大段虚拟地址空间（不占物理显存）
2. 按需把物理显存页”映射进来”
3. 释放时把物理页”取消映射”，但虚拟地址段保留
4. 下次分配时虚拟地址连续（不会碎片化），物理页按需分配

graph TB
    subgraph Traditional["传统 cudaMalloc 模式"]
        T1["第 1 次 alloc 2GB → cudaMalloc(2GB) 物理 + 虚拟连续"]
        T2["第 2 次 alloc 1GB → cudaMalloc(1GB) 与第 1 次不一定相邻"]
        T3["释放第 1 次的 2GB"]
        T4["新请求 3GB → 虚拟空间已碎片化, 找不到 3GB 连续段, OOM"]
    end

    subgraph Expandable["expandable_segments 模式"]
        E1["预留 100 GB 虚拟地址 (无物理)"]
        E2["第 1 次 alloc 2GB → 在虚拟段头部映射 2GB 物理页"]
        E3["第 2 次 alloc 1GB → 紧跟着映射 1GB 物理页"]
        E4["释放第 1 次 → unmap 那 2GB 物理页"]
        E5["新请求 3GB → 虚拟空间还在头部连续, 重新 map 物理页"]
    end

    T4 -.->|fails| OOM[OOM]
    E5 -.->|succeeds| OK[OK]

    style OOM fill:#fee2e2,stroke:#ef4444
    style OK fill:#dcfce7,stroke:#22c55e

实测数据（来自 PyTorch 官方 v2.1 release notes）：

• 训练 70B 模型：OOM 减少 30-50%
• 长序列训练：碎片化导致的吞吐损失从 15% 降到 < 3%

4.8.2 expandable_segments 的内部机制

打开源码 c10/cuda/CUDACachingAllocator.cpp 中 ExpandableSegment 类（约 1500-1800 行附近），核心字段：

struct ExpandableSegment {
    c10::DeviceIndex device_;
    cudaStream_t stream_;
    CUmemGenericAllocationHandle handle_;       // CUDA VM API 句柄
    void* ptr_;                                  // 虚拟地址段起点
    size_t max_handle_size_;                    // 虚拟段总大小
    size_t segment_size_;                        // 每个 chunk 大小（默认 2 MiB）
    std::vector<bool> mapped_;                  // 哪些 chunk 已映射物理页
    ...
};

关键操作：

• map(offset, size)：调 cuMemMap 把物理页映射到虚拟段的某个偏移
• unmap(offset, size)：调 cuMemUnmap 取消映射，但不释放虚拟地址
• extend(size)：扩大虚拟段（保持地址连续）

当一个 free 触发时，如果整段已经没有 mapped chunk，可以 unmap 把物理页归还系统 —— 但虚拟地址仍然挂在池子里，等待下次复用。这就是 expandable 模式能让”看起来连续的虚拟空间”在物理上有”洞”的根本机制。

4.8.3 何时不要开 expandable_segments

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

• 不兼容 torch.cuda.IntTensor.share_memory_() —— cuMemMap 创建的内存无法被 cudaIpcGetMemHandle 处理，跨进程共享会失败。这意味着 DataLoader 的 num_workers > 0 在某些场景会报错
• NCCL 集合通信在某些版本里有兼容问题 —— v2.1 早期版本的 NCCL 配合 expandable 时偶发 hang，社区在 v2.2 之后修复
• Profiler 的内存 trace 解释复杂度上升 —— “alloc 了一段但没真分配物理页”这种状态需要更复杂的可视化

实战建议：先在单机单卡训练上开启 expandable 看效果，确定有收益再推到分布式训练。

但 expandable_segments 不是默认开启的，因为它有几个限制：

• 不兼容 cudaIpcGetMemHandle —— 跨进程显存共享会失败
• 需要 NVIDIA Driver 530+ 与 CUDA 11.4+
• 在某些古老 GPU 上有性能回归

第 17 章 DDP 章会讲怎么诊断”该不该开 expandable_segments”。

4.9 CUDA Graph 集成：PrivatePool 隔离

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

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

PyTorch 的解法：每个 CUDA graph 关联一个 PrivatePool（c10/cuda/CUDACachingAllocator.cpp:973-1010 附近）。在 notifyCaptureBegin 回调时切换到 private pool，期间所有分配走 private pool；notifyCaptureEnded 后切回主池。Private pool 在 graph 销毁前不会释放，保证地址稳定。

源码注释（CUDACachingAllocator.cpp:106-140 附近）专门解释了这套机制：

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: notifyCaptureBegin
    Capturing --> Captured: notifyCaptureAboutToEnd
    Captured --> Replaying: graph.replay()
    Replaying --> Replaying: 多次 replay 复用同样地址
    Captured --> Destroyed: notifyCaptureDestroy<br/>(graph 析构)
    Destroyed --> [*]: 私有池真正释放

注意一个细节：graph 销毁前，private pool 永远不释放。即使你的 graph 只 replay 了一次、或者你已经 del graph，私有池里那些被 graph 固化的地址要等到 Python GC 真正销毁 CUDAGraph 对象后才能释放。这是为什么频繁创建 / 销毁 CUDA graph 容易导致显存压力 —— 每个 graph 短暂的”私有池”都要占据它捕获时的 high-water mark。

torch.compile 的 mode 'reduce-overhead' 大量用 CUDA Graph，所以它在小模型上加速明显，但在显存紧张的大模型场景反而会 OOM。第 15 章会讲 mode='max-autotune' 这个折中。

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

PYTORCH_CUDA_ALLOC_CONF 是大模型训练遇到 OOM 时的第一个救命武器。常用配置：

举个真实例子：训练 Llama-70B 时，常见 OOM 救命组合是：

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:512

这条命令在某些场景能让原本 OOM 的训练直接跑通。第 21 章 profiler 那章会讲怎么用 torch.cuda.memory._dump_snapshot() 的可视化工具找出最佳配置。

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：当显存使用率超过阈值（如 0.6）时，主动触发块回收而不是等 OOM 临门一脚才 release_cached_blocks。

为什么需要主动 GC？因为 PyTorch 的”free 路径”是惰性的 —— 块进入 unallocated 状态后会留在池子里等待下次复用，直到 OOM 才被强制 reclaim。在显存压力大但还没 OOM 的”中间状态”，主动回收能让池子保持松弛，避免突然 OOM 后再做完整 reclaim 的连锁反应。

实战经验：长跑训练（数小时以上）开 garbage_collection_threshold:0.6 通常比不开稳定。短跑训练（< 1 小时）一般不需要。

4.10.3 backend:cudaMallocAsync 是新选择

CUDA 11.2+ 提供了新的 cudaMallocAsync API —— 它本身就是异步的、不会触发 device sync。PyTorch 在 v1.11+ 提供了对这个 API 的实验性 backend。

启用：PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync

它的优势：

• 把 caching 工作交给 CUDA driver，PyTorch 池子简化
• 在某些 NVIDIA 驱动版本上比 PyTorch 自家 caching allocator 略快
• 与 NCCL、CUDA Graph 等的兼容性由 NVIDIA 驱动直接保证

代价：

• 不能用 PyTorch 的 memory_snapshot 工具诊断（NVIDIA 没暴露足够 introspection API）
• 在某些工作负载上比默认 allocator 略慢（特别是大量小张量场景）
• 不支持 expandable_segments（两者互斥）

PyTorch 团队的内部立场是 “默认仍用自家 caching allocator，但欢迎用户用 cudaMallocAsync 跑评估”。第 17 章分布式训练会再讨论这个选择。

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

4.10.4 各配置的优先排查顺序

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

1. 先看 memory_summary()，确认 reserved/allocated 比 + cudaMalloc retries 数字
2. 比值 > 1.3 或 retries > 100 → 优先开 expandable_segments:True
3. 仍 OOM → 加 max_split_size_mb:256（或更小，根据模型权重大小定）
4. 长跑训练再加 garbage_collection_threshold:0.6
5. 仍 OOM → 上 snapshot viewer 找具体张量

按这个顺序走，90% 的训练 OOM 都能解决，不需要改算法或缩 batch size。

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

PyTorch v1.13+ 提供了 torch.cuda.memory._record_memory_history() 与 _dump_snapshot() 一对工具。它们记录每一次 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.5 PyTorch 与 OS swap：一个有趣的不交集

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

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

例外：CUDA 11+ 引入了 cuMemAdvise(CU_MEM_ADVISE_SET_PREFERRED_LOCATION) 与 Unified Memory，理论上能让 GPU”swap”到 CPU 内存。但 PyTorch caching allocator 默认不用 unified memory —— 因为 unified memory 在大模型训练里性能不可预测（页迁移会突然 stall kernel）。第 18 章 FSDP 会讲一种”用 CPU offload 模拟 swap”的策略，那是受控的、按算法明确触发的。

4.12 跨书关联

• 《vLLM 内核探秘》第 4 章 PagedAttention：vLLM 的 KV Cache 也是池化分配，但是用 fixed-size block + PageTable，不切分。设计权衡完全不同 —— 你可以对照看”caching allocator vs paged allocator”两条路
• 《Tokio 异步运行时》第 X 章 内存池：Tokio 的 Vec reuse pool 与 PyTorch 的 caching allocator 思想类似，但 Tokio 不需要应付 stream 同步
• 《Linux io_uring》（未来书）：Linux 内核 SLAB 分配器的 size-class 思想是 PyTorch 两池子设计的”远房亲戚”

4.12.5 一个真实案例：训练突然 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) 让 OS 看到的”显存”留出余量
• 重新设计训练流程把”长寿小张量”集中分配（比如统一用 fused 优化器）

第 21 章会有完整的”OOM 诊断流程图”，囊括这条逻辑。

4.12.6 caching allocator 的设计哲学一句话

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

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

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

4.12.7 几条记账意识

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

• 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 池子里的”完整未切分块”还给 OS，不影响用户持有的张量
• gc.collect()：强制 Python GC，捕获循环引用 —— 这能释放某些被 autograd 反向图保留住的张量
• 进程退出：OS 回收所有显存，不管 PyTorch 池子状态

最常被混淆的两条：

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

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

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

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 在碎片化场景里的”无能”。然后改成 expandable_segments=True 跑同样的代码 —— reserved 的曲线会平滑得多。这个微练习能让本章所有抽象一次性内化。

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

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

各方案的取舍：

• 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.3 与第 3 章 §3.5 提过的 DispatchKey 现在终于要派上用场。