4、多机多卡同步训练实战：多机通信原理、torchrun 启动脚本、环境变量设置、多机训练代码实战

好，咱们终于聊到多机多卡了。

说实话，单机多卡训练，你只要搞定一台机器上的 GPU 通信就行。但多机训练，说白了就是让几台机器像一支军队一样协同作战。我刚开始搞多机训练时，踩过的坑比代码行数还多。今天我把这些经验掰开揉碎了讲给你。

4.1 多机通信原理：机器之间怎么“说话”

多机训练的核心，是让每台机器上的 GPU 能交换梯度。你想想看，如果每台机器只算自己那部分数据，最后不合并梯度，那模型就学歪了。

PyTorch 用的通信后端是 NCCL（NVIDIA Collective Communications Library）。它专门为 GPU 设计，走的是 InfiniBand 或者 RoCE 网络。我个人习惯用 InfiniBand，延迟低、带宽大，但贵。如果你只有千兆以太网，嗯，那训练速度会慢得让你怀疑人生。

通信模式是 AllReduce。每台机器算完梯度后，NCCL 会把所有机器的梯度加起来，再广播回去。这个过程叫“规约-广播”。

这里我画了一张图，帮你理解多机通信的流程：

4.2 torchrun 启动脚本：一键拉起多机训练

PyTorch 官方推荐用 torchrun 来启动分布式训练。它帮你处理了进程管理、环境变量设置、错误重启等脏活累活。

先看一个典型的多机启动命令：

# 在机器 1 上执行（主节点）
torchrun --nnodes=2 --nproc_per_node=4 \
  --rdzv_endpoint=192.168.1.100:29500 \
  --rdzv_backend=c10d \
  train.py

# 在机器 2 上执行（从节点）
torchrun --nnodes=2 --nproc_per_node=4 \
  --rdzv_endpoint=192.168.1.100:29500 \
  --rdzv_backend=c10d \
  train.py

参数说明：

参数	含义	我的建议
--nnodes	总机器数	写死，不要用变量
--nproc_per_node	每台机器的 GPU 数	一般等于 GPU 数量
--rdzv_endpoint	主节点的 IP:端口	用内网 IP，别用 127.0.0.1
--rdzv_backend	发现后端	c10d 就够了

4.3 环境变量设置：MASTER_ADDR, MASTER_PORT, WORLD_SIZE, RANK

torchrun 会自动帮你设置这些环境变量，但理解它们很重要。我见过太多人因为不理解这些变量而出错。

环境变量	含义	示例值
MASTER_ADDR	主节点 IP	192.168.1.100
MASTER_PORT	主节点端口	29500
WORLD_SIZE	总进程数 = nnodes × nproc_per_node	8
RANK	当前进程的全局编号	0 ~ 7
LOCAL_RANK	当前进程在本地机器上的编号	0 ~ 3

举个例子：2 台机器，每台 4 卡，那么：

• 机器 1 的进程 RANK 是 0,1,2,3，LOCAL_RANK 是 0,1,2,3
• 机器 2 的进程 RANK 是 4,5,6,7，LOCAL_RANK 是 0,1,2,3

4.4 多机训练代码实战

好，代码来了。这是一个完整的多机训练脚本，我加了详细的注释。

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler

# 1. 初始化进程组
def init_distributed():
    dist.init_process_group(backend='nccl')
    torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())

# 2. 模拟数据集
class FakeDataset(Dataset):
    def __init__(self, size=1000):
        self.data = torch.randn(size, 3, 224, 224)
        self.labels = torch.randint(0, 10, (size,))
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# 3. 简单模型
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(3*224*224, 10)
    
    def forward(self, x):
        return self.fc(x.view(x.size(0), -1))

# 4. 训练函数
def train():
    init_distributed()
    
    rank = dist.get_rank()
    local_rank = rank % torch.cuda.device_count()
    
    # 每个进程绑定不同的 GPU
    device = torch.device(f'cuda:{local_rank}')
    
    # 模型放到 GPU 上
    model = SimpleModel().to(device)
    # 包装成分布式模型
    model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])
    
    # 数据集：每个进程只看到一部分
    dataset = FakeDataset()
    sampler = DistributedSampler(dataset)
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)
    
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    criterion = nn.CrossEntropyLoss()
    
    for epoch in range(5):
        sampler.set_epoch(epoch)  # 每个 epoch 打乱数据
        for data, target in dataloader:
            data, target = data.to(device), target.to(device)
            
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
        
        # 只在主进程打印
        if rank == 0:
            print(f'Epoch {epoch}, Loss: {loss.item():.4f}')
    
    dist.destroy_process_group()

if __name__ == '__main__':
    train()

启动方式：

# 机器 1
torchrun --nnodes=2 --nproc_per_node=4 --rdzv_endpoint=192.168.1.100:29500 train.py

# 机器 2
torchrun --nnodes=2 --nproc_per_node=4 --rdzv_endpoint=192.168.1.100:29500 train.py

最后说一句，多机训练调试起来比单机麻烦得多。我的习惯是先在一台机器上用 --nnodes=1 跑通，确认代码没问题，再扩展到多机。这样能省下大量排查网络问题的时间。

好了，多机多卡同步训练的核心内容就这些。记住：通信是瓶颈，环境变量是基础，torchrun 是你的好帮手。动手试试吧！