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_DEVICES：

CUDA_VISIBLE_DEVICES=0,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}

这个脚本做了三件关键事情：

1. 通过--nproc_per_node指定每个节点上创建的进程数
2. 每个进程自动获取自己的LOCAL_RANK环境变量
3. 各进程通过NCCL后端进行通信

1.3 新旧版本启动方式对比

特性	MMDetection 2.x	MMDetection 3.x
GPU指定方式	--gpus参数	CUDA_VISIBLE_DEVICES+nproc_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_DEVICES=0,1,2,3 \
python -m torch.distributed.launch \
    --nproc_per_node=4 \
    tools/train.py \
    configs/yolox/yolox_s_8xb8-300e_coco.py \
    --launcher pytorch

关键参数解析：

• --nproc_per_node=4：每个节点创建4个进程
• --launcher pytorch：指定使用PyTorch原生分布式后端

2.3 现代启动方式（推荐）

PyTorch官方推荐使用torchrun替代torch.distributed.launch：

CUDA_VISIBLE_DEVICES=0,1,2,3 \
torchrun \
    --nproc_per_node=4 \
    tools/train.py \
    configs/yolox/yolox_s_8xb8-300e_coco.py \
    --launcher pytorch

改进之处：

1. 自动处理RANK和WORLD_SIZE环境变量
2. 更好的错误处理和进程管理
3. 支持弹性训练（实验性功能）

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 \
    --nnodes=2 \
    --node_rank=0 \
    --nproc_per_node=8 \
    --master_addr="192.168.1.100" \
    --master_port=29500 \
    tools/train.py \
    configs/yolox/yolox_s_8xb8-300e_coco.py \
    --launcher pytorch

从节点执行：

torchrun \
    --nnodes=2 \
    --node_rank=1 \
    --nproc_per_node=8 \
    --master_addr="192.168.1.100" \
    --master_port=29500 \
    tools/train.py \
    configs/yolox/yolox_s_8xb8-300e_coco.py \
    --launcher pytorch

3.2 关键性能参数调优

在config文件中调整这些参数可以提升多GPU训练效率：

# 在config文件中添加/修改
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001)
optimizer_config = dict(grad_clip=None)

# 根据GPU数量调整batch size
data = dict(
    samples_per_gpu=8,  # 单卡batch size
    workers_per_gpu=4,  # 数据加载线程数
)

经验法则：

• samples_per_gpu × nproc_per_node = 总batch size
• workers_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 调试技巧

1. 单卡调试模式：

CUDA_VISIBLE_DEVICES=0 python tools/train.py configs/yolox/yolox_s_8xb8-300e_coco.py

2. 查看分布式环境变量：

import os
print(f"RANK: {os.environ.get('RANK')}")
print(f"WORLD_SIZE: {os.environ.get('WORLD_SIZE')}")
print(f"LOCAL_RANK: {os.environ.get('LOCAL_RANK')}")

3. 强制使用CPU模式（排除GPU问题）：

CUDA_VISIBLE_DEVICES="" torchrun --nproc_per_node=4 tools/train.py config_file.py

4.3 SLURM集群集成

在SLURM环境中，可以直接使用srun命令：

srun -p mm_dev \
    --job-name=mmdet_train \
    --gres=gpu:8 \
    --ntasks=8 \
    --ntasks-per-node=8 \
    --cpus-per-task=5 \
    torchrun \
    --nproc_per_node=8 \
    tools/train.py \
    configs/yolox/yolox_s_8xb8-300e_coco.py \
    --launcher="slurm"

关键参数说明：

• --gres=gpu:8：每节点申请8块GPU
• --ntasks=8：总共8个任务（对应8个GPU）
• --cpus-per-task=5：每个任务分配5个CPU核心

经过多次实践验证，我发现当GPU数量超过4块时，适当增加--cpus-per-task可以显著提升数据加载效率，避免GPU等待数据的情况。例如在8卡训练时，设置为5-8个CPU核心效果最佳。