PyTorch DDP训练环境配置指南

PyTorch Distributed Data Parallel (DDP) 是实现多机多卡训练的核心组件。本文将详细介绍完整的环境配置流程，帮助您构建高效的分布式训练环境。

环境准备

首先确保所有节点具备相同的硬件环境和软件依赖：

# 安装PyTorch（建议使用最新版本）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install torchmetrics

基础配置步骤

1. 网络配置

在所有训练节点上设置环境变量：

export MASTER_ADDR=192.168.1.100
export MASTER_PORT=12345
export WORLD_SIZE=8
export RANK=0

2. 启动脚本示例

创建训练启动脚本 run_distributed.py：

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

# 初始化分布式环境
def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

# 训练函数
def train(rank, world_size):
    setup(rank, world_size)
    
    # 创建模型并移动到GPU
    model = torch.nn.Linear(1000, 10).to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    
    # 定义损失函数和优化器
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.01)
    
    # 训练循环
    for epoch in range(10):
        # 假设数据已分片
        outputs = ddp_model(torch.randn(64, 1000).to(rank))
        loss = criterion(outputs, torch.randint(0, 10, (64,)).to(rank))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    dist.destroy_process_group()

if __name__ == "__main__":
    world_size = 8
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)

3. 启动命令

使用以下命令启动训练：

python run_distributed.py

性能优化建议

• 使用 torch.cuda.amp 进行混合精度训练
• 调整 batch size 以匹配硬件资源
• 确保网络带宽足够支持数据同步