3、PyTorch缓存分配器：核心原理与实战解读
说到显存碎片，很多人第一反应是「我手动调一下分配策略不就行了？」
嗯，没那么简单。
PyTorch 的缓存分配器（Caching Allocator）才是真正的主角。它决定了你的显存是高效复用，还是越用越碎。
我个人习惯把分配器比作「显存管家」。它不直接跟 CUDA 驱动打交道，而是先申请一大块显存，然后自己管理。说白了，就是搞了个「内存池」。
3.1 为什么需要缓存分配器？
直接调用 cudaMalloc 和 cudaFree 有什么问题？
我遇到过好几次：每次分配都要走内核态，开销巨大。而且频繁分配释放，碎片率飙升。
PyTorch 的做法是：
- 先一次性向 GPU 申请大块显存（比如 2GB）
- 内部切成小块，按需分配
- 释放时并不还给 GPU，而是留着复用
这样既减少了系统调用，又能控制碎片。
核心思想： 用空间换时间，用池化换效率。
3.2 Stream 与 Block 管理
PyTorch 的分配器内部维护了多个 Stream（流）。每个 Stream 对应一个 CUDA stream。
每个 Stream 下面，又挂着一堆 Block（缓存块）。
Block 是分配器的最小管理单元。每个 Block 有：
- 起始地址
- 大小
- 所属 Stream ID
- 状态（见下文）
你想想看，如果两个操作在不同的 Stream 上，分配器会尽量让它们使用不同的 Block 池。这样能避免跨 Stream 的同步等待。
避坑指南： 我曾经在项目中同时用了多个 Stream，结果发现显存占用暴涨。后来一查，是分配器为每个 Stream 都保留了独立的缓存池。解决办法是：尽量复用 Stream，或者手动设置
torch.cuda.set_stream。
3.3 缓存块的状态机
每个 Block 在生命周期里会经历几个状态。我画了一张图，帮你理清逻辑：
状态流转其实不复杂：
- 空闲 → 已分配：分配器从空闲池里拿一个 Block 给你
- 已分配 → 分裂：如果 Block 太大，会切成两块，一块给你，一块留着
- 已分配 → 回收：你释放了，Block 回到分配器手里
- 回收 → 合并：相邻的空闲 Block 会合并成更大的块
- 合并 → 空闲：合并完成后，放回空闲池
注意： 分裂和合并是碎片治理的关键。分裂太多会产生小碎片，合并太慢会导致大块不够用。我曾经在一个训练任务里，就是因为频繁分裂产生了大量 128KB 的小块，最后大张量分配失败。
3.4 分配与释放的日志解读
PyTorch 提供了详细的分配器日志。开启方式很简单：
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export TORCH_CUDA_ALLOC_LOG_LEVEL=INFO然后跑你的训练脚本，你会看到类似这样的输出：
[INFO] Allocator: block 0x7f8c4a000000 (size 1024 MB) allocated to stream 0
[INFO] Allocator: block 0x7f8c4a000000 freed, returning to pool
[INFO] Allocator: merging block 0x7f8c4a000000 with adjacent free block
[INFO] Allocator: new free block size 2048 MB我来带你逐行解读：
| 日志内容 | 含义 | 我的建议 |
|---|---|---|
allocated to stream 0 |
这个 Block 分配给了默认 Stream | 如果看到大量跨 Stream 分配，检查你的多流使用 |
freed, returning to pool |
释放了，但没还给 GPU | 正常行为，不用慌 |
merging block ... |
正在合并相邻空闲块 | 合并频繁说明碎片多，可以考虑调大 max_split_size_mb |
new free block size ... |
合并后的新块大小 | 如果合并后还是很小，说明碎片严重 |
实战技巧： 我习惯在训练脚本里加一行
torch.cuda.memory_summary()，每隔几步打印一次显存状态。这样能实时看到分配器的行为，比事后看日志更直观。
另外，如果你发现日志里频繁出现 split 和 merge，说明你的显存分配模式不太理想。可以试试：
- 增大
max_split_size_mb（默认 512 MB） - 使用
torch.cuda.empty_cache()手动触发回收（但别频繁调用） - 调整 batch size，让张量大小更规整
嗯，关于分配器的原理和日志解读，就先聊到这儿。
记住一句话：分配器不是万能的，但不懂分配器是万万不能的。