使用torchrun启动PyTorch分布式训练时遭遇套接字初始化失败与CUDA设备断言错误

看起来你在把单GPU训练代码改成分布式的过程中，遇到了两个缠人的问题：套接字初始化失败和CUDA设备断言错误。我来帮你一步步排查解决，先从最致命的CUDA问题入手。

一、先解决CUDA设备断言错误（进程崩溃的直接原因） #

你的错误日志里有个关键信息：

RuntimeError: device >= 0 && device < num_gpus INTERNAL ASSERT FAILED at "../aten/src/ATen/cuda/CUDAContext.cpp":50, please report a bug to PyTorch. device=1, num_gpus=

这里num_gpus是空的，说明当前进程根本没检测到可用GPU；同时你指定了--nproc_per_node=4，也就是要启动4个进程，每个进程的LOCAL_RANK分别是0、1、2、3，但如果你的机器实际没有4个可用GPU，或者GPU权限/环境变量没设置对，LOCAL_RANK≥1的进程就会找不到对应的GPU，直接崩溃。

快速排查和修复： #

1. 先确认机器的可用GPU数量
   打开终端运行nvidia-smi，看看你的机器到底有多少个空闲GPU。如果只有1个，那你不能用--nproc_per_node=4，得改成--nproc_per_node=1；如果有多个GPU，要确保它们没被其他进程占用。

2. 正确绑定GPU（避免进程找不到GPU）
   如果你只想用特定的GPU（比如GPU 0和1），启动命令要先设置CUDA_VISIBLE_DEVICES，且--nproc_per_node的数量要和你绑定的GPU数一致：

   CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=2 --nnodes=1 train_dist.py
   

   别像之前那样注释掉CUDA_VISIBLE_DEVICES后就不管了，torchrun启动的多进程会严格按照LOCAL_RANK去绑定GPU，数量不匹配就会报错。

3. 优化分布式初始化函数的逻辑
   你当前的init_distributed和main函数里的判断有问题：比如你用torch.cuda.device_count() > 1来决定是否初始化分布式，但torchrun启动的每个进程都会执行这段代码，而device_count在不同进程里可能因为CUDA_VISIBLE_DEVICES的设置而不同。改成下面的写法更稳妥：

   def init_distributed():
       # torchrun会自动帮你设置好RANK、WORLD_SIZE、LOCAL_RANK环境变量，不用手动传
       local_rank = int(os.environ['LOCAL_RANK'])
       # 先绑定设备再初始化进程组
       torch.cuda.set_device(local_rank)
       dist.init_process_group(backend='nccl')  # 自动读取环境变量里的rank和world_size
       rank = dist.get_rank()
       world_size = dist.get_world_size()
       device = torch.device("cuda", local_rank)
       return device, rank, world_size
   
   def main(args):
       # ... 其他初始化代码（日志、随机种子等）...
       # 用环境变量判断是否是分布式训练，而不是看GPU数量
       is_distributed = 'RANK' in os.environ and 'WORLD_SIZE' in os.environ
       if is_distributed:
           device, rank, world_size = init_distributed()
       else:
           device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
           rank = 0
           world_size = 1
   
       # 加载模型并放到设备上
       model = Model(...)  # 你的模型初始化代码
       model.to(device)
       # 只有分布式环境才用DDP包装模型
       if is_distributed:
           model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
   

二、解决套接字初始化失败的问题 #

这个错误The server socket cannot be initialized on [::]:29500 (errno: 97 - Address family not supported by protocol)，是因为PyTorch默认尝试用IPv6地址，但你的系统可能禁用了IPv6，或者网络配置不支持。

修复方法： #

1. 启动命令手动指定IPv4地址
   启动时加上--master_addr=127.0.0.1（本地IPv4回环地址），并禁用IB设备（如果你的机器没有IB网卡）：

   export NCCL_IB_DISABLE=1
   CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --nnodes=1 --master_addr=127.0.0.1 --master_port=29500 train_dist.py
   

2. 或者在代码里提前设置环境变量
   在调用init_distributed前，强制NCCL使用IPv4的本地回环：

   # 放在import之后，init_distributed之前
   os.environ['NCCL_SOCKET_IFNAME'] = 'lo'  # lo是本地回环网卡
   os.environ['NCCL_IB_DISABLE'] = '1'
   

三、分布式训练的其他注意事项（避免踩坑） #

1. 每个epoch要更新DistributedSampler的种子
   不然每个epoch的采样顺序是一样的，分布式训练的数据采样会重复：

   for epoch in range(args.start_epoch, args.epochs):
       if is_distributed:
           sampler_train.set_epoch(epoch)  # 加这行！
       train_one_epoch(model, criterion, optimizer, lr_scheduler, data_loader, device, epoch, args.print_freq)
   

2. 只让主进程（rank=0）打印和保存模型
   不然每个进程都会打印日志、保存模型，会导致文件冲突和重复输出：

   # 替换你原来的保存模型代码
   if cur_acc > acc and rank == 0:
       acc = cur_acc
       path = os.path.join(args.output_dir, f"model_{epoch}_ntu60_DTr.pth")
       torch.save(model.state_dict(), path)
       print("model saved")
       with open('NTU60_epoch.txt', 'a') as f:
           f.write(str(epoch) + '\n')
   

3. 训练结束后销毁进程组
   在main函数的最后（或者异常捕获里）加上：

   if is_distributed:
       dist.destroy_process_group()
   

备注：内容来源于stack exchange，提问作者GPS-999