MMDetection 3.x多GPU训练实战从torch.distributed.launch到torchrun的完整迁移指南最近在部署YOLOX模型训练时发现一个有趣的现象明明通过CUDA_VISIBLE_DEVICES指定了四块GPU但nvidia-smi显示只有GPU0在疯狂工作其他三块卡却在悠闲地围观。这场景就像让四个工人搬砖结果三个人站着看一个人干活——不仅效率低下资源浪费更是让人心疼。经过一番排查才发现这是MMDetection 3.x版本升级后带来的甜蜜陷阱。1. 多GPU训练失效的真相从表面现象到底层原理1.1 为什么CUDA_VISIBLE_DEVICES不再有效在MMDetection 2.x时代我们可以通过--gpus参数轻松指定使用的GPU数量就像这样python tools/train.py configs/yolox/yolox_s_8xb8-300e_coco.py --gpus 4但升级到3.x后这个便捷参数突然消失了。很多开发者包括我的第一反应是改用CUDA_VISIBLE_DEVICESCUDA_VISIBLE_DEVICES0,1,2,3 python tools/train.py configs/yolox/yolox_s_8xb8-300e_coco.py问题在于CUDA_VISIBLE_DEVICES只是环境变量它的作用仅仅是让程序看得见指定的GPU设备但并不会自动启用PyTorch的分布式训练功能。这就好比给工人发放了工作证可见性但没有分配具体任务并行计算。1.2 dist_train.sh背后的魔法打开MMDetection 3.x的tools/dist_train.sh脚本会发现它实际上是通过torch.distributed.launch来启动多进程训练的#!/usr/bin/env bash CONFIG$1 GPUS$2 python -m torch.distributed.launch \ --nproc_per_node$GPUS \ $(dirname $0)/train.py \ $CONFIG \ --launcher pytorch ${:3}这个脚本做了三件关键事情通过--nproc_per_node指定每个节点上创建的进程数每个进程自动获取自己的LOCAL_RANK环境变量各进程通过NCCL后端进行通信1.3 新旧版本启动方式对比特性MMDetection 2.xMMDetection 3.xGPU指定方式--gpus参数CUDA_VISIBLE_DEVICESnproc_per_node启动命令直接调用train.py通过dist_train.sh或torchrun启动分布式后端自动选择需显式指定--launcher pytorch多机支持需要复杂配置标准化参数nnodes, node_rank等2. 正确配置多GPU训练的四步法则2.1 硬件准备与验证在开始之前先用这些命令确认GPU状态nvidia-smi # 查看GPU列表 nvcc --version # 确认CUDA版本 python -c import torch; print(torch.cuda.device_count()) # PyTorch识别的GPU数量注意如果这些命令显示的结果不一致可能是驱动或CUDA环境有问题需要先解决基础环境问题。2.2 传统启动方式即将废弃虽然即将被淘汰但理解原有方式有助于掌握原理CUDA_VISIBLE_DEVICES0,1,2,3 \ python -m torch.distributed.launch \ --nproc_per_node4 \ tools/train.py \ configs/yolox/yolox_s_8xb8-300e_coco.py \ --launcher pytorch关键参数解析--nproc_per_node4每个节点创建4个进程--launcher pytorch指定使用PyTorch原生分布式后端2.3 现代启动方式推荐PyTorch官方推荐使用torchrun替代torch.distributed.launchCUDA_VISIBLE_DEVICES0,1,2,3 \ torchrun \ --nproc_per_node4 \ tools/train.py \ configs/yolox/yolox_s_8xb8-300e_coco.py \ --launcher pytorch改进之处自动处理RANK和WORLD_SIZE环境变量更好的错误处理和进程管理支持弹性训练实验性功能2.4 验证多GPU是否真正工作执行训练后打开另一个终端窗口运行watch -n 1 nvidia-smi健康的多GPU训练应该看到所有指定GPU的显存占用相近所有GPU的Utilization都在波动没有明显的GPU间通信瓶颈如某个GPU的Utilization持续100%3. 高级配置与性能调优3.1 多机多卡配置示例假设有两台机器每台8卡主节点IP为192.168.1.100主节点执行torchrun \ --nnodes2 \ --node_rank0 \ --nproc_per_node8 \ --master_addr192.168.1.100 \ --master_port29500 \ tools/train.py \ configs/yolox/yolox_s_8xb8-300e_coco.py \ --launcher pytorch从节点执行torchrun \ --nnodes2 \ --node_rank1 \ --nproc_per_node8 \ --master_addr192.168.1.100 \ --master_port29500 \ tools/train.py \ configs/yolox/yolox_s_8xb8-300e_coco.py \ --launcher pytorch3.2 关键性能参数调优在config文件中调整这些参数可以提升多GPU训练效率# 在config文件中添加/修改 optimizer dict(typeSGD, lr0.01, momentum0.9, weight_decay0.0001) optimizer_config dict(grad_clipNone) # 根据GPU数量调整batch size data dict( samples_per_gpu8, # 单卡batch size workers_per_gpu4, # 数据加载线程数 )经验法则samples_per_gpu×nproc_per_node 总batch sizeworkers_per_gpu建议设置为GPU数量的1-2倍学习率通常需要随总batch size线性缩放4. 常见陷阱与解决方案4.1 典型错误案例案例一忘记设置--launcher pytorch报错RuntimeError: Expected to have finished reduction in the prior iteration 解决确保命令中包含--launcher pytorch参数案例二端口冲突报错Address already in use 解决更改--master_port默认29500案例三GPU显存不足报错CUDA out of memory 解决减小samples_per_gpu或使用梯度累积4.2 调试技巧单卡调试模式CUDA_VISIBLE_DEVICES0 python tools/train.py configs/yolox/yolox_s_8xb8-300e_coco.py查看分布式环境变量import os print(fRANK: {os.environ.get(RANK)}) print(fWORLD_SIZE: {os.environ.get(WORLD_SIZE)}) print(fLOCAL_RANK: {os.environ.get(LOCAL_RANK)})强制使用CPU模式排除GPU问题CUDA_VISIBLE_DEVICES torchrun --nproc_per_node4 tools/train.py config_file.py4.3 SLURM集群集成在SLURM环境中可以直接使用srun命令srun -p mm_dev \ --job-namemmdet_train \ --gresgpu:8 \ --ntasks8 \ --ntasks-per-node8 \ --cpus-per-task5 \ torchrun \ --nproc_per_node8 \ tools/train.py \ configs/yolox/yolox_s_8xb8-300e_coco.py \ --launcherslurm关键参数说明--gresgpu:8每节点申请8块GPU--ntasks8总共8个任务对应8个GPU--cpus-per-task5每个任务分配5个CPU核心经过多次实践验证我发现当GPU数量超过4块时适当增加--cpus-per-task可以显著提升数据加载效率避免GPU等待数据的情况。例如在8卡训练时设置为5-8个CPU核心效果最佳。