PyTorch DDP分布式训练部署实践

PyTorch DDP分布式训练部署实践

在现代机器学习项目中，PyTorch Distributed Data Parallel (DDP)已成为多机多卡训练的标准方案。本文将分享一个完整的部署实践案例。

环境准备

首先确保环境包含：

pip install torch torchvision torchaudio

核心配置代码

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)
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

# 训练函数
def train(rank, world_size):
    setup(rank, world_size)
    
    # 创建模型并移动到GPU
    model = MyModel().to(rank)
    
    # 包装模型
    ddp_model = DDP(model, device_ids=[rank])
    
    # 数据加载器配置
    train_loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=32,
        shuffle=True,
        num_workers=4,
        pin_memory=True
    )
    
    # 优化器配置
    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=0.001)
    
    # 训练循环
    for epoch in range(10):
        for batch in train_loader:
            optimizer.zero_grad()
            output = ddp_model(batch)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
    
    cleanup()

启动命令

python -m torch.multiprocessing --nproc_per_node=4 train.py

通过合理配置DDP参数，可显著提升训练效率，建议根据硬件配置调整batch size和学习率。