单机训练只有一个 Python 进程。DDP 的基本形态是每张 GPU 一个进程: 每个进程持有一份完整模型,只吃自己那份数据,反向传播后通过集合通信同步梯度。 所以学习分布式的第一件事,不是模型结构,而是认清"我这个进程是谁,它站在哪张卡上,它和谁通信"。

单机 4 卡的最小图

torchrun 的职责是把同一个 Python 脚本启动成多个进程, 并为每个进程设置好 RANK、LOCAL_RANK、 WORLD_SIZE、MASTER_ADDR、MASTER_PORT 等环境变量。 你的脚本只要读取这些变量,初始化进程组,就能进入分布式世界。

# 单机 4 卡启动。--standalone 会自动处理 rendezvous 地址和端口
torchrun --standalone --nproc_per_node=4 train.py

# 单机只用第 0、1 张卡
CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 train.py

# 双机示意。两台机器必须使用同一个 master 地址、端口和 nnodes
torchrun --nnodes=2 --node_rank=0 --nproc_per_node=8 \
  --master_addr=10.0.0.1 --master_port=29500 train.py

torchrun --nnodes=2 --node_rank=1 --nproc_per_node=8 \
  --master_addr=10.0.0.1 --master_port=29500 train.py

训练脚本骨架

import os
import torch
import torch.distributed as dist

def setup_distributed():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    return rank, local_rank, world_size, device

def cleanup():
    dist.destroy_process_group()

if __name__ == "__main__":
    rank, local_rank, world_size, device = setup_distributed()
    print(f"rank={rank} local_rank={local_rank} world_size={world_size} device={device}")
    cleanup()

分布式训练的很多高级技术,都能拆成少数几个 collective。 今天只需要掌握四个:AllReduce 用来同步梯度,AllGather 用来拼完整张量, ReduceScatter 用来"先聚合再切分",Broadcast 用来把 rank 0 的状态发给所有人。

collectives_lab.py

import os
import torch
import torch.distributed as dist

dist.init_process_group("nccl")
rank = dist.get_rank()
world_size = dist.get_world_size()
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)

# 1. AllReduce:所有 rank 最终都拿到 0+1+...+(world_size-1)
x = torch.tensor([float(rank)], device=device)
dist.all_reduce(x, op=dist.ReduceOp.SUM)
print(rank, "all_reduce", x.item())

# 2. AllGather:每个 rank 收集所有 rank 的值
y = torch.tensor([rank], device=device)
gathered = [torch.zeros_like(y) for _ in range(world_size)]
dist.all_gather(gathered, y)
print(rank, "all_gather", [t.item() for t in gathered])

# 3. Broadcast:rank 0 的 2026 发给所有人
z = torch.tensor([2026 if rank == 0 else -1], device=device)
dist.broadcast(z, src=0)
print(rank, "broadcast", z.item())

# 4. ReduceScatter:每个 rank 得到 reduce 后的一块
inp = torch.arange(world_size * 2, device=device, dtype=torch.float32) + rank * 100
out = torch.empty(2, device=device)
dist.reduce_scatter_tensor(out, inp, op=dist.ReduceOp.SUM)
print(rank, "reduce_scatter", out.tolist())

dist.destroy_process_group()

# 运行
torchrun --standalone --nproc_per_node=4 collectives_lab.py

DDP 的关键改动只有四处:初始化进程组、绑定当前 GPU、用 DistributedSampler 切数据、用 DistributedDataParallel 包住模型。之后你仍然写普通的 forward、loss、backward、optimizer step。

import os
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from torchvision import datasets, transforms

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        return self.fc2(x)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--epochs", type=int, default=2)
    parser.add_argument("--batch-size", type=int, default=128)
    args = parser.parse_args()

    dist.init_process_group("nccl")
    rank = dist.get_rank()
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = dist.get_world_size()
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    dataset = datasets.MNIST("./data", train=True, download=True, transform=transform)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        sampler=sampler,
        num_workers=2,
        pin_memory=True
    )

    model = Net().to(device)
    model = DDP(model, device_ids=[local_rank], output_device=local_rank)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

    for epoch in range(args.epochs):
        sampler.set_epoch(epoch)
        model.train()
        total_loss = 0.0
        for step, (data, target) in enumerate(loader):
            data = data.to(device, non_blocking=True)
            target = target.to(device, non_blocking=True)
            optimizer.zero_grad(set_to_none=True)
            logits = model(data)
            loss = F.cross_entropy(logits, target)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        avg = torch.tensor([total_loss / len(loader)], device=device)
        dist.all_reduce(avg, op=dist.ReduceOp.SUM)
        avg = avg.item() / world_size
        if rank == 0:
            print(f"epoch={epoch} avg_loss={avg:.4f}")

    if rank == 0:
        torch.save(model.module.state_dict(), "mnist_ddp.pt")
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

# 运行。没有 torchvision 时先安装对应 PyTorch 版本的 torchvision
torchrun --standalone --nproc_per_node=2 train_ddp_mnist.py --epochs 2

# 常用调试环境变量
NCCL_DEBUG=INFO TORCH_DISTRIBUTED_DEBUG=DETAIL \
torchrun --standalone --nproc_per_node=2 train_ddp_mnist.py

1. 画出单机 4 卡和双机每机 4 卡时,global rank、local rank、node rank、world size 的对应关系。
2. 说明 torchrun --standalone --nproc_per_node=4 train.py 会启动几个进程,每个进程能读到哪些关键环境变量。
3. 用一个小张量例子解释 AllReduce、AllGather、ReduceScatter、Broadcast 四个 collective 的输入输出。
4. DDP 训练里,为什么 forward 不通信,backward 后却会触发梯度通信?
5. 为什么 DDP 训练集需要 DistributedSampler,验证集和测试集在什么情况下也需要分布式采样?
6. 列出 5 个 DDP 卡住或报错时的排查点,并说明你会先看哪一个日志或环境变量。