训练一个 step 里 PyTorch 会做成千上万次显存申请/释放 —— forward 的临时张量、activation、Autograd 保留的中间结果、梯度 buffer…… 如果每次都直接调 cudaMalloc / cudaFree, 性能会差到训练根本跑不动。caching allocator 就是为了把这条路绕开。

三个核心痛点

caching allocator 的三条策略

两个永远要分清的概念

源码在 c10/cuda/CUDACachingAllocator.cpp,是 PyTorch 仓库里 最难也最重要的几千行 C++ 之一。 不需要全读,但要在脑子里有清晰的两级模型:Segment 是 driver 那一级粗粒度块,Block 是应用面细粒度块, 二者通过有序双向链表串起来。

small / large pool — 1 MiB 的分水岭

为了减少碎片,allocator 把所有 Block 按大小分到两个独立的池。 界线是 1 MiB(kSmallSize), Segment 的默认尺寸也对应这个分类。

alloc / free 一次完整流程

// 伪代码:对照 c10/cuda/CUDACachingAllocator.cpp 阅读
Block* alloc(size_t size, cudaStream_t stream) {
  size = round_up(size);                              // 512B / 2MiB 对齐
  auto& pool = (size < 1_MiB) ? small_pool[stream] : large_pool[stream];

  // ① best-fit:在按 size 排序的 std::set 里找第一个 >= size 的 free block
  auto it = pool.lower_bound(BlockKey{size, stream});
  if (it != pool.end()) {
    Block* b = *it;
    pool.erase(it);
    // ② 命中后如果 block 显著大于请求,就把它切两半
    if (should_split(b, size)) {
      Block* tail = split(b, size);                  // tail 留在 free pool
      pool.insert(tail);
    }
    return b;
  }

  // ③ 没命中:向 driver 申请新 Segment
  Block* seg = cuda_malloc(segment_size_for(size));
  pool.insert(seg);
  return alloc(size, stream);                       // 递归 → 这次一定能命中
}

void free(Block* b) {
  // ④ 不还给 driver,只是回池
  auto& pool = pool_for(b);
  // ⑤ 合并:看左右相邻的 block 是不是 free,在同 Segment 内就 merge
  if (b->prev && b->prev->is_free()) b = merge(b->prev, b);
  if (b->next && b->next->is_free()) b = merge(b, b->next);
  pool.insert(b);
}

OOM 不是"显存不够"那么简单。 绝大多数生产 OOM 都是"显存够但拿不出连续块"(碎片)、 "用完没释放"(泄漏)、或者"峰值刚好超线"(瞬时)。 不同病因要不同药方,不能只看 nvidia-smi 拍脑袋。

三件套 + 各自的适用场景

第一步 — 读懂 memory_summary

import torch
print(torch.cuda.memory_summary(device=0, abbreviated=False))

|=============================================================================|
|                  PyTorch CUDA memory summary, device ID 0                   |
|-----------------------------------------------------------------------------|
|        CUDA OOMs: 0           |        cudaMalloc retries: 4                |    ← 重试 = 警讯
|=============================================================================|
|        Metric         | Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed  |
|---------------------------------------------------------------------------|
| Allocated memory      |  18.4 GiB  |  22.1 GiB  | 412.0 GiB  | 393.6 GiB  |
| Active memory         |  18.4 GiB  |  22.1 GiB  | 412.0 GiB  | 393.6 GiB  |
| Requested memory      |  18.3 GiB  |  21.9 GiB  | 411.2 GiB  | 392.9 GiB  |
| GPU reserved memory   |  23.8 GiB  |  23.8 GiB  |  23.9 GiB  | 100.0 MiB  |    ← 一直在涨
| Non-releasable memory |  5436 MiB  |  5440 MiB  | 200.0 GiB  | 194.5 GiB  |    ← 这就是碎片
| Allocations           |    8412    |    9023    |   823541   |   815129   |
| Active allocs         |    8412    |    9023    |   823541   |   815129   |
| GPU reserved segments |     145    |     148    |     156    |     11     |
| Non-releasable allocs |     283    |     291    |    7012    |    6729    |
|---------------------------------------------------------------------------|

每一行的解读

第二步 — 抓 snapshot 找元凶

memory_summary 告诉你"是泄漏还是碎片", 但要回答"是谁分配的",只能靠 snapshot。 它会记录每个 Block 的分配栈,然后用官方 visualizer 画成时间轴。

import torch

# 1) 开启历史记录 —— 训练开始前调一次
torch.cuda.memory._record_memory_history(
    enabled="all",      # 'all' = 同时记录 alloc/free 栈
    max_entries=100000,
)

# 2) 跑你的训练 / 推理代码 ...
train_one_epoch(...)

# 3) OOM 之后(或者 finally 里),把快照 dump 出来
torch.cuda.memory._dump_snapshot("snapshot.pickle")

# 4) 关掉记录,避免长跑时拖性能
torch.cuda.memory._record_memory_history(enabled=None)

把 snapshot.pickle 拖进 pytorch/torch/cuda/_memory_viz.py 就生成一个 HTML 文件(也可以直接打开 pytorch.org/memory_viz), 鼠标悬停每个 Block 都能看到分配时的 Python 栈。

三种典型 OOM 病因 × 对症下药

碎片不是 bug 而是工程取舍,但 PyTorch 给了一组环境变量让你按场景调优。 其中 expandable_segments 是 PyTorch 2.1+ 引入的"用虚拟内存根治碎片"的方案, 值得单独花时间理解。

PYTORCH_CUDA_ALLOC_CONF 速查表

expandable_segments 的工作原理

怎么开 expandable_segments

# 方法一:环境变量(最常用,启动前设)
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# 同时调多个选项,逗号分隔
export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True,max_split_size_mb:512,garbage_collection_threshold:0.8"

# 方法二:Python 端动态设置(必须在第一次 alloc 之前)
import torch
torch.cuda.memory._set_allocator_settings("expandable_segments:True")

完整 OOM 排查脚本

# oom_probe.py — 一键给训练加监控,出 OOM 自动 dump 现场
import torch, atexit, traceback

# 1. 训练开始就开 history
torch.cuda.memory._record_memory_history(enabled="all", max_entries=200_000)

def dump_on_exit():
    # 总会跑,不管是正常退出还是 OOM
    torch.cuda.memory._dump_snapshot("oom_snapshot.pickle")
    print(torch.cuda.memory_summary())

atexit.register(dump_on_exit)

# 2. 训练 loop 里周期性打印水位
def log_mem(step):
    if step % 100 == 0:
        a = torch.cuda.memory_allocated() / 2**30
        r = torch.cuda.memory_reserved()  / 2**30
        print(f"step {step}  alloc={a:.2f}G  reserved={r:.2f}G  frag={r-a:.2f}G")

for step, batch in enumerate(loader):
    loss = model(batch).sum()
    loss.backward()
    optim.step()
    optim.zero_grad(set_to_none=True)   # 关键:不要保留梯度内存
    log_mem(step)

省显存的"标准动作清单"

1. 画一张 caching allocator 的两级模型图:标出 Driver / Segment / Block / Pool / Stream 五个层次,以及它们各自的"边界"是什么(为什么不能跨边界合并)。
2. 给出一段会"看起来明显泄漏"的 Python 代码(比如把 loss 张量塞进 list 用作日志),写一段对应的诊断流程:用哪些 API、看哪几个数字、最终怎么定位元凶。
3. 解释 allocated / reserved / non-releasable / nvidia-smi 四个数之间的不等式关系,并举出一种每两个相邻数都拉开较大差距的场景。
4. 对比 classic Segment 和 expandable Segment 在"段间合并"上的差异,并解释为什么后者能用同一段虚拟地址消除段间碎片。
5. 给定环境变量 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:128,garbage_collection_threshold:0.8,说明每个选项分别针对什么样的"症状",以及组合起来会带来什么副作用或冲突。