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

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

1. 背景与动机

随着模型规模的增长，单机训练已无法满足需求。分布式训练通过多机多卡并行计算，大幅缩短训练时间，是训练大模型的必备技术。

2. 分布式训练基础

2.1 数据并行 vs 模型并行

• 数据并行：每个 GPU 处理不同数据，梯度同步
• 模型并行：模型分片到不同 GPU，数据顺序处理

2.2 PyTorch DDP

import torch import torch.distributed as dist 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 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_replicas=world_size, rank=rank ) train_loader = DataLoader( train_dataset, batch_size=32, sampler=train_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,), nprocs=world_size, join=True)

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_node=8 \ --nnodes=2 \ --node_rank=0 \ --master_addr="192.168.1.1" \ --master_port=12345 \ train.py # Node 1 python -m torch.distributed.launch \ --nproc_per_node=8 \ --nnodes=2 \ --node_rank=1 \ --master_addr="192.168.1.1" \ --master_port=12345 \ train.py

6. 性能优化

7. 结论

分布式训练是训练大模型的关键技术。通过 DDP、混合精度和适当的优化策略，可以充分利用多 GPU 资源，显著缩短训练时间。