torchrun

PyTorch 分布式训练启动器

补充说明

torchrun 是 PyTorch 官方分布式启动工具，替代 torch.distributed.launch。它负责设置 RANK、WORLD_SIZE、LOCAL_RANK 等环境变量，并支持容错重启。

常用示例

单机多卡：

torchrun --standalone --nproc_per_node=4 train.py

多机多卡：

torchrun \
  --nnodes=2 \
  --nproc_per_node=8 \
  --node_rank=0 \
  --master_addr=10.0.0.10 \
  --master_port=29500 \
  train.py

代码中初始化进程组：

import torch.distributed as dist
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

按 rank 控制日志和保存：

if dist.get_rank() == 0:
    print("only rank 0 logs")

容错启动：

torchrun --nnodes=2 --nproc_per_node=8 \
  --rdzv_backend=c10d --rdzv_endpoint=10.0.0.10:29500 \
  --max_restarts=3 train.py

易错点与排错

• master_addr 使用 localhost，多机场景其他节点连不上。
• node_rank 重复或缺失，进程组初始化卡住。
• 每个进程都写日志或保存模型，文件互相覆盖。
• 只设置了 CUDA_VISIBLE_DEVICES，没有按 LOCAL_RANK 设置当前设备。
• NCCL 网卡选择错误，跨机通信走管理网。
• 某个 rank OOM 后整个作业挂起，需要设置超时和监控。
• DataLoader 的 num_workers 过高，CPU 和内存成为瓶颈。
• 没有使用 dist.barrier() 或正确的 samplers，数据重复或遗漏。

读完自测

合上页面，先回答，再展开对照。

1. 多机训练时 master_addr 写成 localhost 会怎样？
2. 只设置 CUDA_VISIBLE_DEVICES 就够了吗？
3. 每个 rank 都保存模型或写日志会有什么问题？