先把当时的情况摆出来。训练过程中 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 只是把空闲块还给驱动,正在被引用的张量不会被释放,看到显存没降要先查引用而不是反复调用。拖得越久,排查成本越高,因为现场状态会被后续操作覆盖。

作者 9823189