先把当时的情况摆出来。训练过程中 torch.cuda.memory_reserved 稳定在 18G,而 allocated 只有 6G,换模型时提示显存不足,但实际计算用量并不高。
原因不止一个,得分开看。缓存分配器保留了空闲块,不主动归还驱动。另一种情况是序列长度变化大,产生大量尺寸不一的碎片。
还有一种情况,变量没有及时释放,引用一直挂在全局对象上。另外,多进程共享显卡时各进程各自缓存。
第一件事是打印显存明细确认差距来源
python -c "import torch;print(torch.cuda.memory_summary())"
然后在安全点主动释放缓存
torch.cuda.empty_cache()
接着把中间结果显式置空并删除引用
del outputs; gc.collect()
等前面都确认过,再设置分配策略减少碎片
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python train.py
排查用的命令都是临时敲的,没有写成脚本,事后才整理。
恢复的判定其实很直接。reserved 与 allocated 的差距缩小到 1G 以内,切换模型时不再报显存不足。
empty_cache 只是把空闲块还给驱动,正在被引用的张量不会被释放,看到显存没降要先查引用而不是反复调用。拖得越久,排查成本越高,因为现场状态会被后续操作覆盖。