下面将对 torchrun 的使用做一个较为详细的介绍，包括命令行参数如何指定，以及在脚本里如何获取分布式相关信息并进行初始化。

torch.distributed.launch 启动器

对于单机多卡的情况：

1
2
CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py
# nproc_per_node: 这个参数是指你使用这台服务器上面的几张显卡

torchrun

torchrun 的基本用法

PyTorch 从 1.10 版本开始推荐使用 torchrun 来启动分布式训练，取代之前的 torch.distributed.launch 脚本。其优点是：

  1. 默认使用 “env://” 方式传递分布式所需的信息（MASTER_ADDR, MASTER_PORT, RANK, LOCAL_RANK 等），不再强制注入 --local_rank=... 命令行参数。
  2. 提供了更简洁、清晰的命令行参数，如 --nproc_per_node, --nnodes, --node_rank 等。
  3. 与后续 PyTorch 版本更兼容，且官方文档更推荐。

命令行示例

1
2
3
4
5
6
# 单机多卡示例
torchrun --standalone \
         --nproc_per_node=4 \
         your_script.py \
         --other_arg1=xxx \
         --other_arg2=yyy
  • --standalone：表示单机模式下自动分配一个可用端口给分布式进程组；若多机训练，需要显式指定 --master_addr 和 --master_port。
  • --nproc_per_node=4：表示在当前节点（机器）上启动 4 个进程（一般对应 4 张 GPU）。
  • your_script.py：你的训练脚本，可以在脚本中通过环境变量获取 LOCAL_RANK、RANK、WORLD_SIZE 等信息。
  • --other_arg1=xxx、--other_arg2=yyy：你自定义的脚本参数。

分布式脚本的基本结构

下面以一个最简化的示例脚本 train.py 为例，说明如何在程序中获取分布式信息并初始化。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
# train.py
import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim

def main():
    # 1. 获取环境变量
    local_rank = int(os.environ["LOCAL_RANK"])   # 当前进程在本机的进程号，从 0 开始
    rank = int(os.environ["RANK"])               # 当前进程在所有进程中的全局 rank
    world_size = int(os.environ["WORLD_SIZE"])   # 总进程数（所有机器、所有卡的进程总和）
    
    # 2. 初始化进程组
    #    backend 通常在 GPU 上用 'nccl'，若 CPU 上可以用 'gloo'。
    dist.init_process_group(backend="nccl")
    
    # 3. 设置当前进程要使用的 GPU
    torch.cuda.set_device(local_rank)
    
    # 4. 构建模型并移动到本地进程对应的 GPU
    model = nn.Linear(10, 1).cuda(local_rank)
    
    # 5. 构建分布式数据并行(DDP)模型 (可选)
    #    如果想让多个进程并行训练同一个模型，需要使用 DDP。
    from torch.nn.parallel import DistributedDataParallel as DDP
    model = DDP(model, device_ids=[local_rank], output_device=local_rank)
    
    # 6. 数据加载、优化器设置等
    optimizer = optim.SGD(model.parameters(), lr=0.001)
    # dataset = ...
    # sampler = ...
    # dataloader = DataLoader(dataset, batch_size=..., sampler=sampler)
    
    # 7. 训练循环 (示意)
    # for epoch in range(num_epochs):
    #     for batch in dataloader:
    #         optimizer.zero_grad()
    #         outputs = model(batch)
    #         loss = ...
    #         loss.backward()
    #         optimizer.step()
    
    # 8. 训练完成后（可选）销毁进程组
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

关键点说明

  1. 获取环境变量
    • LOCAL_RANK：在单机多卡时，表示当前进程对应第几张 GPU（0 ~ nproc_per_node-1）。
    • RANK：在所有进程（包括多机）中的全局排名。
    • WORLD_SIZE：所有进程的总数。
    • 若是多机训练，还会有 MASTER_ADDR、MASTER_PORT 等环境变量，由 torchrun 自动设置。
  2. 初始化进程组
    • 使用 dist.init_process_group(backend='nccl')（GPU 上一般选 nccl）。
    • 不需要再手动指定 init_method，torchrun 会默认用 “env://” 并自动设置主节点地址和端口。
  3. 设置当前进程使用的 GPU
    • torch.cuda.set_device(local_rank) 可以让本进程只操作对应的 GPU。
  4. 分布式数据并行 (DDP)
    • 如果你想真正并行训练，需要用 DistributedDataParallel 包装模型。
    • device_ids=[local_rank]、output_device=local_rank 可以确保梯度在本卡上做同步。
  5. 训练结束
    • 建议调用 dist.destroy_process_group() 结束分布式环境。

多机多卡的命令行示例

如果有多台机器（节点），可以用以下参数指定：

1
2
3
4
5
6
torchrun --nnodes=2 \               # 总共 2 台机器
         --nproc_per_node=4 \       # 每台机器上 4 张 GPU
         --node_rank=0 \            # 当前机器是第 0 台（另一台是 1）
         --master_addr="10.0.0.1" \ # 主节点 IP
         --master_port=12345 \      # 通信端口
         your_script.py

在另一台机器上也启动相同命令，但把 --node_rank=1 改对即可。

与命令行参数的混合使用

除了分布式必需的环境变量外，你的脚本还可以有自己的命令行参数。例如：

1
2
3
4
5
torchrun --standalone \
         --nproc_per_node=4 \
         your_script.py \
         --config ./myconfig.yaml \
         --epochs 100

脚本里可以用 argparse 解析自定义参数，然后再用 os.environ 获取分布式信息：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# train.py
import argparse
import os
import torch.distributed as dist

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, default=None)
    parser.add_argument("--epochs", type=int, default=10)
    return parser.parse_args()

def main():
    args = parse_args()
    
    local_rank = int(os.environ["LOCAL_RANK"])
    dist.init_process_group(backend="nccl")
    # ... 后面跟分布式训练逻辑
    
if __name__ == "__main__":
    main()

常见问题与注意事项

  1. “Unrecognized arguments: –local-rank=xxx”
    • 在 torchrun 下不会自动注入 --local_rank 命令行参数；该问题常见于使用旧版 torch.distributed.launch。
    • torchrun 会改用环境变量的形式传递本地进程号，不会出现这个冲突。
  2. 显式指定端口
    • 如果不想用 --standalone，可以用 --master_addr 和 --master_port 来指定 IP 和端口。
  3. 多机训练网络通信
    • 所有节点都要能通过 MASTER_ADDR 和 MASTER_PORT 通信。
    • 确保端口没被占用、防火墙没拦截。
  4. 梯度同步
    • DDP 训练时，每一步都会在进程间做梯度同步，所以总的 batch size = 单卡 batch size × 卡数。
    • 注意调整学习率，以适应更大的 batch size。
  5. 性能调优
    • 可能需要设置 OMP_NUM_THREADS、torch.backends.cudnn.benchmark = True 等来获得更好性能。

小结

  1. 编写脚本：
    • 在脚本里用 os.environ 读取 LOCAL_RANK、RANK、WORLD_SIZE；
    • 调用 dist.init_process_group(backend='nccl')；
    • 设置 GPU 设备并（可选）用 DDP 封装模型；
    • 进行正常的训练逻辑。
  2. 使用 torchrun 启动：
    • 单机多卡：torchrun --standalone --nproc_per_node=4 your_script.py；
    • 多机多卡：指定 --nnodes, --node_rank, --master_addr, --master_port 等；
    • 自定义的脚本参数直接在 your_script.py 后面追加。
  3. 优势：
    • 不需要在脚本里处理 --local_rank 命令行；
    • 环境变量方式更统一、更简洁；
    • 未来版本 PyTorch 官方会移除老的 torch.distributed.launch，因此推荐迁移到 torchrun。

通过以上步骤，你就可以在本地或多机环境中顺利使用 torchrun 进行分布式训练。祝训练顺利！