pytorch ddp 多机多卡训练示例

下面是一个使用PyTorch DDP（分布式数据并行）进行多机多卡训练的示例：

1.首先，需要使用torch.distributed.launch启动多个进程，每个进程都运行相同的脚本并使用不同的参数。例如，在两台机器上运行以下命令：

# Machine 1
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr="10.0.0.1" --master_port=8888 train.py

# Machine 2
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=1 --master_addr="10.0.0.1" --master_port=8888 train.py

上面的命令将在两台机器上启动4个进程，每个进程使用2个GPU进行训练。

2.在代码中，使用torch.distributed初始化进程组，并将模型和数据加载到每个进程中。例如：

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

# Initialize distributed process group
dist.init_process_group(backend='nccl', init_method='env://')

# Load data and model
train_data = ...
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)

model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

# Distributed model and optimizer
model = nn.parallel.DistributedDataParallel(model)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

这里使用了nn.parallel.DistributedDataParallel将模型包装成分布式模型，使用torch.optim.SGD作为优化器。

3.在训练循环中，每个进程都会收集自己的梯度并将它们聚合到进程组中。然后，所有进程都将使用平均梯度更新模型参数。例如：

for epoch in range(10):
    for data, target in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = nn.functional.cross_entropy(output, target)
        loss.backward()

        # All-reduce gradients
        for param in model.parameters():
            dist.all_reduce(param.grad.data, op=dist.reduce_op.SUM)

        optimizer.step()

在每个批次之后，使用dist.all_reduce将每个进程的梯度聚合到进程组中，然后使用平均梯度更新模型参数。

4.训练完成后，使用dist.destroy_process_group()关闭进程组并释放资源。例如：

dist.destroy_process_group()

这个示例展示了如何使用PyTorch DDP进行多机多卡训练。需要注意的是，使用DDP需要确保所有进程都能够访问相同的数据和模型，并且需要正确设置进程组中的参数。