目录PyTorch torchrun 命令详解1. 基本作用2. torchrun 做了什么3. 常用命令参数3.1 --nproc-per-node3.2 --nnodes3.3 --node-rank3.4 --master-addr3.5 --master-port3.6 --standalone4. torchrun 自动设置的环境变量5. train.py 如何使用这些变量6. DistributedDataParallelDDP7. 数据如何划分8. 单机训练示例9. 多机训练示例Machine AMachine B10. torchrun 工作流程11. 常见参数速查12. 总结PyTorch torchrun 命令详解torchrun是 PyTorch 官方提供的分布式训练启动器Distributed Launcher从 PyTorch 1.9 开始用于替代torch.distributed.launch。它的主要作用是启动多个训练进程配置分布式训练环境为每个进程分配 GPU建立进程间通信NCCL/Gloo配合 DDPDistributedDataParallel完成多 GPU 或多机训练1. 基本作用单 GPU 训练python train.py如果有 8 张 GPU希望每张 GPU 启动一个训练进程torchrun\--nproc-per-node8\train.py实际上会启动 8 个 Python 进程GPU0 - train.py (rank0) GPU1 - train.py (rank1) GPU2 - train.py (rank2) ... GPU7 - train.py (rank7)一个 GPU 对应一个 Python 进程而不是一个线程。2. torchrun 做了什么例如执行torchrun\--nproc-per-node4\train.py实际过程如下torchrun │ ├── python train.py (LOCAL_RANK0) ├── python train.py (LOCAL_RANK1) ├── python train.py (LOCAL_RANK2) └── python train.py (LOCAL_RANK3)虽然运行的是同一个train.py但每个进程拥有不同的环境变量。例如第一个进程LOCAL_RANK0 RANK0 WORLD_SIZE4第二个进程LOCAL_RANK1 RANK1 WORLD_SIZE4第三个进程LOCAL_RANK2 RANK2 WORLD_SIZE4因此训练脚本无需知道是谁启动它只需要读取环境变量即可。3. 常用命令参数3.1 --nproc-per-node每台机器启动多少个训练进程。一般情况下训练进程数 GPU 数例如--nproc-per-node8表示GPU0 GPU1 GPU2 ... GPU7每张 GPU 启动一个 Python 进程。3.2 --nnodes参与训练的机器数量。例如--nnodes2表示Node0 Node1两台服务器共同训练。3.3 --node-rank当前机器编号。例如第一台--node-rank0第二台--node-rank13.4 --master-addr主节点 IP 地址。例如--master-addr10.0.0.1所有节点都会连接到该地址。3.5 --master-port主节点通信端口。例如--master-port29500所有训练进程通过该端口建立通信。3.6 --standalone单机训练模式。例如torchrun\--standalone\--nproc-per-node8\train.py等价于自动设置MASTER_ADDRlocalhost MASTER_PORT随机端口无需手动指定主节点地址。4. torchrun 自动设置的环境变量启动后torchrun会自动设置以下环境变量MASTER_ADDR MASTER_PORT WORLD_SIZE RANK LOCAL_RANK LOCAL_WORLD_SIZE例如WORLD_SIZE16表示整个训练共有16个进程例如LOCAL_WORLD_SIZE8表示当前机器有8个训练进程5. train.py 如何使用这些变量读取当前 GPUimportos local_rankint(os.environ[LOCAL_RANK])绑定 GPUtorch.cuda.set_device(local_rank)例如LOCAL_RANK0对应cuda:0初始化分布式通信importtorch.distributedasdist dist.init_process_group(backendnccl)无需手动指定rankworld_size因为torchrun已经通过环境变量自动配置。6. DistributedDataParallelDDP典型代码modelmodel.cuda(local_rank)modeltorch.nn.parallel.DistributedDataParallel(model,device_ids[local_rank])训练过程Forward ↓ Backward ↓ NCCL AllReduce ↓ 梯度同步 ↓ 参数更新梯度同步由 DDP 自动完成。7. 数据如何划分假设Dataset 10000 images使用DistributedSampler(dataset)4 个 GPU 自动划分为GPU0 0 ~ 2499 GPU1 2500 ~ 4999 GPU2 5000 ~ 7499 GPU3 7500 ~ 9999每个 GPU 只读取自己的数据不会重复。8. 单机训练示例例如训练 LeRobottorchrun\--standalone\--nnodes1\--nproc-per-node8\train.py\--configconfigs/pi0.yaml流程如下torchrun │ ├── GPU0 → train.py ├── GPU1 → train.py ├── GPU2 → train.py ├── ... └── GPU7 → train.py9. 多机训练示例Machine Atorchrun\--nnodes2\--node-rank0\--master-addr10.0.0.1\--master-port29500\--nproc-per-node8\train.pyMachine Btorchrun\--nnodes2\--node-rank1\--master-addr10.0.0.1\--master-port29500\--nproc-per-node8\train.py最终Machine0 GPU0 GPU1 ... GPU7 Machine1 GPU0 GPU1 ... GPU7 ↓ 总共16个训练进程10. torchrun 工作流程torchrun │ ├── 解析命令行参数 │ ├── 设置 MASTER_ADDR、MASTER_PORT │ ├── 设置环境变量 │ ├── RANK │ ├── LOCAL_RANK │ ├── WORLD_SIZE │ └── LOCAL_WORLD_SIZE │ ├── 启动多个 Python 训练进程 │ ├── train.py 调用 init_process_group() │ ├── 建立 NCCL/Gloo 通信组 │ ├── 每个进程绑定一个 GPU │ ├── DistributedSampler 划分数据 │ └── DDP 在反向传播时自动同步梯度11. 常见参数速查参数含义示例--nproc-per-node每台机器启动的训练进程数通常等于 GPU 数8--nnodes参与训练的机器总数2--node-rank当前机器编号从0开始0、1--master-addr主节点 IP 或主机名10.0.0.1--master-port主节点通信端口29500--standalone单机模式自动配置主节点地址和端口常用于单机训练12. 总结torchrun的职责可以概括为三点启动多个训练进程通常每个 GPU 一个进程。配置分布式运行环境设置RANK、WORLD_SIZE、MASTER_ADDR等环境变量。配合torch.distributed和DistributedDataParallel (DDP)实现高效的多 GPU 或多机分布式训练。一句话总结torchrun本身并不负责训练模型而是负责启动训练进程、建立分布式通信环境并将不同的 GPU 和进程组织起来使 DDP 能够正常工作。