PyTorch的CUDA OOM报错包含四个关键数值,分别对应三种不同失败原因,每种有不同解决方案。
torch.cuda.OutOfMemoryError: CUDA out of memory 不是一种错误。PyTorch 在这句话后面打印的数字区分了三种不同的失败情形——单个分配过大、设备确实满了、以及分配器有足够空闲内存但不是一块连续的——这三种情况各有不同的解决方法。
措辞在不同 PyTorch 版本间有变化,但大约从 2.1 起,消息携带四个数字。一个代表性的例子:
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB.
GPU 0 has a total capacity of 23.64 GiB of which 1.18 GiB is free.
Of the allocated memory 19.90 GiB is allocated by PyTorch, and 1.42 GiB is
reserved by PyTorch but unallocated. If reserved but unallocated memory is
large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
按以下顺序阅读,遇到第一个匹配的情况就停下来。
"尝试分配的"容量是否巨大?如果一张只有几 GB 空闲的卡上单个分配就是几 GB,那就是一个张量的问题,而不是很多张量的总和。注意力计算的临时空间和完整词表的 logit 是常见的嫌疑对象,两者都与 batch size 乘以序列长度成正比。把这两个因素中的一个减半,分配量也就减半。
"已预留但未分配"是否大于"尝试分配"?那么你并不是真的内存耗尽,而是碎片化了:分配器持有足够总量的空间,但没有足够大的单个连续块。这在使用变长序列时会发生,而这恰恰是推理服务的典型场景。在进程启动前的环境变量中设置 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True——它让分配器可以扩展段而不是需要一块连续的空间,在变长工作负载上可以通过一个环境变量恢复数 GB 的内存。按桶排序序列长度使分配以相同大小重复也有同样效果,但更费力。
"空闲"是否接近零且"已分配"接近容量?这是诚实的情况:工作负载确实放不下。去下文做算术,决定该收缩什么。
"PyTorch 已分配"是否远小于容量减去空闲?如果是,那么卡上有别的东西在占用——另一个进程、一个僵尸进程、一个显示服务器、或者你自己的 worker 的第二份副本。见最后一节。
你可以计算需求,而不是等到崩溃才发现,算术很短,手算就够了。下面的每个数字都是推导出来的,不是直接引用的;用你自己模型的 config.json 中的数字代入。
参数数量乘以每个参数的字节数。FP16 和 BF16 是 2 字节,FP32 是 4 字节,8 位是 1 字节,4 位大约是 0.5 字节加上量化和零点的小额开销。
8e9 params x 2 bytes = 16.0 GB (bf16)
8e9 params x 1 byte = 8.0 GB (int8)
8e9 params x 0.5 byte = 4.0 GB (4-bit, before quantisation overhead)
一张 24 GB 的卡在 bf16 下能放下一个 8B 模型还剩 8 GB,根本放不下一个 70B 的 bf16 模型——光权重就有 140 GB,还没算其他东西。
上下文中的每个 token 都在每一层存储一个 key 向量和一个 value 向量,在整个序列存活期间都保留。每个 token 的量:
bytes_per_token = 2 (K and V)
x layers
x kv_heads # 不是 attention heads,如果模型用了 GQA
x head_dim
x bytes_per_element
代入:32 层、8 个 KV head、head_dim 128、bf16 的模型:
2 x 32 x 8 x 128 x 2 = 131,072 bytes = 128 KiB per token
8,192-token context -> 1.0 GiB per sequence
32,768-token context -> 4.0 GiB per sequence
8 concurrent sequences at 8k -> 8.0 GiB
最后一行是终结推理服务实验的那一行。权重放得下,一个请求放得下,第八个并发请求就放不下了。注意分组查询注意力为你省下的乘数:同样的模型如果用 32 个 KV head 而不是 8,每 token 需要 512 KiB,是四倍。关于 KV 缓存的机制详见 how the KV cache works 和 VRAM requirements sizing。
使用 Adam 进行全量微调,每个参数要保存:权重、梯度、两个优化器矩——通常还要加一份 FP32 主拷贝。混合精度下通常说的是激活之前每个参数约 16 字节:
2 bytes bf16 weights
2 bytes bf16 gradients
4 bytes fp32 master weights
4 bytes Adam first moment
4 bytes Adam second moment
---------
16 bytes per parameter
8B parameters x 16 = 128 GB,还没算一个激活。
这就是为什么 8B 模型的全量微调不会在一张消费级卡上发生,而 LoRA 可以:只训练适配器意味着 99% 的参数去掉了梯度和两个矩。
减少最大并发序列数或最大上下文。两者都直接乘以 KV 项,而服务引擎将它们暴露为明确的限制而不是留给运气。这是第一个可以动手的杠杆,因为它的效果是已知且可以计算的。
对权重做量化。8 位使权重项减半,4 位使权重项变成四分之一,质量代价取决于方法和模型——选择一种量化就是在权衡。注意这不会影响 KV 缓存,所以长上下文工作负载在量化后仍然可能失败。
对 KV 缓存做量化。如果服务引擎支持 8 位 KV,它会使随上下文和并发增长的项减半,而那正是增长的项。
设置 expandable_segments:True。免费,而且对服务产生的变长分配模式特别有效。
检查你不是持有了第二份拷贝。将 checkpoint 加载到 CPU 然后再移到 GPU 可能瞬时需要两份;直接加载到设备,或用低 CPU 内存的加载路径,可以避免这个峰值。
降低每设备 batch size 并增加梯度累积来补偿。有效 batch size 不变,但内存不变。这是第一个要尝试的,代价只是墙上时钟时间。
缩短最大序列长度。激活内存随它增长,而对于物化完整注意力矩阵的注意力实现,它随序列长度的平方增长。检查你的样本实际需要你设置的长度的比例是多少——通常只是一小部分尾部的样本。
开启梯度检查点。在反向传播时重新计算激活而不是存储它们。对于大约 20–30% 更多的计算量,可以大量节省内存,在大多数训练器里只是一个 flag 的事。
切换到 LoRA 或 QLoRA。去掉了冻结权重的优化器状态和梯度,按上面的算术,那是需求的大部分。
使用内存高效的注意力 kernel。FlashAttention 家族的任何实现都避免物化完整的注意力矩阵,从而去掉了二次激活项。
最后把优化器状态卸载到 CPU。这有用但很慢,所以应该放在不花代价的选项之后。
如果 nvidia-smi 显示大部分内存被占用而你相信什么都没在跑,那就是以下情况之一:
nvidia-smi # 谁在占内存,按 PID
nvidia-smi --query-compute-apps=pid,used_memory --format=csv
ps -o pid,user,cmd -p <PID> # 那个 PID 实际上是什么
一个崩溃后没被回收的进程。一个在 CUDA 上下文存活时被 kill 的 Python 进程可以一直持有内存直到进程真正消失。杀死上面列表中的 PID,不是启动它的那 个终端。
一个 notebook kernel。Jupyter 在 cell 执行完并在你关闭标签页后仍然保持 kernel 存活,从而也保持了张量。重启 kernel;仅 del model 不会释放仍被 traceback 或输出 cell 引用的内存。
一张卡上跑了两个 worker。用两个 worker 启动的服务器会加载两次模型。这表现为正好是预期分配量的两倍,一旦你算出了预期数字就很容易发现。
一个显示服务器或其他用户。在共享机器或桌面上,在你开始之前几 GB 可能就没了。
在一个 Python 进程内,torch.cuda.empty_cache() 将已缓存但未分配的内存归还给驱动。它不会释放活跃的张量,也不会挽救真正的短缺——如果它看起来修复了你的 OOM,那你实际遇到的是碎片化,而 expandable_segments 是更好的答案。
dev.to · AI