Pytorch 内存分配与 max_split_size_mb

训练 Pytorch 模型时会遇到 CUDA Out of Memory 的问题，大部分情况下是模型本身占用显存超过硬件极限，但是有时是Pytorch 内存分配机制导致预留显存太多，从而报出显存不足的错误，针对这种情况，本文记录 Pytorch 内存分配机制，与通过配置 max_split_size_mb 来解决上述问题。

问题复现

假如我们当前的显存分配如上图所示，假设当前想分配 800MB 显存，虽然空闲的总显存有 1000MB，但是上方图的空闲显存由地址不连续的两个 500MB 的块组成，不够分配这 800MB 显存；而下方的图中，如果两个 500MB 的空闲块地址连续，就可以通过显存碎片的整理组成一个 1000MB 的整块，足够分配 800MB。上方图的这种情况就被称为显存碎片化。

PyTorch 显存管理

基础概念

Block

• 分配 / 管理内存块的基本单位，(stream_id, size, ptr) 三元组可以特异性定位一个 Block，即 Block 维护一个 ptr 指向大小为 size 的内存块，隶属于 stream_id 的 CUDA Stream。
• 所有地址连续的 Block（不论是否为空闲，只要是由 Allocator::malloc 得来的）都被组织在一个双向链表里，便于在释放某一个 Block 时快速检查前后是否存在相邻碎片，若存在可以直接将这三个 Block 合成为一个。

BlockPool

• 内存池，用 std::set 存储 Block 的指针，按照 (cuda_stream_id -> block size -> addr) 的优先级从小到大排序。所有保存在 BlockPool 中的 Block 都是空闲的。

• DeviceCachingAllocator 中维护两种 BlockPool (large_blocks, small_blocks)，分别存放较小的块和较大的块（为了分别加速小需求和大需求），简单地将 <= 1MB 的 Block 归类为小块，> 1MB 的为大块。

Block 在 Allocator 内有两种组织方式，一种是显式地组织在 BlockPool（红黑树）中，按照大小排列；另一种是具有连续地址的 Block 隐式地组织在一个双向链表里（通过结构体内的 prev, next 指针），可以以 O(1) 时间查找前后 Block 是否空闲，便于在释放当前 Block 时合并碎片。

显存申请

Pytorch 申请显存需要用到 malloc 函数：返回一个可用的 Block（L466）

表述定义：

• size：请求分配的显存大小
• alloc_size：需要 cudaMalloc 时的显存大小

hard-coded：

根据 size 决定实际上的 alloc_size（get_allocation_size 函数，L1078）：

• 为小于 1MB 的 size 分配 2MB；
• 为 1MB ～ 10MB 的 size 分配 20MB；
• 为 >= 10MB 的 size 分配 { size 向上取整至 2MB 倍数 } MB。

申请步骤

Pytorch 在申请显存时会寻找是否有合适的 block， 该过程有五个步骤，如果这五个步骤都没找到合适的 Block，就会报经典的 [CUDA out of memory. Tried to allocate …] 错误。

步骤一：get_free_block 函数（L1088）

TLDR：尝试在 Allocator 自己维护的池子中找一个大小适中的空闲 Block 返回。
*** TLDR = Too Long; Didn’t Read**

• 用当前的 (size, stream_id) 这二元组制作 Block Key 在对应的 BlockPool 中查找；
• 环境变量 PYTORCH_CUDA_ALLOC_CONF 中指定了一个阈值 max_split_size_mb，有两种情况不会在此步骤分配：

1. 需要的 size 小于阈值但查找到的 Block 的比阈值大（避免浪费block）；
2. 两方都大于阈值但 block size 比需要的 size 大得超过了 buffer（此处是 20MB，这样最大的碎片不超过 buffer 大小）。

• 这里的这个阈值 max_split_size_mb 涉及一个有趣的特性，后面会讲到。
• 若成功分配一个 Block，则将这个 Block 从 BlockPool 中删除。后续用完释放（free）时会再 insert 进去，见free_block : L976）

步骤二：trigger_free_memory_callbacks 函数（L1106）

TLDR：手动进行一波垃圾回收，回收掉没人用的 Block，再执行步骤一。

• 若第一步的 get_free_block 失败，则在第二步中先调用 trigger_free_memory_callbacks，再调用一次第一步的 get_free_block；
• trigger_free_memory_callbacks 本质上通过注册调用了 CudaIPCCollect 函数，进而调用 CudaIPCSentDataLimbo::collect 函数（torch/csrc/CudaIPCTypes.cpp : L64）；
• CudaIPCSentDataLimbo 类管理所有 reference 不为 0 的 Block，所以实际上 collect 函数的调用算是一种懒更新，直到无 Block 可分配的时候才调用来清理那些 reference 已经为 0 的 Block（值得一提的是，该类的析构函数会首先调用 collect 函数，见 torch/csrc/CudaIPCTypes.cpp : L58）；
• 相关源码可以看 torch/csrc/CudaIPCTypes.h & .cpp。

步骤三：alloc_block 函数（L1115）

TLDR：Allocator 在已有的 Block 中找不出可分配的了，就调用 cudaMalloc 创建新的 Block。

• 步骤一、二中重用 block 失败，于是用 cudaMalloc 分配内存，大小为 alloc_size；
• 注意有一个参数 set_fraction 会限制可分配的显存为当前剩余的显存 * fraction（若需要分配的超过这一限制则失败），但还没搞清楚哪里会指定这个（TODO）；
• 新分配的内存指针会被用于创建一个新 Block，新 Block 的 device 与 cuda_stream_id 与 caller 保持一致。

上面几个步骤都是试图找到一些空闲显存，下面是两个步骤是尝试进行碎片整理，凑出一个大块显存

步骤四：release_available_cached_blocks 函数（L1175）

TLDR：先在自己的池子里释放一些比较大的 Block，再用 cudaMalloc 分配看看

• 如果上面的 alloc_block 失败了，就会尝试先调用这一函数，找到比 size 小的 Block 中最大的，由大至小依次释放 Block，直到释放的 Block 大小总和 >= size（需要注意，这一步骤只会释放那些大小大于阈值 max_split_size_mb 的 Block，可以理解为先释放一些比较大的）；
• 释放 block 的函数见 release_block（L1241），主要就是 cudaFree 掉指针，再处理一些 CUDA graphs 的依赖，更新其他数据结构等等，最后把这个 Block 从 BlockPool 中移除；
• 当前整个 BlockPool（作为提醒，Allocator 有两个 BlockPool，这里指的是最初根据 size 指定的大 pool 或小 pool）中，可以释放的 Block 需要满足两个条件：cuda_stream_id 相同的，且大小要大于阈值 max_split_size_mb。如果将这样的 Block 全部释放的空间仍比 size 小，那么这一步就会失败。
• 释放掉了一批 Block 之后，再次执行步骤三中的 alloc_block 函数，创建新 Block。

步骤五：release_cached_blocks 函数（L1214）

• 如果释放一些 Block 还不够分配，则把整个 Allocator 中的 large / small pool 全部释放掉（同样调用 release_block：L1241），再次调用 alloc_block 函数。

分配失败

如果经历上述5个步骤还是没有找到合适的块用于显存申请，则会报出经典的 CUDA out of memory. Tried to allocate ... 错误，例如：

CUDA out of memory. Tried to allocate 1.24 GiB (GPU 0; 15.78 GiB total capacity; 10.34 GiB already allocated; 435.50 MiB free; 14.21 GiB reserved in total by PyTorch)

• Tried to allocate：指本次 malloc 时预计分配的 alloc_size；
• total capacity：由 cudaMemGetInfo 返回的 device 显存总量；
• already allocated：由统计数据记录，当前为止请求分配的 size 的总和；
• free：由 cudaMemGetInfo 返回的 device 显存剩余量；
• reserved：BlockPool 中所有 Block 的大小，与已经分配的 Block 大小的总和。
  即 [reserved] = [already allocated] + [sum size of 2 BlockPools]

注意，reserved + free 并不等同于 total capacity，因为 reserved 只记录了通过 PyTorch 分配的显存，如果用户手动调用 cudaMalloc 或通过其他手段分配到了显存，是没法在这个报错信息中追踪到的（又因为一般 PyTorch 分配的显存占大部分，分配失败的报错信息一般也是由 PyTorch 反馈的）。

在这个例子里，device 只剩 435.5MB，不够 1.24GB，而 PyTorch 自己保留了 14.21GB（储存在 Block 里），其中分配了 10.3GB，剩 3.9GB。那为何不能从这 3.9GB 剩余当中分配 1.2GB 呢？原因肯定是碎片化了，而且是做了整理也不行的情况。

max_split_size_mb

我们已经经历了努力寻找显存但是没有成功的过程，那么当我们空闲显存明显比需要申请的显存大很多时候这个过程就存在不太合理的地方了。解决问题的关键在于 CUDA 中的 max_split_size_mb 变量设置。

官方定义

官方文档：

max_split_size_mb prevents the native allocator from splitting blocks larger than this size (in MB). This can reduce fragmentation and may allow some borderline workloads to complete without running out of memory. Performance cost can range from ‘zero’ to ‘substantial’ depending on allocation patterns. Default value is unlimited, i.e. all blocks can be split. The memory_stats() and memory_summary() methods are useful for tuning. This option should be used as a last resort for a workload that is aborting due to ‘out of memory’ and showing a large amount of inactive split blocks. max_split_size_mb is only meaningful with backend:native. With backend:cudaMallocAsync, max_split_size_mb is ignored.

max_split_size_mb 可以防止本机分配程序拆分大于此大小的块(以 MB 为单位)，以此减少显存碎片，并且可能允许在不耗尽内存的情况下完成一些边界工作负载。

实现逻辑

根据 知乎大神 的Pytorch 显存分配源码解读，max_split_size_mb 的作用应该是限制分配显存时连续空闲显存块的大小的，通过这个阈值降低分配显存时直接拆分大块连续显存的概率。

他起作用的核心在于 get_free_block 函数中：我当前需要申请 size 大小的显存，阈值为 max_split_size_mb，此时我找到了 Block 大小的空闲块：

1. 如果 size < max_split_size_mb (一个小块），但是 Block > max_split_size_mb 则不会直接执行显存分配（避免拆分大块空闲显存，导致 Block 浪费）
2. 如果 size > max_split_size_mb 并且 Block > max_split_size_mb，但是 Block - size > 20MB，也不会执行显存分配，这个机制使得大于 20MB 的显存碎片不那么容易产生

至于 max_split_size_mb 影响数据、模型拆分、是内存分配的最大值等说法，个人不敢苟同。

错误信息

典型的使用 max_split_size_mb 可以大概率解决的错误信息类似这种：

RuntimeError: CUDA out of memory. Tried to allocate 6.18 GiB (GPU 0; 24.00 GiB total capacity; 11.39 GiB already allocated; 3.43 GiB free; 17.62 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

这里 Pytorch 保留显存 17.62 GiB，已分配内存 11.39 GiB，中间 6 个多 g 显存没有充分利用，表示当前碎片化比较严重，这种情况可以尝试降低 max_split_size_mb 的值来降低碎片出现的概率。

修改 max_split_size_mb

直接修改环境变量即可，建议在 Python 运行过程中临时修改，避免不必要的性能降低

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:256'

单位 MB

参考资料

• https://pytorch.org/docs/stable/notes/cuda.html
• https://zhuanlan.zhihu.com/p/486360176
• https://blog.csdn.net/MirageTanker/article/details/127998036

文章链接：
https://www.zywvvd.com/notes/environment/cuda/cuda-outof-memory/cuda-outof-memory/