PyTorch DDP实战：从单卡到多卡，手把手教你搭建分布式训练环境

1. 为什么需要分布式训练当你用单卡训练一个ResNet模型时可能觉得速度还行。但当你开始训练BERT这类大模型时等待训练完成就像看着油漆慢慢变干。这时候分布式训练就像给你的模型装上了火箭推进器——它能让你同时使用多张GPU卡把训练时间从几天压缩到几小时。想象你是个包工头原来只能指挥一个工人搬砖单卡训练现在突然有了8个工人8卡训练。如果让他们各自为战效率反而可能下降。分布式训练的精妙之处在于它让这些工人像训练有素的蚂蚁军团一样协同工作——每张卡处理不同的数据批次但会定期同步学到的知识。2. DDP环境搭建实战2.1 硬件准备与驱动检查首先确认你的机器有至少两块NVIDIA GPU别笑我真见过有人对着集成显卡折腾CUDA。运行nvidia-smi应该能看到类似这样的输出----------------------------------------------------------------------------- | NVIDIA-SMI 525.85.12 Driver Version: 525.85.12 CUDA Version: 12.0 | |--------------------------------------------------------------------------- | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | | | | MIG M. | || | 0 NVIDIA RTX 3090 On | 00000000:17:00.0 Off | Off | | 30% 45C P0 120W / 350W| 2345MiB / 24576MiB | 30% Default | | | | N/A | --------------------------------------------------------------------------- | 1 NVIDIA RTX 3090 On | 00000000:65:00.0 Off | Off | | 25% 38C P0 110W / 350W| 1234MiB / 24576MiB | 25% Default | | | | N/A | ---------------------------------------------------------------------------如果看到多块GPU但CUDA不可用先检查驱动安装。我遇到过最坑的情况是服务器自动更新了内核但没重装NVIDIA驱动症状是nvidia-smi能看但torch.cuda.is_available()返回False。2.2 PyTorch与NCCL安装安装支持DDP的PyTorch版本建议1.8。用conda安装时记得指定cudatoolkit版本conda install pytorch torchvision torchaudio pytorch-cuda12.1 -c pytorch -c nvidia验证NCCL支持DDP的后端通信库import torch print(torch.cuda.nccl.version()) # 应该输出类似(2, 14, 3)的版本号如果报错可能需要单独安装NCCLconda install -c conda-forge nccl3. 单卡代码改造指南3.1 初始化分布式环境这是最容易踩坑的部分。假设我们有两张卡需要在每个进程中执行import torch.distributed as dist def setup(rank, world_size): # 关键配置所有进程必须一致 os.environ[MASTER_ADDR] localhost # 单机多卡就用localhost os.environ[MASTER_PORT] 12355 # 随便选个没被占用的端口 # 初始化进程组 dist.init_process_group( backendnccl, # NVIDIA显卡就用nccl rankrank, world_sizeworld_size ) torch.cuda.set_device(rank) # 每个进程绑定不同GPU注意MASTER_PORT冲突是常见问题。我有次在服务器上跑实验总是卡住后来发现是同事用了相同端口。3.2 改造DataLoader单卡时代的DataLoader要升级为分布式版本from torch.utils.data.distributed import DistributedSampler def get_dataloader(dataset, batch_size): sampler DistributedSampler( dataset, num_replicasworld_size, rankrank, shuffleTrue # 分布式环境下也要保持shuffle ) return DataLoader( dataset, batch_sizebatch_size, samplersampler, num_workers4, pin_memoryTrue # 这个对加速很关键 )重点说明每个进程只会看到部分数据自动分区batch_size是单卡的batch大小总batch_size是batch_size*world_size必须设置sampler否则不同卡会处理相同数据3.3 包装模型为DDP这是最神奇的部分——只需一行代码就能让模型支持分布式from torch.nn.parallel import DistributedDataParallel as DDP model MyAwesomeModel().to(rank) model DDP(model, device_ids[rank])但有两个陷阱要注意必须在DDP包装后再创建optimizer否则optimizer看不到完整的参数结构保存模型时要取model.module.state_dict()因为DDP给原始模型套了壳4. 启动训练的正确姿势4.1 单机多卡启动推荐使用PyTorch自带的启动器python -m torch.distributed.launch \ --nproc_per_node2 \ # 每台机器进程数GPU数 --nnodes1 \ # 机器数量 --node_rank0 \ # 当前机器序号 --master_addrlocalhost \ --master_port12355 \ train.py更简洁的torchrun方式PyTorch 1.9torchrun --nproc_per_node2 train.py4.2 多机多卡配置假设有两台机器IP为192.168.1.101和192.168.1.102在每台机器上执行# 机器0上执行 torchrun \ --nproc_per_node4 \ --nnodes2 \ --node_rank0 \ --master_addr192.168.1.101 \ --master_port12355 \ train.py # 机器1上执行只有node_rank不同 torchrun \ --nproc_per_node4 \ --nnodes2 \ --node_rank1 \ --master_addr192.168.1.101 \ --master_port12355 \ train.py5. 实战中的性能优化技巧5.1 梯度累积实现大batch当显存不足时可以用梯度累积模拟大batchoptimizer.zero_grad() for i, (inputs, targets) in enumerate(dataloader): outputs model(inputs) loss criterion(outputs, targets) loss.backward() if (i1) % accumulation_steps 0: # 每累积N步更新一次 optimizer.step() optimizer.zero_grad()5.2 混合精度训练DDP完美支持AMP自动混合精度scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs model(inputs) loss criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()我在V100上测试混合精度能使训练速度提升2倍以上而且精度损失可以忽略。5.3 通信优化参数DDP默认在所有进程间同步梯度但有些情况可以优化model DDP( model, device_ids[rank], output_devicerank, broadcast_buffersFalse, # 不广播BN层统计量 find_unused_parametersTrue # 动态图可能需要 )6. 常见坑与解决方案6.1 死锁问题当某个进程报错退出时其他进程会卡住。解决方案是加异常处理try: train() except Exception as e: print(fRank {rank} failed: {e}) dist.destroy_process_group() raise6.2 内存泄漏排查DDP的内存泄漏通常来自没有正确释放进程组确保在训练结束后调用dist.destroy_process_group()在循环中重复创建DDP模型应该只创建一次6.3 日志处理技巧所有进程都打印日志会非常混乱。推荐只在rank 0打印if dist.get_rank() 0: print(fEpoch {epoch} loss: {loss.item()})或者用torch.distributed的日志工具import torch.distributed as dist dist.all_reduce(loss, opdist.ReduceOp.SUM) avg_loss loss / dist.get_world_size()7. 完整代码示例最后放一个可直接运行的完整示例基于CIFAR10import os import torch import torch.distributed as dist import torchvision from torch.nn.parallel import DistributedDataParallel as DDP class CIFAR10Model(torch.nn.Module): def __init__(self): super().__init__() self.conv1 torch.nn.Conv2d(3, 6, 5) self.pool torch.nn.MaxPool2d(2, 2) self.conv2 torch.nn.Conv2d(6, 16, 5) self.fc1 torch.nn.Linear(16 * 5 * 5, 120) self.fc2 torch.nn.Linear(120, 84) self.fc3 torch.nn.Linear(84, 10) def forward(self, x): x self.pool(torch.nn.functional.relu(self.conv1(x))) x self.pool(torch.nn.functional.relu(self.conv2(x))) x torch.flatten(x, 1) x torch.nn.functional.relu(self.fc1(x)) x torch.nn.functional.relu(self.fc2(x)) x self.fc3(x) return x def setup(rank, world_size): os.environ[MASTER_ADDR] localhost os.environ[MASTER_PORT] 12355 dist.init_process_group(nccl, rankrank, world_sizeworld_size) torch.cuda.set_device(rank) def cleanup(): dist.destroy_process_group() def get_dataloader(rank, world_size): transform torchvision.transforms.Compose([ torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) dataset torchvision.datasets.CIFAR10( root./data, trainTrue, downloadTrue, transformtransform ) sampler torch.utils.data.distributed.DistributedSampler( dataset, num_replicasworld_size, rankrank, shuffleTrue ) return torch.utils.data.DataLoader( dataset, batch_size32, samplersampler, num_workers4, pin_memoryTrue ) def train(rank, world_size): setup(rank, world_size) model CIFAR10Model().to(rank) ddp_model DDP(model, device_ids[rank]) criterion torch.nn.CrossEntropyLoss() optimizer torch.optim.SGD(ddp_model.parameters(), lr0.001, momentum0.9) dataloader get_dataloader(rank, world_size) for epoch in range(10): ddp_model.train() dataloader.sampler.set_epoch(epoch) # 保证shuffle有效性 for inputs, labels in dataloader: inputs, labels inputs.to(rank), labels.to(rank) optimizer.zero_grad() outputs ddp_model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() if rank 0: print(fEpoch {epoch} completed) cleanup() if __name__ __main__: world_size torch.cuda.device_count() torch.multiprocessing.spawn( train, args(world_size,), nprocsworld_size, joinTrue )运行命令torchrun --nproc_per_node2 train.py这个示例包含了DDP的核心要素分布式初始化、数据并行采样、模型包装和正确启动方式。你可以在此基础上扩展更复杂的模型和训练逻辑。