深度学习分布式训练详解：PyTorch DDP 实战

深度学习分布式训练详解PyTorch DDP 实战1. 背景与动机随着模型规模的增长单机训练已无法满足需求。分布式训练通过多机多卡并行计算大幅缩短训练时间是训练大模型的必备技术。2. 分布式训练基础2.1 数据并行 vs 模型并行数据并行每个 GPU 处理不同数据梯度同步模型并行模型分片到不同 GPU数据顺序处理2.2 PyTorch DDPimport torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def setup(rank, world_size): dist.init_process_group(nccl, rankrank, world_sizeworld_size) def cleanup(): dist.destroy_process_group()3. DDP 实战def train(rank, world_size): setup(rank, world_size) # 创建模型并移动到指定 GPU model MyModel().to(rank) model DDP(model, device_ids[rank]) # 数据加载器 train_sampler torch.utils.data.distributed.DistributedSampler( train_dataset, num_replicasworld_size, rankrank ) train_loader DataLoader( train_dataset, batch_size32, samplertrain_sampler ) optimizer torch.optim.Adam(model.parameters()) for epoch in range(num_epochs): train_sampler.set_epoch(epoch) for data, target in train_loader: data, target data.to(rank), target.to(rank) optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() optimizer.step() cleanup() # 启动多进程 import torch.multiprocessing as mp if __name__ __main__: world_size 4 mp.spawn(train, args(world_size,), nprocsworld_size, joinTrue)4. 混合精度分布式训练from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for data, target in train_loader: data, target data.to(rank), target.to(rank) with autocast(): output model(data) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()5. 多机训练# Node 0 python -m torch.distributed.launch \ --nproc_per_node8 \ --nnodes2 \ --node_rank0 \ --master_addr192.168.1.1 \ --master_port12345 \ train.py # Node 1 python -m torch.distributed.launch \ --nproc_per_node8 \ --nnodes2 \ --node_rank1 \ --master_addr192.168.1.1 \ --master_port12345 \ train.py6. 性能优化优化技术加速比说明DDP3.5x (4卡)数据并行Mixed Precision2-3x混合精度训练Gradient Accumulation-模拟大批量Gradient Checkpointing30%节省显存7. 结论分布式训练是训练大模型的关键技术。通过 DDP、混合精度和适当的优化策略可以充分利用多 GPU 资源显著缩短训练时间。