nvidia-smi 显示 8GB 空闲，为什么 PyTorch 报 CUDA out of memory？——CUDA 缓存分配器底层原理

2026 年 6 月，PyTorch 官方发布了一篇 devlog：《When does fragmentation occur in the CUDA caching allocator?》。这篇文章解释了每个 AI 开发者都遇到但几乎没人真正理解的问题------"明明 nvidia-smi 显示还有 8GB 空闲显存，为什么 PyTorch 还是报 OOM？"

---

一、nvidia-smi 和 PyTorch 看到的不是同一个"显存"

打开终端，跑两行：

$ nvidia-smi
| 0  NVIDIA RTX 4090    On  | 00000000:01:00.0 Off |                  Off |
| 30%   45C    P2    72W / 450W |  15360MiB / 24564MiB |     62%      Default |

24564 MiB 是 GPU 物理显存。15360 MiB 是 nvidia-smi 报告的"已使用"。

但 PyTorch 告诉你：

>>> torch.cuda.memory_allocated() / 1024**3
8.2   # GB
>>> torch.cuda.memory_reserved() / 1024**3
11.5  # GB

allocated 是 PyTorch 实际在用的。reserved 是 PyTorch 从 CUDA 驱动"预支"但可能空闲的。

这三个数字的关系：

指标	含义	工具
GPU 物理总量	硬件固定值	nvidia-smi
PyTorch reserved	从驱动申请的段（segment），不释放	torch.cuda.memory_reserved()
PyTorch allocated	段内实际分配给张量的块	torch.cuda.memory_allocated()

关键矛盾 ：PyTorch reserved 的内存不还给驱动。即使 Python 删了所有张量、调了 gc.collect()，nvidia-smi 仍然显示"已使用"。因为 PyTorch 的缓存分配器缓存了这些段------它在等下次分配时复用，而不是还给 CUDA 驱动。

这就是为什么你明明 del 了一个 20GB 的模型，nvidia-smi 还是显示 20GB 被占用------PyTorch 把它藏在缓存里了。

---

二、段（Segment）和块（Block）：分配器的两层结构

PyTorch CUDA 缓存分配器的核心数据结构：

cudaMalloc → Segment（大块连续显存）
                  ├── Block A（已分配，active=true）
                  ├── Block B（空闲，active=false）
                  └── Block C（已分配，active=true）

• 段（Segment） ：通过 cudaMalloc 或 cuMemMap 从 CUDA 驱动获取的连续显存区域。段之间不连续。
• 块（Block）：从一个段上切分出来的子区域，服务于具体的张量分配。
• 分裂（Splitting）：当一个空闲块比请求大时，前面部分分配出去，剩余部分作为新的空闲块。
• 合并（Merging）：两个相邻的空闲块可以合并为一个更大的空闲块。

关键规则：只有同一个段内的相邻空闲块才能合并。不同段之间的块永远不能合并。

这是碎片化问题的根源。

---

三、碎片化：为什么"有空闲"但"分配不了"

看一个例子。8 个 16 MiB 的张量，释放后想分配 4 个 32 MiB 的张量：

import torch
MiB = 1024 * 1024

# 分配 8 个 16 MiB 张量
small = [torch.empty(16 * MiB, dtype=torch.uint8, device='cuda') for _ in range(8)]
# 此时：8 个独立的 16 MiB 段，共 128 MiB reserved

# 释放全部
small.clear()
# 此时：8 个段各有 1 个 16 MiB 空闲块，但 GPU 仍占 128 MiB reserved

# 尝试分配 4 个 32 MiB
large = [torch.empty(32 * MiB, dtype=torch.uint8, device='cuda') for _ in range(4)]
# 💣 CUDA OOM！

发生了什么：

1. 8 次 cudaMalloc 创建了 8 个独立的段，每个 16 MiB
2. 释放后，8 个段各有 1 个 16 MiB 空闲块------但它们分属不同段，无法合并
3. 32 MiB 的请求在任何一个段里都找不到 ≥ 32 MiB 的连续空闲块
4. 分配器调用新的 cudaMalloc 分配 4 个新的 32 MiB 段
5. 总共需要 128 MiB（旧的 8 个 16 MiB 段）+ 128 MiB（新的 4 个 32 MiB 段）= 256 MiB reserved

但如果你反过来分配------先分配大的，再分配小的：

large = [torch.empty(32 * MiB, dtype=torch.uint8, device='cuda') for _ in range(4)]
# 4 个 32 MiB 段 → 128 MiB reserved
large.clear()
# 4 个 32 MiB 空闲段
small = [torch.empty(16 * MiB, dtype=torch.uint8, device='cuda') for _ in range(8)]
# ✅ 从已有的 32 MiB 段上分裂出 16 MiB 块！无需新的 cudaMalloc

这就是碎片化的本质：分配顺序决定了显存利用率。

---

四、expandable_segments：一个虚拟大段解决碎片化

PyTorch 2.x 引入了 expandable_segments。不再为每个 cudaMalloc 创建独立段，而是使用 cuMemMap 创建一个虚拟地址空间：

cuMemMap → ExpandableSegment（虚拟 1TB 连续地址空间）
                ├── Block A（16 MiB，物理显存已提交）
                ├── Block B（32 MiB，物理显存已提交）
                ├── Block C（空闲，虚拟地址已预留）
                └── ...

关键 ：所有 Block 都在同一个虚拟段内------相邻空闲块可以合并。

同一个"先小后大"的场景，用 expandable_segments=True：

# 设置环境变量后重启
# export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# 先 8 个 16 MiB，释放，再 4 个 32 MiB
small = [torch.empty(16 * MiB, dtype=torch.uint8, device='cuda') for _ in range(8)]
small.clear()
large = [torch.empty(32 * MiB, dtype=torch.uint8, device='cuda') for _ in range(4)]
# ✅ 不崩！因为释放的 16 MiB 块在同一个虚拟段内，相邻的已经合并成大块了

但这不免费 ：cuMemMap 的虚拟地址管理有开销。PyTorch 官方建议 CUDA Graph 场景用 expandable_segments:True，普通推理用默认值。

---

五、max_split_size_mb：你一直在用但可能不理解

CSDN 上大量文章教你设：

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

但它到底做了什么？在分配器的 maybe_split_block 函数中：

Block maybe_split_block(Pool pool, size_t size, Block block) {
    remaining = block.size - size;
    should_split = (size < 1MB && remaining > 512) ||
                   (size >= 1MB && remaining > 1MB);
    // max_split_size_mb 控制的是这里的逻辑
    if (!should_split) return block;
    block, rest = split(block, size);
    pool.add(rest);
    return block;
}

max_split_size_mb 设的是"允许分裂的最大剩余块大小"。默认没有上限。当你设为 128 时：如果分裂后剩余块 > 128 MiB，不允许分裂------整个大块直接分配给请求。

为什么这能缓解碎片化？因为分裂产生的小块是最难合并的碎片源。限制分裂 = 减少小碎片的产生。

但这也会浪费显存 ------一个 500 MiB 的块分配给 100 MiB 的请求时，如果 max_split_size_mb=128，剩余 400 MiB 不能分裂出来给别人用。

---

六、CUDA Graph 与分配器的致命交互

CUDA Graph 捕获期间，PyTorch 分配器会记录所有 tensor 的内存地址。回放时，图必须使用相同的地址------这意味着捕获期间的显存分配不能被释放。

这就是为什么你在 vLLM 中看到：

AssertionError: Workspace is locked but allocation requires 0.76 MB.
Workspace growth is not allowed after locking.

CUDA Graph 捕获完成后，分配器锁定了 workspace。任何新的分配请求------即使是 0.76 MB------都会触发断言失败。

expandable_segments 在这里有帮助：虚拟地址空间预留了位置，物理显存可以按需提交。但这是两刃剑------物理显存不够时仍然会 OOM。

---

七、四类 OOM 的分配器级诊断

下次看到 CUDA OOM，先判断是哪一类：

OOM 类型	allocated vs reserved	nvidia-smi	根因
真实 OOM	allocated ≈ GPU 总量	接近 100%	模型太大或 batch 太大
碎片化 OOM	allocated ≪ reserved ≪ GPU	< 100% 但报 OOM	段间碎片化，空闲块不连续
CUDA Graph OOM	reserved ≈ GPU	~95%	Graph workspace 锁定时新分配
缓存 OOM	allocated 正常，reserved 暴涨	忽高忽低	大量小块分配产生碎片

诊断命令：

# 看 reserved vs allocated 缺口
print(f"allocated: {torch.cuda.memory_allocated() / 1024**3:.1f} GB")
print(f"reserved:  {torch.cuda.memory_reserved() / 1024**3:.1f} GB")
print(f"gap:       {(torch.cuda.memory_reserved() - torch.cuda.memory_allocated()) / 1024**3:.1f} GB")

# 看段和块的分布
print(torch.cuda.memory_summary())

# 看碎片化程度
snap = torch.cuda.memory_snapshot()
for seg in snap:
    free_blocks = sum(1 for b in seg['blocks'] if b['state'] == 'free')
    total_blocks = len(seg['blocks'])
    print(f"seg {seg['total_size']//1024**2}MiB: {free_blocks}/{total_blocks} free blocks")

---

八、总结

PyTorch CUDA 分配器的核心矛盾：缓存策略（不还显存给驱动）提升性能，但制造碎片化假象。

环境变量	作用	何时用
expandable_segments:True	虚拟大段，消除段间碎片	CUDA Graph + vLLM/SGLang
max_split_size_mb:128	限制分裂，减少小块碎片	碎片化 OOM
roundup_power2_divisions:4	减少对齐浪费	大量不规则 size 的推理

下次你看到"CUDA out of memory, 11 GiB free"时，你知道那不是显存不够------是分配器的段无法合并了。

---

本文参考了 PyTorch DevLog (2026-06-01)、Zach DeVito's Blog (2022-08-04) 以及 PyTorch 源码 CUDACachingAllocator.cpp。