单机器双GPU训练遇DistStoreError超时问题求助

问题描述 #

在单机双GPU环境下使用PyTorch训练模型时，调用_init_dist_pytorch('nccl')环节卡在TCPStore初始化步骤：

• 设置world_size=1时，进程组可正常初始化；
• 设置world_size=2时，触发超时错误：

DistStoreError: Timed out after 300 seconds waiting for clients. 1/2 clients joined.

已尝试的无效方案 #

• 调整MASTER_PORT值（测试过23455、8090、12345等）；
• 配置防火墙（开放对应端口）或直接关闭防火墙（执行sudo ufw disable）；
• 使用python -m torch.distributed.launch命令启动脚本。

额外信息 #

• 在Jupyter Notebook中直接调用TCPStore构造函数，world_size=2可正常执行；但通过设置环境变量调用_init_dist_pytorch时触发超时；
• 环境依赖：Python 3.10.14，PyTorch相关生态库。

解决建议 #

1. 检查环境变量传递完整性 #

用分布式启动命令时，确保每个进程能正确获取MASTER_ADDR、MASTER_PORT、RANK、WORLD_SIZE这些核心环境变量。可在脚本开头添加打印代码验证：

import os
print("MASTER_ADDR:", os.environ.get("MASTER_ADDR"))
print("MASTER_PORT:", os.environ.get("MASTER_PORT"))
print("RANK:", os.environ.get("RANK"))
print("WORLD_SIZE:", os.environ.get("WORLD_SIZE"))

手动设置环境变量时，需保证rank0和rank1进程的RANK分别设为0和1，WORLD_SIZE统一设为2。

2. 改用torchrun替代torch.distributed.launch #

torch.distributed.launch已被官方弃用，torchrun能更可靠地管理进程与环境变量传递，启动命令示例：

torchrun --nproc_per_node=2 your_training_script.py

无需手动设置RANK、WORLD_SIZE等变量，torchrun会自动分配。

3. 排查localhost解析问题 #

检查/etc/hosts文件，确保127.0.0.1与localhost正确映射。若localhost被解析到IPv6地址（::1），可能导致TCPStore通信失败，可强制设置MASTER_ADDR=127.0.0.1而非localhost。

4. 确认GPU设备可见性 #

执行以下代码验证两个GPU是否都能被PyTorch识别：

import torch
print(torch.cuda.device_count())
print([torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])

若仅识别到一个GPU，会导致world_size=2时无法启动第二个进程。

5. 手动调试TCPStore初始化 #

在脚本中模拟_init_dist_pytorch的逻辑手动初始化TCPStore，定位问题环节：

import os
import torch.distributed as dist
from torch.distributed import TCPStore

master_addr = os.environ.get("MASTER_ADDR", "127.0.0.1")
master_port = int(os.environ.get("MASTER_PORT", 29500))
world_size = 2
rank = int(os.environ.get("RANK", 0))

store = TCPStore(master_addr, master_port, world_size, rank == 0, timeout=300)
print("TCPStore initialized successfully")
dist.init_process_group(backend="nccl", store=store, rank=rank, world_size=world_size)

若手动初始化正常，说明_init_dist_pytorch内部的环境变量处理存在问题，需检查该函数的实现逻辑。

6. 更新PyTorch及相关依赖 #

旧版本PyTorch可能存在TCPStore的已知bug，尝试升级到稳定版本：

pip install --upgrade torch torchvision torchaudio

内容的提问来源于stack exchange，提问作者孙悟空