PyTorch CUDA 缓存分配器内存碎片化：触发条件、检测与池化策略优化

背景：缓存分配器的设计目标

PyTorch 的 CUDA 缓存分配器（CUDA Caching Allocator）核心目标是让程序达到稳态运行—— 即无需频繁调用 cudaMalloc 和 cudaFree 即可满足内存分配需求。这一设计基于一个关键观察：CUDA 内存 API 会引入设备同步，打断 CPU 提前下发 GPU 指令的流水线，而 Python 解释器的延迟可以通过让 CPU 执行领先于 GPU 来隐藏。

为实现这一目标，分配器从 CUDA 申请大块内存后，在内部维护一个内存池（Block Pool），通过分割与复用策略满足张量分配请求。然而，这种池化策略在特定 workload 下会产生内存碎片化，导致可用内存被分割成无法有效利用的小块，最终触发 OOM（Out of Memory）错误。

内存碎片化的触发条件

理解碎片化触发条件是优化的前提。根据 PyTorch 分配器的设计，以下场景容易导致碎片化：

1. 动态 Shape 与变长序列

当模型处理变长序列或动态 batch size 时，每次迭代申请的内存大小不一致。分配器采用最佳适配（best-fit）策略从池中查找可用块，若找不到精确匹配的块，会分割更大的块。频繁的不规则分割会导致池中积累大量无法合并的小块。

2. 块分割与 inactive_split_bytes

分配器将内存分为两个池：小于 1MB 的 small_pool 和大于等于 1MB 的 large_pool。当从池中取出的块远大于请求大小时，分配器会执行 maybe_split_block 操作，将剩余部分返回池中。这些被分割但未激活的块计入 inactive_split_bytes 指标，是碎片化的直接度量。

关键风险在于：如果一个大块的一部分仍被活跃张量占用，整块都无法返回给 CUDA。这种 "部分占用" 场景在长训练任务中尤为常见。

3. 跨 Stream 延迟释放

分配器按 CUDA Stream 维护独立的内存池。当张量在一个 stream 上分配但在另一个 stream 使用时，需要调用 record_stream() 标记跨 stream 依赖。释放时，分配器必须等待所有依赖 stream 上的事件完成，才能将内存返回池中。这种延迟释放机制会暂时占用内存，加剧碎片化。

4. OOM 时的 cudaMalloc Retry

当分配器无法满足请求时，会触发 return_our_possibly_fragmented_memory_to_cuda()，将缓存的空闲块通过 cudaFree 返回 CUDA 驱动，然后重试分配。cudaMalloc retries 计数器反映这一过程的发生频率 —— 频繁重试是碎片化严重的信号。

碎片化检测方法

PyTorch 提供了丰富的内存统计接口，关键指标如下：

检测代码示例：

import torch

# 在疑似碎片化发生后打印内存摘要
print(torch.cuda.memory_summary(device=0, abbreviated=False))

当 OOM 错误信息中出现 "reserved in total by PyTorch" 远大于 "already allocated" 时，即可确认碎片化是主要元凶。

池化策略优化

1. max_split_size_mb 调优

环境变量 PYTORCH_CUDA_ALLOC_CONF 中的 max_split_size_mb 控制块分割的阈值。设置该参数可防止分配器将大块过度分割成小块。

# 设置最大分割块大小为 512MB
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512

调优策略：

• 若观察到 inactive_split_bytes 持续上升，尝试增大 max_split_size_mb（如 256MB 或 512MB）
• 对于固定 batch size 的训练，较小的值（如 128MB）可能更合适
• 默认值行为下，小于 1MB 的请求进入 small_pool，大于等于 1MB 的进入 large_pool

2. roundup_power2_divisions 对齐策略

默认情况下，分配大小向上对齐到 512 字节的倍数。通过设置 roundup_power2_divisions:N，可使分配大小对齐到 2 的幂次分割点，减少因微小尺寸差异导致的无法复用。

# 在 2 的幂次之间使用 4 个分割点
export PYTORCH_CUDA_ALLOC_CONF=roundup_power2_divisions:4

权衡：更大的 N 值减少内部碎片，但会增加对齐开销（N=1 平均浪费 1/4 内存，N=2 浪费 1/8）。

3. 显式缓存释放时机

torch.cuda.empty_cache() 将分配器缓存的空闲内存返回给 CUDA。应在以下时机考虑调用：

• 训练阶段切换（如从训练进入评估）
• 动态 batch size 变化前
• 长时运行服务中定期执行（权衡性能开销）

注意：频繁调用会降低分配效率，因为后续分配需要重新执行 cudaMalloc。

4. 跨 Stream 内存管理

避免不必要的跨 stream 张量共享。若必须使用，确保及时调用 record_stream()，并在合适的同步点检查内存回收情况。

实践配置清单

针对典型训练场景，建议按以下流程排查与优化：

1. 基线监控：在训练脚本中定期打印 torch.cuda.memory_summary()，建立内存使用基线
2. 碎片化识别：当出现 OOM 且 reserved >> allocated 时，确认碎片化问题
3. 参数调优：
   • 首先尝试 max_split_size_mb:256 或 max_split_size_mb:512
   • 对于变长序列任务，尝试 roundup_power2_divisions:2 或 4
4. 释放策略：在阶段切换点插入 torch.cuda.empty_cache()，观察是否缓解
5. 验证效果：监控 inactive_split_bytes 和 cudaMalloc_retries 是否下降

总结

PyTorch CUDA 缓存分配器的内存碎片化是动态 workload 下的常见问题，其根源在于块分割策略与不规则分配模式的冲突。通过理解 max_split_size_mb、roundup_power2_divisions 等配置参数的作用机制，结合 inactive_split_bytes 和 cudaMalloc_retries 等指标进行监控，可以有效缓解碎片化导致的 OOM 风险。关键在于在内存复用效率与碎片化程度之间找到适合自身 workload 的平衡点。

参考来源

• A guide to PyTorch's CUDA Caching Allocator, zdevito, 2022
• Mitigating CUDA GPU memory fragmentation and OOM issues, PyTorch Forums, 2021

systems