### PyTorch DistributedDataParallel (DDP) 多GPU训练教程 #### 1. 基本原理 DistributedDataParallel（简称 DDP）是 PyTorch 提供的一种高效的分布式训练工具，能够显著提升多 GPU 训练的速度和效率。相比于 DataParallel，DDP 支持更灵活的场景，不仅限于单机多卡环境，还扩展到了多机多卡的支持[^1]。 #### 2. 单机多卡配置下的基本流程 以下是构建基于 DDP 的单机多卡训练框架的关键步骤： - **初始化进程组** 使用 `torch.distributed.init_process_group` 初始化进程组，指定通信后端以及设备数量。 - **数据分片** 利用 `torch.utils.data.DistributedSampler` 对数据集进行切分，确保每个 GPU 只处理一部分数据子集。 - **模型并行化** 将模型实例传递给 `torch.nn.parallel.DistributedDataParallel` 进行封装，从而实现梯度同步和其他必要的操作。 - **启动多个进程** 需要在脚本中通过 Python 的 `multiprocessing.spawn` 或者命令行的方式创建多个独立的训练进程。 --- #### 3. 示例代码 以下是一个完整的单机多卡 DDP 实现示例： ```python import os import torch import torch.nn as nn import torch.optim as optim from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, Dataset, DistributedSampler class ToyDataset(Dataset): def __init__(self): self.data = [(i, i * 2) for i in range(100)] def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx] def setup(rank, world_size): """ Initialize the process group """ os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '12355' # initialize the process group torch.distributed.init_process_group( backend='nccl', init_method='env://', world_size=world_size, rank=rank ) def cleanup(): """ Destroy the process group after training is done """ torch.distributed.destroy_process_group() class SimpleModel(nn.Module): def __init__(self): super(SimpleModel, self).__init__() self.fc = nn.Linear(2, 1) def forward(self, x): return self.fc(x) def train_fn(rank, world_size): setup(rank, world_size) model = SimpleModel().to(rank) ddp_model = DDP(model, device_ids=[rank]) dataset = ToyDataset() sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank) dataloader = DataLoader(dataset, batch_size=8, shuffle=False, sampler=sampler) criterion = nn.MSELoss().to(rank) optimizer = optim.SGD(ddp_model.parameters(), lr=0.01) for epoch in range(5): # Train for a few epochs sampler.set_epoch(epoch) for data, target in dataloader: inputs, labels = data.to(rank).float(), target.to(rank).unsqueeze(-1).float() outputs = ddp_model(inputs) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() print(f"Rank {rank}, Epoch [{epoch}], Loss: {loss.item()}") cleanup() if __name__ == "__main__": WORLD_SIZE = torch.cuda.device_count() # Number of GPUs available on your machine torch.multiprocessing.spawn(train_fn, args=(WORLD_SIZE,), nprocs=WORLD_SIZE, join=True) ``` 上述代码展示了如何利用 DDP 构建一个多 GPU 并行化的训练过程[^1][^2]。 --- #### 4. 关键点解析 - **进程管理** 每个 GPU 上运行一个单独的进程，因此需要显式地设置每台机器上的 CUDA_VISIBLE_DEVICES 和 MASTER 地址/端口信息。 - **数据采样器** `DistributedSampler` 是为了保证不同进程中加载的数据互不重叠而设计的组件。如果忽略该部分，则可能导致重复计算或错误收敛行为。 - **性能优化建议** - 合理调整批量大小以充分利用硬件资源； - 开启混合精度训练 (`amp`) 来减少内存占用并加速运算速度； - 考虑使用第三方库简化复杂逻辑，比如 Horovod、PyTorch Lightning 等[^4]。 ---