【PyTorch】torchrun：分布式训练的启动和管理命令行工具

torchrun 是 PyTorch 提供的一个命令行工具，用于简化分布式训练的启动和管理。它在 PyTorch 1.10 中引入，作为 torch.distributed.launch 的升级替代品，提供了更简洁的接口、更好的弹性支持和容错能力，适用于单节点多 GPU 或多节点分布式训练。以下是关于 torchrun 的详细说明：

什么是 torchrun？

torchrun 是一个 Python 脚本执行工具，封装了分布式训练的初始化流程。它通过设置环境变量（如 RANK、LOCAL_RANK、WORLD_SIZE 等）自动管理进程组的启动和通信，消除了手动传递参数的复杂性。它支持 PyTorch 的分布式数据并行（DDP, DistributedDataParallel）和其他分布式训练模式，适用于 CPU 和 GPU 训练。

主要特点

• 简化使用：相比 torch.distributed.launch，torchrun 减少了命令行参数的复杂性，自动处理进程管理和排名。
• 弹性训练：支持动态调整节点或进程数量（弹性模式），能在进程失败时自动重启或重新分配资源。
• 容错性：如果某个进程崩溃，torchrun 可以检测并尝试恢复训练。
• 跨平台：支持单节点多 GPU、多节点训练，兼容 Linux、Windows 等系统。

基本命令格式

torchrun [参数] 你的训练脚本.py [脚本参数]

常用参数

1. --nproc_per_node

   • 每个节点上启动的进程数，通常设置为 GPU 数量（GPU 训练）或 CPU 核心数。
   • 示例：--nproc_per_node=4

2. --nnodes

   • 参与训练的节点（机器）总数，范围可以是固定值或 min:max 格式（弹性模式）。
   • 示例：--nnodes=1（单节点）或 --nnodes=2:4（弹性，2 到 4 个节点）

3. --node_rank

   • 当前节点的排名，从 0 到 nnodes-1。单节点时设为 0，多节点时每个节点需不同排名。
   • 示例：--node_rank=0

4. --rdzv_endpoint

   • 主节点的地址和端口，格式为 host:port，用于节点间通信协调。
   • 示例：--rdzv_endpoint=192.168.1.1:1234

5. --rdzv_backend

   • 分布式协调的后端，默认是 c10d，其他选项包括 etcd（用于弹性训练）。
   • 示例：--rdzv_backend=c10d

6. --rdzv_id

   • 任务的唯一标识符，用于区分多个分布式任务，特别是在共享集群上运行时。
   • 示例：--rdzv_id=123

7. --max_restarts (弹性模式)

   • 进程失败后最大重启次数，用于容错。
   • 示例：--max_restarts=3

使用示例

单节点多 GPU 训练

假设你有一个训练脚本 train.py，需要在单节点上的 4 个 GPU 上运行：

torchrun --nproc_per_node=4 --nnodes=1 --node_rank=0 --rdzv_endpoint=localhost:1234 train.py --batch_size=64 --epochs=10

• --nproc_per_node=4：启动 4 个进程，每个进程绑定到一个 GPU。
• --rdzv_endpoint=localhost:1234：指定主节点地址和端口。
• train.py --batch_size=64 --epochs=10：训练脚本及其参数。

多节点多 GPU 训练

假设有 2 个节点，每个节点有 4 个 GPU，节点 1 的 IP 为 192.168.1.1，节点 2 的 IP 为 192.168.1.2：

• 在节点 1（主节点）上运行：

  torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 --rdzv_endpoint=192.168.1.1:1234 train.py --batch_size=64 --epochs=10
  

• 在节点 2 上运行：

  torchrun --nproc_per_node=4 --nnodes=2 --node_rank=1 --rdzv_endpoint=192.168.1.1:1234 train.py --batch_size=64 --epochs=10
  

• 主节点（node_rank=0）协调通信，其他节点连接到 rdzv_endpoint。

弹性训练

支持动态节点数，例如允许 2 到 4 个节点参与：

torchrun --nproc_per_node=4 --nnodes=2:4 --node_rank=0 --rdzv_endpoint=localhost:1234 --max_restarts=3 train.py --batch_size=64

• --nnodes=2:4：节点数可在 2 到 4 之间动态调整。
• --max_restarts=3：允许失败后重启 3 次。

训练脚本适配

你的 train.py 需要支持分布式训练，关键步骤如下：

1. 初始化进程组：
   使用 torch.distributed.init_process_group 初始化分布式环境，torchrun 会自动设置环境变量。
2. 获取排名和设备：
   通过环境变量（如 LOCAL_RANK）确定进程的本地排名和绑定的 GPU。
3. 模型和数据：
   使用 torch.nn.parallel.DistributedDataParallel 包裹模型，搭配 DistributedSampler 分割数据。

示例 train.py

import torch
import torch.distributed as dist
import torch.nn as nn
import os

def main():
    # 初始化进程组，torchrun 自动设置环境变量
    dist.init_process_group(backend='nccl')  # NCCL 用于 GPU 训练

    # 获取排名和设备
    local_rank = int(os.environ['LOCAL_RANK'])  # 本地排名
    rank = int(os.environ['RANK'])  # 全局排名
    torch.cuda.set_device(local_rank)
    device = torch.device(f'cuda:{local_rank}')

    # 定义模型并移动到指定 GPU
    model = YourModel().to(device)
    model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

    # 数据加载器
    from torch.utils.data import DataLoader, DistributedSampler
    dataset = YourDataset()
    sampler = DistributedSampler(dataset)
    loader = DataLoader(dataset, batch_size=64, sampler=sampler)

    # 训练循环
    for epoch in range(10):
        sampler.set_epoch(epoch)  # 确保每个 epoch 数据洗牌不同
        for data, target in loader:
            data, target = data.to(device), target.to(device)
            # 训练步骤...
            if rank == 0:  # 仅主进程保存模型或打印日志
                print(f"Epoch {epoch}, Loss: ...")

    # 清理进程组
    dist.destroy_process_group()

if __name__ == '__main__':
    main()

环境变量

torchrun 自动设置以下关键环境变量，供脚本使用：

• RANK：进程的全局排名（0 到 world_size-1）。
• LOCAL_RANK：进程在当前节点内的本地排名（0 到 nproc_per_node-1）。
• WORLD_SIZE：总进程数（nnodes * nproc_per_node）。
• MASTER_ADDR 和 MASTER_PORT：从 rdzv_endpoint 解析的主节点地址和端口。

注意事项

• 后端：GPU 训练用 nccl 后端，CPU 训练用 gloo 后端。
• I/O 控制：保存模型、日志等操作应仅在主进程（RANK == 0）执行，避免冲突。
• 端口：确保 --rdzv_endpoint 的端口未被占用。
• 兼容性：检查 PyTorch 版本，torchrun 在 1.10 及以上版本可用。

与 torch.distributed.launch 的区别

• 命令简化：无需显式传递 --local_rank，torchrun 通过环境变量管理。
• 弹性支持：torchrun 支持动态节点数和容错，launch 不支持。
• 现代性：torch.distributed.launch 已弃用，torchrun 是官方推荐工具。

总结

torchrun 是 PyTorch 分布式训练的现代启动工具，简化了多 GPU 或多节点训练的配置，支持弹性训练和容错。使用时，只需提供节点数、进程数、通信端点等参数，搭配适配的训练脚本即可。