现场是这样的。PyTorch 2.4 训练脚本报 RuntimeError: CUDA out of memory. Tried to allocate 2.00 GiB,但 nvidia-smi 显示还有 6GB free,torch.cuda.mem_get_info 也显示空闲,重启进程后仍然一样。

中间走过弯路,把「没开 expandable_segments,分配器无法合并」当成主因,改完现象没变。

把日志和配置对完,能解释这个现象的基本就这几处。显存被碎片化成很多小块,凑不出连续的 2GB;其他进程占用但 nvidia-smi 按 PID 没统计到;模型并行组各自预留了显存,实际可用比显示少。

第一件事是打印分配器的显存明细

python -c "import torch;print(torch.cuda.mem_get_info());print(torch.cuda.memory_summary())"

接着开启可扩展段后重跑同一 batch

python -c "import os;os.environ['PYTORCH_CUDA_ALLOC_CONF']='expandable_segments:True';import torch;print(torch.cuda.is_available())"

再往下找出真正占用显存的进程

nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv

等前面都确认过,再把 batch 拆小验证是否是连续块问题

python train.py --batch-size 4 --grad-accum 4

同一份数据我导出了两份,一份留底,一份拿来试。

怎么算真的好了,看的是同一批现象。同样的 batch 配置能正常分配,memory_summary 里 reserved 与 allocated 的差值缩小,不再出现 Tried to allocate。

free 是总量,能不能分配要看有没有连续块;expandable_segments 能缓解碎片,但根本办法还是减少峰值和拆 batch。线上能稳定复现的问题都不算难,难的是只在大流量下才出现的那种。

作者 9823189