别再手动写rank和world_size了！用torch.distributed.launch和torchrun启动PyTorch分布式训练（保姆级教程）

告别手动配置用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(backendnccl)这种方式存在几个明显的问题容易出错手动指定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_node4 train.py这个简单的命令会自动启动4个进程对应4张GPU为每个进程设置正确的RANK和LOCAL_RANK配置WORLD_SIZE为4设置默认的MASTER_ADDR和MASTER_PORT2.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_node4 \ --nnodes2 \ --node_rank0 \ --master_addr192.168.1.100 \ --master_port29500 \ train.py从节点192.168.1.101:python -m torch.distributed.launch \ --nproc_per_node4 \ --nnodes2 \ --node_rank1 \ --master_addr192.168.1.100 \ --master_port29500 \ train.py注意多机训练时需要确保节点间网络畅通防火墙开放指定端口3. torchrun更现代的解决方案从PyTorch 1.9开始官方推荐使用torchrun替代torch.distributed.launch。它提供了更简洁的语法和更强的功能。3.1 torchrun的优势自动故障恢复worker失败时会自动重启弹性训练支持可以动态调整worker数量更简单的参数合并了一些冗余选项更好的错误处理提供更清晰的错误信息3.2 基本使用示例单机4卡训练torchrun --nproc_per_node4 train.py多机训练2节点每节点4卡# 主节点 torchrun --nproc_per_node4 --nnodes2 --node_rank0 --master_addr192.168.1.100 --master_port29500 train.py # 从节点 torchrun --nproc_per_node4 --nnodes2 --node_rank1 --master_addr192.168.1.100 --master_port29500 train.py3.3 弹性训练配置torchrun支持弹性训练允许worker数量在运行时变化。创建一个配置文件elastic_config.json{ min_size: 2, max_size: 8, nproc_per_node: 4 }然后启动torchrun --rdzv_confelastic_config.json train.py4. 在代码中获取分布式信息使用启动工具后我们的训练脚本可以简化为import torch.distributed as dist def main(): # 初始化分布式环境 dist.init_process_group(backendnccl) # 获取自动设置的参数 rank dist.get_rank() world_size dist.get_world_size() local_rank int(os.environ[LOCAL_RANK]) print(fRank {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_node4 train.py5.2 数据加载注意事项在分布式训练中数据加载需要特殊处理from torch.utils.data.distributed import DistributedSampler dataset MyDataset() sampler DistributedSampler(dataset, shuffleTrue) dataloader DataLoader(dataset, batch_size64, samplersampler)关键点每个epoch开始前调用sampler.set_epoch(epoch)保证shuffle正确性不要在自己的Dataset中实现shuffle5.3 日志记录策略在分布式环境中直接打印日志会导致混乱。推荐做法if dist.get_rank() 0: print(只有rank 0会打印这条消息)或者使用更专业的日志库import logging logging.basicConfig( levellogging.INFO if dist.get_rank() 0 else logging.WARN ) logger logging.getLogger(__name__)6. 性能优化技巧6.1 后端选择PyTorch支持多种分布式后端后端适用场景特点NCCL多GPU训练针对NVIDIA GPU优化性能最佳GlooCPU训练或多机训练稳定性好支持CPUMPIHPC环境需要系统支持MPI通常推荐使用NCCLdist.init_process_group(backendnccl)6.2 梯度压缩对于带宽受限的环境可以考虑梯度压缩model torch.nn.parallel.DistributedDataParallel( model, device_ids[local_rank], gradient_as_bucket_viewTrue # 启用梯度分桶 )6.3 通信重叠通过调整bucket_cap_mb参数优化通信model torch.nn.parallel.DistributedDataParallel( model, device_ids[local_rank], bucket_cap_mb25 # 调整桶大小 )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_node4 --use_env train.py迁移后torchrun --nproc_per_node4 train.py在实际项目中我发现torchrun的自动故障恢复功能特别有用尤其是在长时间训练任务中。曾经有一次训练因为节点重启而中断torchrun自动恢复了训练进度节省了大量时间。