一.引言

[ torch.cuda.OutOfMemoryError: CUDA out of memory. ] 相信这个错误大家都不陌生，在使用 GPU 进行单机单卡，单机多卡的训练任务中，经常遇到该报错。由于 OOM 是很明显的内存不足告警，通常情况下笔者发现错误就直接去 kill 其他任务或者加内存了，很少关心任务中具体的报错信息，今天趁着训练又报错来整理下。

二.报错日志

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 50.00 MiB 
(GPU 0; 79.35 GiB total capacity; 78.25 GiB already allocated; 10.12 MiB free; 
78.46 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

1.日志数据

torch.cuda.OutOfMemoryError: CUDA out of memory. 
Tried to allocate 50.00 MiB 
(GPU 0; 
    79.35 GiB total capacity; 
    78.25 GiB already allocated; 
    10.12 MiB free;
    78.46 GiB reserved in total by PyTorch)

为了日志更加清晰，这里我把日志逐行分开：

torch.cuda.OutOfMemoryError - 表示发生了 CUDA 内存不足的错误
Tried to allocate 50.00 MB - 尝试分配的内存大小为 50.00 MB
GPU 0 - 指示错误发生在 Device-0 即第一个 GPU 设备上
79.35 GIB total capacity - 表示该 GPU 的总内存容量为 79.35 GB
78.25 GIB already allocated - 已被其他任务或操作占用显存大小为 78.25 GB
10.12 MiB free - 当前可用的空闲显存大小为 10.12 MB
78.46 GiB reserver in total by PyTorch - 由 PyTorch 保留的总显存大小为 78.46 MIB

这个报错信息代表在进行 CUDA 操作时，GPU 内存不足以分配所需要的 50 MB 所以报错。

2.日志提示

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

除了给出详细的内存占用数据外，报错信息中还给出了该报错的可能解决方案：

◆ max_split_size_mb

如果保留内存 reserver memory 远远大于已分配内存 allocated memory，可以尝试调整 max_split_size_mb 参数来避免内存碎片化。该方法通常用语大规模模型和数据集，意在减少内存碎片导致的性能问题。这里较小的值可能会增加内存分配次数，但减少了单个分配的内存大小，有助于缓解内存碎片化问题。较大的值则可以减少内存分配次数，但可能导致更大的内存碎片。

import os

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

◆ See documentation

除了尝试调整 max_split_size_mb 适配自己的模型和数据外，也可以到 Pytorch 文档了解更多内存管理的知识：CUDA semantics — PyTorch 2.0 documentation

有兴趣的同学也可以阅读文档获取更多相关的 PyTorch 内存信息。

◆ 更多解决方案

模型 / 数据集大小 - 减少模型、数据集大小，减少对 GPU 内存的需求
batch_size - 调整训练过程中的 batch_size 降低内存使用量
分布式 - 如果有多个 GPU 设备，可以使用多卡分布式训练
代码优化 - 例如减少重复计算，使用更高效的实现

这里最暴力的方法就是堆资源增加显存。

三.日志扩展

1.Reserved Memory 是什么？

在 PyTorch 中，保留内存 Reserved Memory 是指由 PyTorch 在 GPU 上预先分配和保留的显存。这部分显存不会被释放用于其他目的，而是用于缓冲、缓存和其他内部操作，以提高计算性能和效率。

当使用 PyTorch 进行训练或推断时，PyTorch 会为了优化内存管理而保留一定量的显存。这样，即使释放了一些不再需要的变量或张量，这部分显存也不会立即被释放，而是保留给PyTorch 供后续使用。保留内存的目的是避免频繁的内存分配和释放操作，从而减少内存碎片化和提高计算性能。然而，过多的保留内存可能导致整体可用内存减少，进而导致出现 "CUDA out of memory" 错误。如果你的模型特别大或数据集很大，保留内存可能会占据较大的一部分显存，因此可能需要根据具体情况进行调整。