1、显存溢出初探:什么是显存溢出?为什么会出现?常见的错误信息解读(CUDA out of memory)
大家好,我是你们的老朋友。今天咱们来聊聊深度学习工程师最头疼的问题之一——显存溢出。
说实话,我入行头两年,被这个错误折磨得够呛。每次看到那个红色的报错,心里就咯噔一下。后来摸爬滚打久了,才慢慢摸清它的脾气。今天就把这些经验分享给你。
1.1 什么是显存溢出?
显存溢出,英文叫 GPU Out of Memory (OOM)。说白了就是:你的显卡内存不够用了。
你想想看,GPU 就像一个小仓库。模型参数、中间计算结果、梯度、优化器状态……这些东西全都要塞进去。仓库满了,再往里塞东西,自然就爆了。
我习惯把显存比作一个「临时工作台」。CPU 是厨房的储物柜,GPU 是灶台。你炒菜时,不可能把所有食材都堆在灶台上——放不下。同样,训练模型时,GPU 只能放当前需要的数据。
核心概念:显存溢出 = GPU 显存容量 < 当前任务所需显存
1.2 为什么会发生显存溢出?
原因其实就几个,我帮你捋一捋。
1.2.1 模型太大
这是最直接的原因。模型参数量越大,占的显存就越多。比如一个 7B 参数的模型,光参数就要占 14GB 左右(FP16 精度)。
我记得有一次,团队里的小伙子直接加载了一个 13B 的模型,显卡只有 24GB。结果可想而知——OOM。我当时跟他说:「兄弟,你这就像用自行车拉大象,不爆才怪。」
1.2.2 批次大小(batch size)太大
很多人以为 batch size 越大越好。其实不然。每个样本在 GPU 里都会产生中间激活值,batch size 翻倍,显存占用也差不多翻倍。
我见过有人把 batch size 设成 256,显卡只有 8GB。结果训练了 3 分钟就崩了。后来改成 32,跑得稳稳的。
1.2.3 序列长度太长
对于 NLP 模型,序列长度直接影响显存。Transformer 的注意力机制复杂度是 O(n²),序列长度翻倍,显存占用翻四倍。
嗯,这里要注意:很多人只关注模型参数量,忽略了序列长度的影响。我曾经踩过这个坑,调了三天才发现是序列长度的问题。
1.2.4 显存泄漏
这个比较隐蔽。代码里有些变量没释放,或者 DataLoader 的 num_workers 设置不当,导致显存一点点被吃掉,最后爆掉。
我曾经遇到过一个 case:训练到第 100 个 epoch 时突然 OOM。排查了半天,发现是某个自定义的损失函数里创建了临时张量,但没有及时删除。
1.3 常见的错误信息解读
咱们来看看最常见的报错长什么样。
1.3.1 CUDA out of memory
这是最经典的报错。通常长这样:
RuntimeError: CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 8.00 GiB total capacity; 6.32 GiB already allocated; 1.56 GiB free; 7.88 GiB reserved in total by PyTorch)
我来帮你拆解一下这个信息:
| 字段 | 含义 |
|---|---|
| Tried to allocate 2.00 GiB | 当前操作试图申请 2GB 显存 |
| 8.00 GiB total capacity | 你的显卡总共有 8GB 显存 |
| 6.32 GiB already allocated | 已经用了 6.32GB |
| 1.56 GiB free | 还剩 1.56GB 空闲 |
| 7.88 GiB reserved in total | PyTorch 总共预留了 7.88GB |
看到没?虽然显示有 1.56GB 空闲,但 PyTorch 已经预留了 7.88GB。它想再申请 2GB,但预留空间不够了,所以就报错。
小技巧:报错里的「reserved」不等于「allocated」。PyTorch 会预先申请一大块显存,然后慢慢分配。所以即使显示有空闲,也可能申请失败。
1.3.2 其他常见报错
- CUDA error: out of memory——底层 CUDA 驱动报的,比 PyTorch 的报错更底层
- RuntimeError: CUDA error: an illegal memory access was encountered——这个可能是显存泄漏或者野指针
- Process finished with exit code -1073740791 (0xC0000409)——Windows 下的显存崩溃,通常伴随 OOM
1.4 知识体系总览
为了让你更直观地理解显存溢出的全貌,我画了一张图:
1.5 避坑指南
我曾经犯过的错:刚入行时,我习惯把所有的数据都搬到 GPU 上,包括验证集和测试集。结果显存很快就爆了。后来才明白,验证和测试时不需要保留梯度,可以用 torch.no_grad() 来节省显存。
另外,我建议你养成一个好习惯:每次训练前,先用 nvidia-smi 看看显存占用。如果发现显存被其他进程占着,先杀掉再跑。
实用命令:
# 查看显存占用
nvidia-smi
# 实时监控(每秒刷新)
watch -n 1 nvidia-smi
# 查看 PyTorch 显存分配
import torch
print(torch.cuda.memory_summary())
好了,这一章就到这里。显存溢出其实没那么可怕,关键是要理解它的成因,学会看报错信息。下一章咱们会深入聊聊如何用工具定位显存问题。
想要更多资料,联系可可