通过设置PYTORCH_CUDA_ALLOC_CONF中的max_split_size_mb解决Pytorch的显存碎片化导致的CUDA:Out Of Memory问题

问题的出现

最近在基友的带动下开始投身ai绘画的大潮，于是本地部署了stable diffusion web ui，利用手上的24G显存开始了愉快的跑高分辨率图片之旅。然而某天在用inpaint功能修图扩图过程中突然收到了如下的报错消息：

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

Time taken: 4.06sTorch active/reserved: 17917/18124 MiB, Sys VRAM: 24576/24576 MiB (100.0%)

也就是喜闻乐见的CUDA的OOM.

---

解决过程

因为问题是突然出现，使用的跑图参数在之前也是可行的，所以首先想到是运行环境出了问题，对比了repo文件的改动，尝试重新部署，重装系统然后重新部署，均无果。这时才将注意力转回到报错消息本身：

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

Time taken: 4.06sTorch active/reserved: 17917/18124 MiB, Sys VRAM: 24576/24576 MiB (100.0%)

注意到reserved - allocated = 17.62 - 11.39 = 6.23 > 6.18，而free = 3.43，符合消息中的reserved memory is >> allocated memory.显然显存是够的，但是因为碎片化无法分配。所以开始调查PYTORCH_CUDA_ALLOC_CONF这个环境变量设置。
首先检索到的是一文读懂 PyTorch 显存管理机制一文中的这一段：

关于阈值max_split_size_mb，直觉来说应该是大于某个阈值的 Block 比较大，适合拆分成稍小的几个 Block，但这里却设置为小于这一阈值的 Block 才进行拆分。个人理解是，PyTorch 认为，从统计上来说大部分内存申请都是小于某个阈值的，这些大小的 Block 按照常规处理，进行拆分与碎片管理；但对大于阈值的 Block 而言，PyTorch 认为这些大的 Block 申请时开销大（时间，失败风险），可以留待分配给下次较大的请求，于是不适合拆分。默认情况下阈值变量max_split_size_mb为 INT_MAX，即全部 Block 都可以拆分。

当时的理解是，既然默认值就是最大值，此前没有专门设置过这个变量，那么所有的显存请求都可以被拆分，应该不会出现OOM才对（ 注意，这里的理解是错误的！ 文末会有反思环节，这是后话 ）。于是进一步调查，按照pytorch显存管理文章的思路，在每次Generate的逻辑前后执行empty_cache操作并回显显存使用情况(webui.py，L49)：

def wrap_gradio_gpu_call(func, extra_outputs=None):
    def f(*args, **kwargs):
        os.system('nvidia-smi -i 0')
        print(torch.cuda.memory_allocated())
        print(torch.cuda.max_memory_allocated())
        print(torch.cuda.memory_reserved())
        print(torch.cuda.max_memory_reserved())
        print(torch.cuda.memory_stats())
        print(torch.cuda.memory_snapshot())
        torch.cuda.empty_cache()
        time.sleep(1)
        os.system('nvidia-smi -i 0')
        print(torch.cuda.memory_allocated())

        shared.state.begin()

        with queue_lock:
            res = func(*args, **kwargs)

        shared.state.end()

        return res

    return modules.ui.wrap_gradio_call(f, extra_outputs=extra_outputs, add_stats=True)

观察结果是一切正常，没有什么建设性的发现。empty_cache操作前后显存用量基本没有变化，维持在载入模型后的7~8G左右。
苦思冥想不得解，这时候基友提示把这个PYTORCH_CUDA_ALLOC_CONF环境变量也记录下来确认一下：

os.system("echo %PYTORCH_CUDA_ALLOC_CONF%")

得到的结果是：

%PYTORCH_CUDA_ALLOC_CONF%

……
也就是说Pytorch、gradio和sdweb三者的逻辑中都没有设置这个变量。抱着试试看的心态，在启动脚本前手动设置此变量值为int32类型的上限值，也就是官网和文献1中提到的默认值：

@echo off
set PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:2147483647
set PYTHON=
set GIT=
set VENV_DIR=
set COMMANDLINE_ARGS=--no-half --api --theme dark
call webui.bat

此时再使用同样的参数，逐步提高分辨率测试了多次，都可以正常运行。正当笔者满心欢喜以为问题得解的时候，OOM又出现了。之后再测试，OOM开始高概率随机出现，极少数时候可以正常运行。这时才开始怀疑可能对max_split_size_mb参数值和实际行为的关系理解有误，开始尝试将其值调小为32MB（CUDA报错:Out of Memory），竟然就可以稳定正常运行了，OOM不再出现，问题解决了。

---

结论

TLDR：对于显存碎片化引起的CUDA OOM，解决方法是将PYTORCH_CUDA_ALLOC_CONF的max_split_size_mb设为较小值。

set PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:32

---

逻辑探究

问题虽然解决了，但对于其中的逻辑还是完全一头雾水。由于是改小了才可以正常运行，遂怀疑实际行为与参数的字面意思相反，即大于该值的显存请求才会被拆分。然后就使用参数值逐级逼近报错消息中的显存请求大小6.18GB的方式来测试其实际行为，对一系列参数值进行测试：

max_split_size_mb	测试结果
4096	正常
5120	正常
6144	正常
7168	OOM
8192	OOM
2147483647	OOM

现在拿到了确凿的事实：当参数值小于显存请求的大小时就可以正常运行。这和之前的理解明显是不符的，因为小的显存请求明显是不用分割就可以分配，大的请求才需要分割。这时才回去一文读懂 PyTorch 显存管理机制文章详细阅读前文的分析，发现文中提及的Block是指空闲Block，而不是显存请求。max_split_size_mb分割的对象也是空闲Block（这里有个暗含的前提：pytorch显存管理机制中，显存请求必须是连续的）。
这里实际的逻辑是：由于默认策略是所有大小的空闲Block都可以被分割，所以导致OOM的显存请求发生时，所有大于该请求的空闲Block有可能都已经被分割掉了。而将max_split_size_mb设置为小于该显存请求的值，会阻止大于该请求的空闲Block被分割。如果显存总量确实充足，即可保证大于该请求的空闲Block总是存在，从而避免了分配失败的发生。笔者的情况，显存总量是24G，那么最理想的条件下，大于6.18G的空闲Block最多也只能有3个，这就解释了为什么OOM是高概率随机出现。而报错信息中的”3.43 GiB free”实际上是指pytorch所能找到的最大的空闲Block的大小，而非总的空闲空间大小，所以会出现reserved - allocated > free的现象。CUDA报错:Out of Memory一文中将max_split_size_mb称为“一次分配的最大单位”也是错误的，这个值实际上决定了最大的空闲Block可能的最小值，这个最小值为显存总量 - max_split_size_mb。所以这个变量应该命名为最小空闲Block保留大小语义才更为明确。

---

进一步优化

有了上述结论，就可以导出最优设置策略：将max_split_size_mb设置为小于OOM发生时的显存请求大小最小值的最大整数值，就可以在保证跑大图的可行性的同时最大限度照顾性能。
笔者的观测OOM请求最小值是6.18GB，所以最终选择了6144作为最优设置：

set PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:6144

---

总结

这个故事告诉我们：

1. TLDR（太长不看）不可取，很容易断章取义。快速提取关键信息固然重要，但对于复杂系统还是需要更谨慎一些；
2. 进行逻辑和机制分析的时候尽量使每个段落的内容自洽，尽量不依赖上下文和概念定义。这有助于读者快速检索时得到正确的信息；
3. 亲手验证之前不能相信默认值。只有确认了实际行为才能消除表述和理解之间可能存在的差异，不能想当然。

---

参考文献

1. 一文读懂 PyTorch 显存管理机制
2. CUDA报错:Out of Memory
3. pytorch显存管理
4. CUDA semantics — PyTorch 1.13 documentation