torchrun 实现多卡GPU训练，怎么使用

torchrun 是 PyTorch 提供的一个命令行工具，用于在多节点、多 GPU 环境下启动分布式训练任务。它是 torch.distributed.launch 的替代方案，提供了更简洁、更灵活的启动方式，能够自动处理进程组的初始化和管理。

torchrun --nproc_per_node=5 LoraQwenMulKa.py

主要功能

• 自动处理进程组初始化：torchrun 会自动处理分布式训练中进程组的初始化，你无需在代码中手动设置复杂的环境变量。
• 支持多节点和多 GPU：可以方便地在多个节点（多台机器）和多个 GPU 上启动训练任务。
• 故障恢复：当某个进程失败时，torchrun 可以自动重启整个训练任务。

使用示例

下面将分别介绍单节点多 GPU 和多节点多 GPU 的使用示例。

单节点多 GPU 示例

假设你有一个简单的 PyTorch 训练脚本 train.py，代码如下：

import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)


def cleanup():
    dist.destroy_process_group()


class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def demo_basic(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)

    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10).to(rank))
    labels = torch.randn(20, 5).to(rank)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup()


def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             join=True)


if __name__ == "__main__":
    n_gpus = torch.cuda.device_count()
    assert n_gpus >= 2, f"Requires at least 2 GPUs to run, but got {n_gpus}"
    world_size = n_gpus
    run_demo(demo_basic, world_size)

要在单节点的多个 GPU 上运行这个脚本，可以使用以下 torchrun 命令：

torchrun --nproc_per_node=2 train.py

这里的 --nproc_per_node=2 表示在每个节点上启动 2 个进程，每个进程对应一个 GPU。

多节点多 GPU 示例

假设你有两个节点（机器），节点 0 的 IP 地址是 192.168.1.100，节点 1 的 IP 地址是 192.168.1.101，每个节点有 2 个 GPU。

在节点 0 上运行以下命令：

torchrun --nnodes=2 --nproc_per_node=2 --node_rank=0 --master_addr="192.168.1.100" --master_port=12355 train.py

在节点 1 上运行以下命令：

torchrun --nnodes=2 --nproc_per_node=2 --node_rank=1 --master_addr="192.168.1.100" --master_port=12355 train.py

• --nnodes：表示总的节点数。
• --nproc_per_node：表示每个节点上的进程数，通常等于该节点上的 GPU 数。
• --node_rank：表示当前节点的编号，从 0 开始。
• --master_addr：表示主节点的 IP 地址。
• --master_port：表示主节点的端口号。

通过以上命令，你可以在多节点多 GPU 的环境下启动分布式训练任务。