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

📅 2026/7/2 16:58:37
nvidia-smi 显示 8GB 空闲,为什么 PyTorch 报 CUDA out of memory?——CUDA 缓存分配器底层原理
nvidia-smi 显示 8GB 空闲为什么 PyTorch 报 CUDA out of memory——PyTorch 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|0NVIDIA RTX4090On|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**38.2# GBtorch.cuda.memory_reserved()/1024**311.5# GBallocated是 PyTorch 实际在用的。reserved是 PyTorch 从 CUDA 驱动预支但可能空闲的。这三个数字的关系指标含义工具GPU 物理总量硬件固定值nvidia-smiPyTorch 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已分配activetrue ├── Block B空闲activefalse └── Block C已分配activetrue段Segment通过cudaMalloc或cuMemMap从 CUDA 驱动获取的连续显存区域。段之间不连续。块Block从一个段上切分出来的子区域服务于具体的张量分配。分裂Splitting当一个空闲块比请求大时前面部分分配出去剩余部分作为新的空闲块。合并Merging两个相邻的空闲块可以合并为一个更大的空闲块。关键规则只有同一个段内的相邻空闲块才能合并。不同段之间的块永远不能合并。这是碎片化问题的根源。三、碎片化为什么有空闲但分配不了看一个例子。8 个 16 MiB 的张量释放后想分配 4 个 32 MiB 的张量importtorch MiB1024*1024# 分配 8 个 16 MiB 张量small[torch.empty(16*MiB,dtypetorch.uint8,devicecuda)for_inrange(8)]# 此时8 个独立的 16 MiB 段共 128 MiB reserved# 释放全部small.clear()# 此时8 个段各有 1 个 16 MiB 空闲块但 GPU 仍占 128 MiB reserved# 尝试分配 4 个 32 MiBlarge[torch.empty(32*MiB,dtypetorch.uint8,devicecuda)for_inrange(4)]# CUDA OOM发生了什么8 次cudaMalloc创建了 8 个独立的段每个 16 MiB释放后8 个段各有 1 个 16 MiB 空闲块——但它们分属不同段无法合并32 MiB 的请求在任何一个段里都找不到 ≥ 32 MiB 的连续空闲块分配器调用新的cudaMalloc分配 4 个新的 32 MiB 段总共需要 128 MiB旧的 8 个 16 MiB 段 128 MiB新的 4 个 32 MiB 段256 MiB reserved但如果你反过来分配——先分配大的再分配小的large[torch.empty(32*MiB,dtypetorch.uint8,devicecuda)for_inrange(4)]# 4 个 32 MiB 段 → 128 MiB reservedlarge.clear()# 4 个 32 MiB 空闲段small[torch.empty(16*MiB,dtypetorch.uint8,devicecuda)for_inrange(8)]# ✅ 从已有的 32 MiB 段上分裂出 16 MiB 块无需新的 cudaMalloc这就是碎片化的本质分配顺序决定了显存利用率。四、expandable_segments一个虚拟大段解决碎片化PyTorch 2.x 引入了expandable_segments。不再为每个cudaMalloc创建独立段而是使用cuMemMap创建一个虚拟地址空间cuMemMap → ExpandableSegment虚拟 1TB 连续地址空间 ├── Block A16 MiB物理显存已提交 ├── Block B32 MiB物理显存已提交 ├── Block C空闲虚拟地址已预留 └── ...关键所有 Block 都在同一个虚拟段内——相邻空闲块可以合并。同一个先小后大的场景用expandable_segmentsTrue# 设置环境变量后重启# export PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True# 先 8 个 16 MiB释放再 4 个 32 MiBsmall[torch.empty(16*MiB,dtypetorch.uint8,devicecuda)for_inrange(8)]small.clear()large[torch.empty(32*MiB,dtypetorch.uint8,devicecuda)for_inrange(4)]# ✅ 不崩因为释放的 16 MiB 块在同一个虚拟段内相邻的已经合并成大块了但这不免费cuMemMap的虚拟地址管理有开销。PyTorch 官方建议 CUDA Graph 场景用expandable_segments:True普通推理用默认值。五、max_split_size_mb你一直在用但可能不理解CSDN 上大量文章教你设exportPYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:128但它到底做了什么在分配器的maybe_split_block函数中Blockmaybe_split_block(Pool pool,size_t size,Block block){remainingblock.size-size;should_split(size1MBremaining512)||(size1MBremaining1MB);// max_split_size_mb 控制的是这里的逻辑if(!should_split)returnblock;block,restsplit(block,size);pool.add(rest);returnblock;}max_split_size_mb设的是允许分裂的最大剩余块大小。默认没有上限。当你设为 128 时如果分裂后剩余块 128 MiB不允许分裂——整个大块直接分配给请求。为什么这能缓解碎片化因为分裂产生的小块是最难合并的碎片源。限制分裂 减少小碎片的产生。但这也会浪费显存——一个 500 MiB 的块分配给 100 MiB 的请求时如果max_split_size_mb128剩余 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 类型allocatedvsreservednvidia-smi根因真实 OOMallocated ≈ GPU 总量接近 100%模型太大或 batch 太大碎片化 OOMallocated ≪ reserved ≪ GPU 100% 但报 OOM段间碎片化空闲块不连续CUDA Graph OOMreserved ≈ GPU~95%Graph workspace 锁定时新分配缓存 OOMallocated 正常reserved 暴涨忽高忽低大量小块分配产生碎片诊断命令# 看 reserved vs allocated 缺口print(fallocated:{torch.cuda.memory_allocated()/1024**3:.1f}GB)print(freserved:{torch.cuda.memory_reserved()/1024**3:.1f}GB)print(fgap:{(torch.cuda.memory_reserved()-torch.cuda.memory_allocated())/1024**3:.1f}GB)# 看段和块的分布print(torch.cuda.memory_summary())# 看碎片化程度snaptorch.cuda.memory_snapshot()forseginsnap:free_blockssum(1forbinseg[blocks]ifb[state]free)total_blockslen(seg[blocks])print(fseg{seg[total_size]//1024**2}MiB:{free_blocks}/{total_blocks}free blocks)八、总结PyTorch CUDA 分配器的核心矛盾缓存策略不还显存给驱动提升性能但制造碎片化假象。环境变量作用何时用expandable_segments:True虚拟大段消除段间碎片CUDA Graph vLLM/SGLangmax_split_size_mb:128限制分裂减少小块碎片碎片化 OOMroundup_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。