知识花园

torchrun

6分钟阅读

torchrun是PyTorch分布式训练命令行工具，支持单机多卡和多机多卡训练。

基本用法

torchrun [options] your_script.py [args]

常用参数

参数说明
--nnodes=N总节点数（默认1）
--nproc_per_node=G每节点进程数（通常等于GPU数量）
--node_rank=R当前节点排名（从0开始）
--master_addr=HOST主节点地址（默认127.0.0.1）
--master_port=PORT主节点端口（默认29500）
--standalone单机模式，简化配置

示例

单机多卡（4 GPU）：

torchrun --standalone --nproc_per_node=4 your_script.py --args

多机多卡（2节点各4 GPU）：

节点0：

torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 \
  --master_addr="node0_ip" --master_port=12345 your_script.py

节点1：

torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 \
  --master_addr="node0_ip" --master_port=12345 your_script.py

脚本配置

import torch
import torch.distributed as dist
 
def main():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    # 创建分布式数据加载器和模型...
 
if __name__ == "__main__":
    main()

注意事项

  • 使用--standalone时无需设置master_addr和node_rank
  • 确保所有节点间网络互通
  • 脚本需兼容分布式训练（使用DistributedDataLoader）

torchrun 完整说明

1. torchrun 是否会随PyTorch自动安装？

✅ 是的，安装pytorch之后自带，不需要pip额外装。

  • torchrun 本质是一个命令行入口脚本，内部等价于：
torchrun train.py
# 和下面这条完全等价
python -m torch.distributed.run train.py
  • 它属于 torch.distributed.run 模块，是PyTorch源码内置，安装完pytorch就可以直接在终端调用。

⚠️ 注意：conda/pip装好pytorch后，如果终端敲torchrun报command not found，说明你的shell环境没有激活对应的python虚拟环境，找不到该脚本。

历史：旧版本废弃的 python -m torch.distributed.launch，现在全部替换为torchrun。

2. torchrun 的启动原理

核心一句话

torchrun 是进程管理器，本身不做模型训练，不做张量通信；它负责批量拉起N个子进程，给每个子进程注入分布式环境变量，之后你的训练脚本内部调用dist.init_process_group()完成真正分布式通信初始化。

完整执行流程（单机多卡示例 torchrun --nproc_per_node=2 train.py）

  1. 父进程（torchrun主进程）启动 读取参数：--nproc_per_node=2，代表本机启动2个worker子进程。单机模式会自动生成MASTER_ADDR=127.0.0.1、自动分配MASTER_PORT。

  2. fork/spawn 生成N份完全相同的训练脚本子进程 两份子进程都会运行你写的train.py，两份跑一模一样的代码，但是每个子进程的环境变量被父进程修改注入：

环境变量进程0进程1
RANK（全局rank）01
LOCAL_RANK（本机内编号）01
WORLD_SIZE（总进程数）22
MASTER_ADDR / MASTER_PORT相同相同

这就是为什么你脚本里写 os.environ["LOCAL_RANK"] 可以读到值，不是dist库生成的，是torchrun父进程注入的环境变量。

  1. 每个子进程内部执行你的代码 每个worker进程走到：
dist.init_process_group(backend="nccl", init_method="env://")

init_process_group读取当前进程的环境变量（RANK/WORLD_SIZE/MASTER_*），完成进程组初始化，建立NCCL通信通道，之后dist.all_reduce、DDP才可以工作。

init_method="env://"含义：从环境变量读取分布式配置，这就是和torchrun配套的初始化方式。

  1. 业务训练逻辑运行；子进程退出后，父进程负责收集子进程返回码，做错误处理、弹性重启（torchrun独有能力）。

多机场景

多台机器都执行torchrun，通过rdzv（rendezvous，会合机制），所有节点上的worker通过主节点的地址端口相互发现，分配全局RANK，拉起全部进程。

torchrun对比旧launch的关键升级

  1. 默认使用环境变量传参，不再需要--local_rank命令行参数；旧launch需要传命令行参数，还需要加--use_env开关。
  2. 原生支持弹性训练：进程挂掉可以自动重启（--max_restarts），支持动态增减worker数量。旧launch没有容错能力，一个进程挂全部直接崩溃。
  3. 内置rendezvous会合服务，多机作业启动更简单。

3. 重要易错点

  1. ❗torchrun只负责拉起进程+注入环境变量，不会自动调用init_process_group！ 你代码必须手动写dist.init_process_group()，否则dist所有API全部报错。

  2. 每一个GPU对应一个独立Python进程，不是多线程。各个进程之间内存完全隔离，只能靠dist通信原语交换张量。

  3. 单机示例最小启动命令

# 单机2卡
torchrun --nproc_per_node=2 train.py
# 等价写法
python -m torch.distributed.run --nproc_per_node=2 train.py

4. 流程图极简版

终端执行 torchrun → torchrun父进程
    ↓
生成N个子进程，每个子进程注入 RANK LOCAL_RANK WORLD_SIZE MASTER_*环境变量
    ↓
每个子进程运行train.py
    ↓
脚本中 dist.init_process_group(init_method="env://") 读取环境变量初始化进程组
    ↓
dist / DDP 开始分布式训练

关系图谱