先把当时的情况摆出来。8 卡 A100 跑 DDP,每张卡比模型加优化器多占 3GB,torch.cuda.memory_allocated 看不到这部分,只有 nvidia-smi 有。

排查方向其实有限,先把可能性排一排。NCCL 通信缓冲区不计入 PyTorch 分配器。另一种情况是每个进程都建了完整通信域。

少数情况下是bucket 太大,梯度同步缓冲翻倍。不容易想到的是,NCCL 共享内存路径没设,走了主机内存映射。

第一件事是对比两种显存口径

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

限制 NCCL 缓冲大小

export NCCL_BUFFSIZE=1048576

然后打开通信域初始化日志

export NCCL_DEBUG=INFO; python train.py --dry-run

剩下的就是按卡统计实际占用

nvidia-smi --query-gpu=index,memory.used --format=csv

改之前先确认了最近有没有人动过配置,避免白查一遍。

这次没动代码,只在配置和运行环境上找原因。

隔了一天再看,现象没变,说明不是偶发。

别只看报错消失,验证要多做一步。每卡多出的显存降到 1GB 以内,训练吞吐不变,OOM 不再在第二步出现。

补一句。nvidia-smi 和 torch 的两套账本来就不相等,差值里通常有 NCCL 和 CUDA context。监控比日志更早给出信号,事后回看经常能提前半小时发现异常。

作者 9823189