问题出现得很突然,日志倒是留得挺全。同一时刻 nvidia-smi 显示进程占用九点六 GB,torch 里 allocated 只有九 GB,差的六百 MB 一直存在,怀疑是显存泄漏;连续跑了六个小时,差值没有变大,但团队里没人能说清这块是谁的。

刚开始盯着「nvidia-smi 统计的是进程总量,包含其他库的分配」,折腾了半个多小时,排掉了。

真正麻烦的地方在于,几个原因的现象几乎一样。CUDA 上下文和驱动自身要占几百 MB,torch 不计入。还有一种是cuDNN 与 cuBLAS 的 workspace 由算子库自己管理。

不容易想到的是,确实有张量没释放,只是被引用计数藏住了。

同时取两边的数值做对比

nvidia-smi --query-compute-apps=pid,used_memory --format=csv; python -c "import torch;print(torch.cuda.memory_allocated()/2**20)"

接着看保留量与分配量的差额有多大

python -c "import torch;print(torch.cuda.memory_reserved()/2**20 - torch.cuda.memory_allocated()/2**20)"

随后按分钟采样观察差值是否增长

python watch_mem.py --interval 60 --duration 3600 --out mem.csv

剩下的就是用分段汇总排除异常大块

python -c "import torch;print(torch.cuda.memory_summary(abbreviated=True))"

别只看报错消失,验证要多做一步。六小时里两边差值稳定在六百 MB 上下,保留量与分配量同步回落,没有单调增长,可以判定不是泄漏。

torch 只统计自己分配器管的内存,驱动上下文和算子库工作区都在外面,几百 MB 的固定差值属于正常。同一份配置在测试环境跑通,不代表线上没事,两边的默认值经常不一样。

作者 9823189