这类问题我又碰上了一次。4 卡 A10 跑 DDP 训练,单卡能跑 batch 8,4 卡一起跑 batch 8 就报 NCCL out of memory,nvidia-smi 里每卡都还剩 3GB 却起不来。

第一反应是「每卡进程还各载了一份完整权重与优化器」,验证下来不成立。

真正麻烦的地方在于,几个原因的现象几乎一样。NCCL 预分配通信缓冲,剩余显存被算作不可用;bucket 过大导致 allreduce 时额外拷贝一份梯度;P2P 通道开启后显存映射额外占用。

上来先看每卡 NCCL 相关环境变量是否被设过

python -c "import os;print({k:v for k,v in os.environ.items() if k.startswith('NCCL')})"

调小 bucket 上限减少梯度缓冲

python -m torch.distributed.run --nproc_per_node 4 train.py --bucket-cap-mb 25

然后关掉 P2P 与 IB 再试排除通道占用

NCCL_P2P_DISABLE=1 NCCL_IB_DISABLE=1 python -m torch.distributed.run --nproc_per_node 4 train.py

降到 batch 6 并观察是否仍 OOM

python -m torch.distributed.run --nproc_per_node 4 train.py --batch 6 --accum 2

那几天正好在做别的变更,先把时间线对齐才敢动手。

排查用的命令都是临时敲的,没有写成脚本,事后才整理。

改完先别撤,确认这几项。4 卡 batch 8 能正常起跑,allreduce 阶段不再 OOM,每卡显存峰值留有 2GB 以上余量。

NCCL 的缓冲是在剩余显存里预留的,表面看还有空间其实已经被通信占住,调 bucket 和关通道比加卡更直接。监控比日志更早给出信号,事后回看经常能提前半小时发现异常。

作者 9823189