GPU 显存碎片化深度实战：从 Buddy 分配器、expandable_segments 到碎片率治理与 CUDA OOM 修复

18 阅读 0 评论 0 点赞

GPU 显存碎片化深度实战：从 Buddy 分配器、expandable_segments 到碎片率治理与 CUDA OOM 修复

引言：有空闲却 OOM 的诡异现场

每一个在 A100/H100 上跑过大模型训练或长上下文推理的工程师，大概率都见过这样一个让人抓狂的报错：


torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.20 GiB.
GPU has a total capacity of 80.00 GiB of which 42.10 GiB is free.

注意最后一句：42 GiB 是空闲的，却连 1.2 GiB 都分配不出来。

这不是 PyTorch 在撒谎，也不是驱动坏了。真正的原因是显存碎片化（VRAM Fragmentation）：那 42 GiB 空闲显存被打散成了无数个互不相邻的小块，任何一个单独的块都不足以容纳这次请求所需的连续显存。本文将从 GPU 显存分配器的底层机制出发，讲清楚碎片是怎么产生的、如何量化观测、以及如何用 expandable_segments、max_split_size_mb、内存池与 MIG/MPS 隔离等手段系统性治理，最后给出一个可复现的"碎片率 60% → OOM → 治理后稳定"的完整案例。

一、显存碎片化的两种形态与量化定义

1.1 外部碎片 vs 内部碎片

和任何内存分配系统一样，GPU 显存也存在两类碎片：

类型	含义	在 GPU 上的典型成因
外部碎片（External）	总空闲量足够，但被切成不连续的零散块，无法凑出一块满足请求大小的连续区域	变长激活值、动态 KV-Cache、多流交错分配释放
内部碎片（Internal）	分配器为了满足对齐/粒度要求，给一块请求分配了比实际需要更大的块，多出的部分被浪费	CUDA 分配粒度取整、PyTorch caching allocator 的 rounding

外部碎片是导致"假 OOM"的元凶；内部碎片则悄悄吞噬可用容量、抬高 reserved 水位。

1.2 碎片率的工程化定义

对于一个分配器视角，最实用的碎片率定义为：


fragmentation_ratio = 1 - (max_contiguous_free_bytes / total_free_bytes)

total_free_bytes：当前所有空闲块字节之和
max_contiguous_free_bytes：最大的单一连续空闲块字节数

当 fragmentation_ratio 趋近 0，说明空闲显存高度连续，分配大块毫无压力；当趋近 1，说明空闲被切碎，极易触发 OOM。业界的经验红线是：碎片率 > 0.4 时，分配大块张量的失败风险显著上升；> 0.6 时基本必然偶发 OOM。

二、GPU 显存分配器的底层：Buddy 与 Caching Allocator 双层结构

要治理碎片，必须先理解显存请求是走了哪两层分配器。

2.1 驱动层：CUDA 的 Buddy/Segment 分配器

当你调用 cudaMalloc（PyTorch 的第一次分配会触发它），请求的其实是 CUDA 驱动里的分配器。它的核心特征：

采用 段（segment）+ buddy 风格的切分：驱动从显存中切出一段 segment，再按 2 的幂（或固定 bin）递归二分满足请求大小；
释放时尝试与相邻的"伙伴块"合并，归还更大的连续区域；
分配粒度存在最小对齐与大块阈值两条边界，小块走小 bin、大块走大 bin。

这就解释了为什么"交替分配大块和小块"最伤显存：大块把 segment 切出一个洞，小块填充洞的缝隙，等大块释放后那个洞又被周围的小块包围，无法与别的大块合并。

2.2 框架层：PyTorch Caching Allocator

PyTorch 不会每次张量释放都立刻 cudaFree 还给驱动——那样太慢。它维护了自己的 caching allocator：

释放的张量不立即归还驱动，而是缓存进 per-stream 的 block 池；
下次分配先去池里找大小合适的空闲 block 复用，找不到才向驱动 cudaMalloc 新 segment；
持有但未使用的 block 记为 reserved（已向驱动预留），已分配给张量的记为 allocated。

于是出现了两个视角的"空闲"：

视角	查看方式	含义
驱动层空闲	`nvidia-smi` / `cudaMemGetInfo`	驱动认为还给它的显存，包含所有框架 reserved 但未 allocated 的部分
分配器层空闲	`torch.cuda.memory_stats()`	caching allocator 池里可复用的 block，这才是真正能立刻满足 Python 端请求的量

诡异 OOM 的本质：nvidia-smi 显示 42 GiB free（驱动层），但 caching allocator 池里没有一块够大的连续 block（分配器层），于是 Python 端 OOM。两个数字对不上，就是碎片在作怪。

三、碎片是怎么被"养"出来的

3.1 变长生命周期交错

训练/推理中，以下几类对象的生命周期与体积差异极大，最容易制造碎片：

激活值（activation）：前向产生、反向用完即弃，体积随 batch/序列长度剧烈波动；
KV-Cache：长上下文推理中随生成长度线性增长，且不同请求长度不一；
梯度与优化器状态：生命周期贯穿整个 step，体积固定且巨大；
workspace / 临时 buffer：算子内部申请，用完即释。

当"巨大且长寿"的对象与"零碎且短寿"的对象在同一上下文里交错分配释放，空洞就被钉死了。

3.2 多进程共享上下文（MPS）

nvidia-smi -c COMPUTE 之外的 MPS（Multi-Process Service） 让多个进程共享同一个 CUDA context 与同一片显存。问题来了：进程 A 释放的 block 在进程 B 的视角里未必能复用，而 context 级的预留又是共享的——跨进程的碎片无法被单个进程回收，整体碎片率会系统性偏高。

3.3 分配顺序敏感性（可复现反例）

下面这段脚本刻意制造碎片：先交替分配大块/小块，再释放所有小块，留下"大块之间的缝隙"，最后尝试分配一个比原来大块还大的连续张量。


import torch

def frag_demo():
    torch.cuda.empty_cache()
    big, small = [], []
    # 交替分配，制造交错空洞
    for i in range(20):
        big.append(torch.empty((600, 1024, 1024), dtype=torch.float16, device="cuda"))   # ~1.2 GiB
        small.append(torch.empty((10, 1024, 1024), dtype=torch.float16, device="cuda"))  # ~20 MiB
    # 只释放小块，大块之间的缝隙被永久钉死
    del small
    torch.cuda.empty_cache()
    stats = torch.cuda.memory_stats()
    reserved = stats["reserved_bytes.all.current"]
    allocated = stats["allocated_bytes.all.current"]
    frag = 1.0 - allocated / reserved if reserved else 0.0
    print(f"reserved={reserved/1e9:.2f}GiB allocated={allocated/1e9:.2f}GiB "
          f"allocator_frag={frag*100:.1f}%")
    # 尝试分配一个略小于单块、但当前无连续空间的张量
    try:
        torch.empty((580, 1024, 1024), dtype=torch.float16, device="cuda")
        print("alloc OK")
    except RuntimeError as e:
        print("alloc FAILED ->", str(e)[:80])

frag_demo()

观察：释放小块后 nvidia-smi 可能显示大量 free，但 allocator_frag 居高不下，且最后的分配很容易失败——这正是生产环境 OOM 的迷你版。

四、如何量化与观测碎片

4.1 从 memory_stats 取出关键指标

torch.cuda.memory_stats() 返回的字典里，下面几个键是碎片治理的"仪表盘"：


import torch, json

stats = torch.cuda.memory_stats()
keys_of_interest = [
    "allocated_bytes.all.current",
    "reserved_bytes.all.current",
    "active_bytes.all.current",
    "requested_bytes.all.current",
    "num_alloc_retries",   # 分配重试次数：>0 说明发生过碎片导致的重试
    "num_ooms",            # 已发生的 OOM 次数
    "segment.all.current", # 当前 segment 数量（越多越碎）
]
for k in keys_of_interest:
    print(f"{k:38s} = {stats.get(k)}")

# 分配器层碎片率（allocated vs reserved 的近似）
reserved = stats["reserved_bytes.all.current"]
allocated = stats["allocated_bytes.all.current"]
frag = 1.0 - allocated / reserved if reserved else 0.0
print(f">>> allocator fragmentation ratio = {frag*100:.1f}%")

num_alloc_retries > 0：caching allocator 曾经找不到合适 block，被迫向驱动要新 segment——碎片的早期信号，比 OOM 出现得早得多。
segment.all.current 持续攀升：说明 allocator 在不断向驱动开新段，旧段因碎片无法复用。

4.2 驱动层 vs 分配器层双视角探针


import torch
import pynvml  # pip install nvidia-ml-py

def probe():
    free, total = torch.cuda.mem_get_info()          # 驱动层
    stats = torch.cuda.memory_stats()                # 分配器层
    drv_free_gib = free / 1e9
    alloc_free_gib = (stats["reserved_bytes.all.current"]
                      - stats["allocated_bytes.all.current"]) / 1e9
    print(f"driver free   = {drv_free_gib:6.2f} GiB")
    print(f"allocator free= {alloc_free_gib:6.2f} GiB  (caching pool 可复用)")
    print(f"gap           = {drv_free_gib - alloc_free_gib:6.2f} GiB  (reserved 但未 allocated，被碎片锁定)")

probe()

两者之差就是"被碎片锁死、名义空闲却用不上"的那部分显存。监控这个 gap 比监控 nvidia-smi 的 free 更有意义。

4.3 内存快照可视化

PyTorch 2.x 提供快照能力，可把分配事件导出为 Chrome-trace 风格的 JSON，用官方 visualizer 看碎片分布：


import torch

torch.cuda.memory._record_memory_history(max_entries=100000)
# ... 跑一段训练/推理 ...
torch.cuda.memory._dump_snapshot("vram_snapshot.pickle")
torch.cuda.memory._record_memory_history(enabled=None)

快照能直观看到哪些 segment 被切成碎块、哪些大块长期驻留，是定位"谁在制造碎片"的终极工具。

五、三大治理手段

5.1 expandable_segments：让 segment 可增长（PyTorch 2.1+）

传统 caching allocator 的 segment 大小在创建时就固定了；当预留的 segment 被碎片填满，只能再开一个新 segment，进一步加剧碎片。

expandable_segments:True 让已有 segment 按需向驱动扩展，而不是另起炉灶开新段：


export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True


# 或在代码里设置（需在首次 CUDA 分配前）
import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
import torch  # 必须在设置后 import / 首次分配前生效

适用场景：长时运行、显存占用随 step 缓慢增长（如训练 epoch 中激活峰值波动）的任务。它在多数大模型训练场景可把碎片率从 40%~60% 压到 10% 以内。

5.2 max_split_size_mb：限制 block 被切碎的上限

当一个大 block 被切给小请求后，它内部的小块很难再合并回大块。设置上限可防止 allocator 把大块无脑切碎：


export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:256

含义：大于 256 MiB 的空闲 block 不会被进一步拆分去满足小请求；小请求会去别处找合适的块，或触发新分配。这让"大块保持完整"，将来大张量还能用上。

常见取值：变长 KV-Cache 场景可设 128~512 MiB；固定 shape 的训练可设更大（如 1024）以减少 allocator 频繁开段。

5.3 主动 Defrag 与内存池

torch.cuda.empty_cache()：释放 caching allocator 持有的、所有未被张量占用的 reserved block，还给驱动。代价是下次分配要重新 cudaMalloc，有延迟尖刺——适合在 step 间隙、请求边界调用，不要在每个迭代都调。
显式预分配 workspace 复用：对推理服务的临时 buffer、attention 的 workspace，预先分配一个大的可复用 tensor，手动切片管理，避免反复申请释放制造碎片。
bucketing（分桶）：对变长序列，把请求按长度分桶，桶内 shape 对齐，分配/释放的块大小一致，天然降低碎片（HuggingFace static_kv_cache、vLLM PagedAttention 都是此思想）。

组合示例：


export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:512

六、实战案例：从 60% 碎片率 OOM 到稳定

6.1 复现场景

一个长上下文推理服务，动态 batch + 变长 KV-Cache，偶尔在请求洪峰后 OOM，但 nvidia-smi 显示显存还够。治理前观测：

指标	治理前
驱动层 free	38 GiB
allocator free（可复用）	9 GiB
allocator fragmentation ratio	58%
`num_alloc_retries`（/小时）	120+
`num_ooms`（/天）	3~5

6.2 治理步骤

开启 expandable_segments:True，让 segment 可扩展，减少另开新段；
设置 max_split_size_mb:512，保住大块完整性；
推理请求按序列长度分桶，KV-Cache 用固定分桶复用，消除长短请求交错；
在每轮调度空闲时调用一次 empty_cache() 做轻度 defrag（非每请求）。

6.3 治理后对比

指标	治理前	治理后
驱动层 free	38 GiB	36 GiB
allocator free（可复用）	9 GiB	30 GiB
allocator fragmentation ratio	58%	9%
`num_alloc_retries`	120+/小时	0
`num_ooms`	3~5/天	0
P99 首 token 延迟	波动大	稳定

关键结论：总空闲几乎没变，但可复用连续块从 9 GiB 提升到 30 GiB，碎片率从 58% 降到 9%，OOM 彻底消失。这再次印证——问题不在"容量"，而在"连续性"。

七、MPS / MIG：隔离即治理

当多任务必须共享一张卡时，隔离碎片域本身就是治理手段：

MIG（Multi-Instance GPU）：硬件级切分，每个实例有独立的物理显存与调度，碎片被物理隔离在不同实例，互不影响。适合多租户推理。
MPS（Multi-Process Service）：共享单一 context，不隔离碎片，跨进程碎片不可回收，前文已述风险。仅在确实需要共享 context（如多进程协作同一模型）时使用，并配合 max_split_size_mb 约束。
独立 context / 独立进程 + CUDA_VISIBLE_DEVICES：最朴素的隔离，碎片域自然分离，代价是显存无法跨进程共享。

经验法则：多租户推理优先 MIG；单模型多副本训练用独立进程；MPS 只用于真正需要共享 context 的协作场景。

八、训练 vs 推理的不同策略

维度	训练	推理
主要碎片源	激活值波动、梯度检查点、变长样本	动态 KV-Cache、动态 batch
首选手段	`expandable_segments` + 激活 offload + 梯度检查点	分桶 + PagedAttention + `max_split_size_mb`
进阶	activation recomputation（用计算换连续）	连续批处理（continuous batching）摊薄碎片
监控重点	`num_alloc_retries` 随 step 趋势	每请求前后碎片率 delta

特别地，PagedAttention（vLLM 等）本身就是一种碎片治理架构：它把 KV-Cache 切成固定大小的 page，像操作系统虚拟内存一样按需分页，从根本上消除了"需要一整块连续 KV"的需求——这是推理侧对抗显存碎片的最优雅解法。

九、把碎片率纳入可观测性

生产环境应把碎片率作为一级指标。下面一个最小 Prometheus exporter 示例：


from prometheus_client import Gauge, start_http_server
import torch, time

FRAG = Gauge("vram_allocator_fragmentation_ratio", "CUDA caching allocator fragmentation ratio")
RETRY = Gauge("vram_alloc_retries_total", "CUDA allocator alloc retries")
GAP = Gauge("vram_locked_by_fragmentation_bytes", "Driver-free minus allocator-free bytes")

def collect():
    s = torch.cuda.memory_stats()
    reserved = s["reserved_bytes.all.current"]
    allocated = s["allocated_bytes.all.current"]
    free, _ = torch.cuda.mem_get_info()
    frag = 1.0 - allocated / reserved if reserved else 0.0
    FRAG.set(round(frag, 4))
    RETRY.set(s.get("num_alloc_retries", 0))
    GAP.set(int(free - (reserved - allocated)))

if __name__ == "__main__":
    start_http_server(9102)
    while True:
        collect()
        time.sleep(5)

告警规则建议：vram_allocator_fragmentation_ratio > 0.4 持续 5 分钟即触发，先于 OOM 介入治理。

十、总结与落地 Checklist

GPU 显存碎片化的核心认知只有三句话：

OOM 看的是"连续性"，不是"总空闲"——nvidia-smi 的 free 会骗人；
碎片分两层：驱动 buddy 分配器 + 框架 caching allocator，治理要针对 allocator 层；
可扩展段 + 限切上限 + 分桶复用 + 隔离（MIG/独立进程）是四件套。

落地 Checklist：

[ ] 用 torch.cuda.memory_stats() 的 num_alloc_retries / fragmentation 建立基线监控
[ ] 长时任务优先 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
[ ] 变长场景加 max_split_size_mb（128~512 起调）
[ ] 推理服务做 KV-Cache 分桶 / 启用 PagedAttention
[ ] 多任务共享卡优先 MIG，慎用 MPS
[ ] 只在请求边界/step 间隙 empty_cache()，切勿每迭代调用
[ ] 把碎片率接入 Prometheus，>0.4 告警先于 OOM

把"显存够不够"的朴素问题，升级为"显存连不连续"的工程问题，是每一个 AI 基础设施工程师走向成熟的标志。

点赞(0) 打赏

本文分类：GPU并行计算
本文标签：GPU显存显存碎片化 CUDA OOM PyTorch AI基础设施显存治理
浏览次数：18 次浏览
发布日期：2026-10-08 02:22:58
本文链接：https://ybb.press/gpu-compute/15798.html

GPU 显存碎片化深度实战：从 Buddy 分配器、expandable_segments 到碎片率治理与 CUDA OOM 修复

GPU 显存碎片化深度实战：从 Buddy 分配器、expandable_segments 到碎片率治理与 CUDA OOM 修复

引言：有空闲却 OOM 的诡异现场

一、显存碎片化的两种形态与量化定义

1.1 外部碎片 vs 内部碎片

1.2 碎片率的工程化定义

二、GPU 显存分配器的底层：Buddy 与 Caching Allocator 双层结构

2.1 驱动层：CUDA 的 Buddy/Segment 分配器

2.2 框架层：PyTorch Caching Allocator

三、碎片是怎么被"养"出来的

3.1 变长生命周期交错

3.2 多进程共享上下文（MPS）

3.3 分配顺序敏感性（可复现反例）

四、如何量化与观测碎片

4.1 从 memory_stats 取出关键指标

4.2 驱动层 vs 分配器层双视角探针

4.3 内存快照可视化

五、三大治理手段

5.1 expandable_segments：让 segment 可增长（PyTorch 2.1+）

5.2 max_split_size_mb：限制 block 被切碎的上限

5.3 主动 Defrag 与内存池

六、实战案例：从 60% 碎片率 OOM 到稳定

6.1 复现场景

6.2 治理步骤

6.3 治理后对比

七、MPS / MIG：隔离即治理

八、训练 vs 推理的不同策略

九、把碎片率纳入可观测性

十、总结与落地 Checklist

评论列表共有 0 条评论

发表评论取消回复

GPU 显存碎片化深度实战：从 Buddy 分配器、expandable_segments 到碎片率治理与 CUDA OOM 修复

GPU 显存碎片化深度实战：从 Buddy 分配器、expandable_segments 到碎片率治理与 CUDA OOM 修复

引言：有空闲却 OOM 的诡异现场

一、显存碎片化的两种形态与量化定义

1.1 外部碎片 vs 内部碎片

1.2 碎片率的工程化定义

二、GPU 显存分配器的底层：Buddy 与 Caching Allocator 双层结构

2.1 驱动层：CUDA 的 Buddy/Segment 分配器

2.2 框架层：PyTorch Caching Allocator

三、碎片是怎么被"养"出来的

3.1 变长生命周期交错

3.2 多进程共享上下文（MPS）

3.3 分配顺序敏感性（可复现反例）

四、如何量化与观测碎片

4.1 从 memory_stats 取出关键指标

4.2 驱动层 vs 分配器层双视角探针

4.3 内存快照可视化

五、三大治理手段

5.1 expandable_segments：让 segment 可增长（PyTorch 2.1+）

5.2 max_split_size_mb：限制 block 被切碎的上限

5.3 主动 Defrag 与内存池

六、实战案例：从 60% 碎片率 OOM 到稳定

6.1 复现场景

6.2 治理步骤

6.3 治理后对比

七、MPS / MIG：隔离即治理

八、训练 vs 推理的不同策略

九、把碎片率纳入可观测性

十、总结与落地 Checklist

反向传播算法深度实战：从链式法则到现代深度学习引擎

GPU 显存碎片化深度实战：从 Buddy 分配器、expandable_segments 到碎片率治理与 CUDA OOM 修复

Linux内核 EEVDF 调度器：从 CFS 到最早截止时间优先的架构革命

AI 训练与推理混部场景下 Linux I/O 优先级工程实战：从 cgroup iocost 到 io_uring sqthread 的端到端隔离

评论列表 共有 0 条评论

发表评论 取消回复

评论列表共有 0 条评论

发表评论取消回复