告别手动配置：用torch.distributed.launch和torchrun轻松启动PyTorch分布式训练
如果你曾经尝试过手动配置PyTorch分布式训练环境，一定对那些繁琐的环境变量设置记忆犹新——RANK、WORLD_SIZE、MASTER_ADDR、MASTER_PORT...每次启动训练都需要小心翼翼地设置这些参数，稍有不慎就会导致进程间通信失败。更糟糕的是，在多机多卡场景下，这种手动管理方式几乎是一场噩梦。幸运的是，PyTorch为我们提供了两个强大的工具来简化这一过程：torch.distributed.launch脚本和torchrun命令（PyTorch 1.9+）。本文将带你深入了解这些工具的使用方法，让你彻底告别手动配置的烦恼。
1. 为什么需要自动化启动工具
在传统的PyTorch分布式训练中，我们需要手动设置大量环境变量和参数。以一个简单的双卡训练为例，通常需要这样配置：
import os import torch.distributed as dist os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '29500' os.environ['RANK'] = '0' # 手动指定当前进程的rank os.environ['WORLD_SIZE'] = '2' # 总进程数 dist.init_process_group(backend='nccl')这种方式存在几个明显的问题：
- 容易出错：手动指定RANK容易混淆，特别是在多机环境下
- 管理困难：每个进程需要单独启动，并正确传递参数
- 扩展性差：当GPU数量变化时，需要修改大量代码
- 可维护性低：不同环境下的启动脚本难以复用
torch.distributed.launch和torchrun正是为了解决这些问题而生的工具。它们能够：
- 自动计算并设置RANK和WORLD_SIZE
- 统一管理所有训练进程
- 简化多机多卡配置
- 提供更优雅的错误处理和日志记录
2. torch.distributed.launch使用指南
torch.distributed.launch是PyTorch早期提供的分布式训练启动脚本，虽然在新版本中逐渐被torchrun取代，但仍然是许多现有项目的选择。
2.1 基本用法
单机多卡训练（4卡）的启动命令如下：
python -m torch.distributed.launch --nproc_per_node=4 train.py这个简单的命令会自动：
- 启动4个进程（对应4张GPU）
- 为每个进程设置正确的RANK和LOCAL_RANK
- 配置WORLD_SIZE为4
- 设置默认的MASTER_ADDR和MASTER_PORT
2.2 关键参数解析
| 参数 | 说明 | 示例值 |
|---|---|---|
--nproc_per_node | 每个节点上的进程数（通常等于GPU数） | 4 |
--nnodes | 总节点数 | 2 |
--node_rank | 当前节点的rank（多机时需要） | 0 |
--master_addr | 主节点IP地址 | 192.168.1.100 |
--master_port | 主节点端口号 | 29500 |
--use_env | 通过环境变量传递RANK等信息 | (无值) |
2.3 多机多卡配置示例
假设有两台机器，每台有4张GPU：
主节点（192.168.1.100）:
python -m torch.distributed.launch \ --nproc_per_node=4 \ --nnodes=2 \ --node_rank=0 \ --master_addr="192.168.1.100" \ --master_port=29500 \ train.py从节点（192.168.1.101）:
python -m torch.distributed.launch \ --nproc_per_node=4 \ --nnodes=2 \ --node_rank=1 \ --master_addr="192.168.1.100" \ --master_port=29500 \ train.py注意：多机训练时需要确保节点间网络畅通，防火墙开放指定端口
3. torchrun：更现代的解决方案
从PyTorch 1.9开始，官方推荐使用torchrun替代torch.distributed.launch。它提供了更简洁的语法和更强的功能。
3.1 torchrun的优势
- 自动故障恢复：worker失败时会自动重启
- 弹性训练支持：可以动态调整worker数量
- 更简单的参数：合并了一些冗余选项
- 更好的错误处理：提供更清晰的错误信息
3.2 基本使用示例
单机4卡训练：
torchrun --nproc_per_node=4 train.py多机训练（2节点，每节点4卡）：
# 主节点 torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr="192.168.1.100" --master_port=29500 train.py # 从节点 torchrun --nproc_per_node=4 --nnodes=2 --node_rank=1 --master_addr="192.168.1.100" --master_port=29500 train.py3.3 弹性训练配置
torchrun支持弹性训练，允许worker数量在运行时变化。创建一个配置文件elastic_config.json：
{ "min_size": 2, "max_size": 8, "nproc_per_node": 4 }然后启动：
torchrun --rdzv_conf="elastic_config.json" train.py4. 在代码中获取分布式信息
使用启动工具后，我们的训练脚本可以简化为：
import torch.distributed as dist def main(): # 初始化分布式环境 dist.init_process_group(backend='nccl') # 获取自动设置的参数 rank = dist.get_rank() world_size = dist.get_world_size() local_rank = int(os.environ['LOCAL_RANK']) print(f"Rank {rank}/{world_size} (local: {local_rank}) is ready") # 确保每张GPU处理不同的数据 torch.cuda.set_device(local_rank) # 构建模型并移动到当前GPU model = build_model().cuda() # 使用DistributedDataParallel包装模型 model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank]) # 训练逻辑... if __name__ == "__main__": main()关键点说明：
- LOCAL_RANK：工具会自动设置这个环境变量，表示当前节点上的GPU索引
- 设备设置：使用
torch.cuda.set_device确保每个进程使用正确的GPU - 模型包装：
DistributedDataParallel会自动处理梯度同步
5. 常见问题与最佳实践
5.1 端口冲突问题
当多个训练任务同时运行时，可能会遇到端口冲突。解决方法：
- 显式指定不同的
--master_port - 使用脚本自动选择可用端口：
MASTER_PORT=$((29500 + RANDOM % 1000)) torchrun --master_port=$MASTER_PORT --nproc_per_node=4 train.py5.2 数据加载注意事项
在分布式训练中，数据加载需要特殊处理：
from torch.utils.data.distributed import DistributedSampler dataset = MyDataset() sampler = DistributedSampler(dataset, shuffle=True) dataloader = DataLoader(dataset, batch_size=64, sampler=sampler)关键点：
- 每个epoch开始前调用
sampler.set_epoch(epoch)保证shuffle正确性 - 不要在自己的Dataset中实现shuffle
5.3 日志记录策略
在分布式环境中，直接打印日志会导致混乱。推荐做法：
if dist.get_rank() == 0: print("只有rank 0会打印这条消息")或者使用更专业的日志库：
import logging logging.basicConfig( level=logging.INFO if dist.get_rank() == 0 else logging.WARN ) logger = logging.getLogger(__name__)6. 性能优化技巧
6.1 后端选择
PyTorch支持多种分布式后端：
| 后端 | 适用场景 | 特点 |
|---|---|---|
| NCCL | 多GPU训练 | 针对NVIDIA GPU优化，性能最佳 |
| Gloo | CPU训练或多机训练 | 稳定性好，支持CPU |
| MPI | HPC环境 | 需要系统支持MPI |
通常推荐使用NCCL：
dist.init_process_group(backend='nccl')6.2 梯度压缩
对于带宽受限的环境，可以考虑梯度压缩：
model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[local_rank], gradient_as_bucket_view=True # 启用梯度分桶 )6.3 通信重叠
通过调整bucket_cap_mb参数优化通信：
model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[local_rank], bucket_cap_mb=25 # 调整桶大小 )7. 从launch迁移到torchrun
如果你现有的项目使用torch.distributed.launch，迁移到torchrun非常简单。主要变化：
- 去掉
--use_env参数（torchrun默认使用环境变量） - 将
python -m torch.distributed.launch替换为torchrun - 检查代码中对
LOCAL_RANK的使用（torchrun会确保设置这个变量）
例如，原来的启动命令：
python -m torch.distributed.launch --nproc_per_node=4 --use_env train.py迁移后：
torchrun --nproc_per_node=4 train.py在实际项目中，我发现torchrun的自动故障恢复功能特别有用，尤其是在长时间训练任务中。曾经有一次训练因为节点重启而中断，torchrun自动恢复了训练进度，节省了大量时间。
