CUDA 显存碎片问题：诊断与解决方案

1. 什么是显存碎片

深度学习训练过程中，PyTorch 通过 CUDA Caching Allocator 管理 GPU 显存。该 allocator 会向 CUDA driver 申请大块 segment，然后在内部按需切分为更小的 block 分配给各个 tensor。

显存碎片指的是：一个大 block 被中间插入的小 block 切分后，即使大 block 对应的 tensor 已释放，由于小 block 仍然存活，整个 segment 无法合并回收。这种碎片并非物理设备上真实的地址碎片，而是 Caching Allocator 层面的逻辑碎片。

典型场景：

Step 1: allocator 申请一个大 segment
        [LLLLLLLLLLLLLLLLLL]

Step 2: 大 tensor 释放，但小 tensor 卡在中间
        [____ss____________]
             ↑ 小 block 存活，segment 无法合并

Step 3: 下一轮需要同样大小的连续空间 → 找不到 → OOM

当这种情况以固定周期重复发生时（例如每个训练 step），碎片会不断累积，最终导致 OOM——即使 nvidia-smi 显示仍有大量空闲显存。

2. 如何判断 OOM 是否由碎片引起

2.1 使用 PyTorch Memory Visualizer

PyTorch 提供了官方的显存分析工具：https://docs.pytorch.org/memory_viz

使用方式：

# 在代码中启用显存快照记录
torch.cuda.memory._record_memory_history(max_entries=100000)

# ... 训练代码 ...

# 导出快照
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")
torch.cuda.memory._record_memory_history(None)

将生成的 .pickle 文件拖入 https://docs.pytorch.org/memory_viz 即可可视化。

2.2 碎片的典型特征

在 Memory Visualizer 中，碎片问题通常表现为：

• 大面积灰色空白区域：表示已释放但无法合并的空闲 block
• 空白区域中间夹有细小的彩色条：存活的小 tensor 阻止了 segment 合并
• 每个 step 都产生新的碎片化 segment：说明碎片是周期性的

如果在 OOM 发生时，torch.cuda.memory_summary() 显示 reserved 远大于 allocated，且存在大量 inactive split block，基本可以确认是碎片问题。

3. 解决方案

3.1 expandable_segments（推荐，零代码改动）

import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

原理：使用 cuMemMap 虚拟地址映射代替 cudaMalloc，segment 可以按需扩展和收缩。相邻的空闲 block 即使跨 segment 也能合并，从根本上减少碎片。

适用条件：需要 CUDA >= 11.4 且 GPU driver 支持虚拟内存管理（大多数现代 NVIDIA GPU 均支持）。

3.2 max_split_size_mb（简单，无需改代码）

import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:256"

原理：超过指定大小的 block 不会被 split。小 tensor 会被分配到另一个 segment，大 block 释放后保持完整。

使用建议：

1. 通过 Memory Visualizer 确定碎片 block 的大小
2. 将 max_split_size_mb 设为略小于碎片 block 大小的值
3. 测试验证不会因此导致新的 OOM

局限性：

• 可能增加总显存占用（小 tensor 不能复用大 block 的尾部空间）
• 极端情况下，当显存不足触发 empty_cache 时，释放的大 segment 可能因物理地址不连续而无法重新分配出同样大小的连续空间，仍然 OOM
• 需要针对具体场景反复调参

3.3 MemPool 隔离（推荐，代码改动极小）

PyTorch 2.5 起引入了 MemPool API，支持将特定分配路由到独立的私有内存池。

各版本功能支持：

原理：将大 tensor 的分配隔离到独立 pool 中，小 tensor 走 default pool，两者互不干扰。pool 内的 segment 不会被 pool 外的分配 split。

关键特性：

• torch.cuda.empty_cache() 不会清除 MemPool 的缓存，pool 内的缓存可以安全地跨 step 复用
• pool 在 Python 对象生命周期内一直存在，del pool 后才会清理

代码示例：

# 在循环外创建一次
heavy_pool = torch.cuda.MemPool()

for step in range(num_steps):
    # 大 tensor 在独立 pool 中，不会被小 tensor split
    with torch.cuda.use_mem_pool(heavy_pool):
        big_tensor = compute_large_block(...)

    # 小 tensor 走 default pool，互不干扰
    small_result = some_small_ops(big_tensor)

    del big_tensor
    # heavy_pool 的 segment 完整释放回缓存，下一轮直接复用

改动量评估：只需定位产生大 block 的代码位置，在其外层套一个 with torch.cuda.use_mem_pool(pool):，原有计算逻辑零改动。

3.4 手动复用 Buffer（改动大，效果最强）

对于能够预先确定 shape 的大 tensor，可以在循环外预分配 buffer，每轮复用同一块内存。

# 循环外预分配
activation_buffer = torch.empty(batch, seq_len, hidden, device="cuda")

for step in range(num_steps):
    # 用 out= 写入预分配 buffer，不产生新分配
    torch.mm(input, weight, out=activation_buffer)
    ...

前提条件：

• 算子需要支持 out= 参数或 in-place 操作
• tensor 的 shape 每轮固定或有已知上界

局限性：代码侵入性大，需要逐个修改算子调用。适用于对性能有极致要求的场景。

4. 方案对比

推荐策略：优先尝试 expandable_segments（零改动）；如果环境不支持，使用 MemPool 隔离（改动极小）；max_split_size_mb 作为快速止血手段；手动 buffer 复用留给极端场景。