### PyTorch DistributedDataParallel 分布式训练教程 #### 使用场景与动机 为了加速大规模深度学习模型的训练过程以及更高效地利用硬件资源，分布式数据并行成为一种常见策略。通过将批量的数据划分给不同的进程处理，可以在多个GPU甚至跨机器之间分摊工作负载，从而显著缩短单次迭代所需时间。 #### 整体架构概述 在PyTorch中实现这一目标主要依赖`torch.nn.parallel.DistributedDataParallel`模块[^1]。此接口允许开发者轻松设置一个多GPU或多节点环境下的同步SGD算法执行流程。它内部实现了参数梯度平均化机制，确保所有参与者的权重更新一致；同时也提供了灵活配置选项支持如自动批标准化等特性。 #### 关键组件介绍 - **初始化通信后端** 需要指定使用的通讯库（例如NCCL对于NVIDIA GPU），并通过环境变量或命令行参数定义集群拓扑结构信息，包括但不限于总设备数(`WORLD_SIZE`)、当前rank编号(RANK)及master地址/端口(MASTER_ADDR, MASTER_PORT)[^2]。 - **自定义采样器(DistributedSampler)** 当前版本推荐使用内置工具类`distributed.DistributedSampler`, 它能够依据全局worker数量合理分割输入dataset成互不重叠的小批次供各子任务独立读取而不重复访问相同样本点。 - **模型封装** 将待训练网络实例传入DDP构造函数完成包装操作即可开启后续协作模式下的反向传播计算逻辑。 #### 最小可行代码示例 下面给出一段简单的MNIST手写数字识别案例作为入门指南： ```python import os import torch.distributed as dist from torchvision import datasets, transforms from torch.utils.data.dataloader import DataLoader from torch.nn.parallel import DistributedDataParallel as DDP from torch.optim import SGD from torch import nn def setup(rank, world_size): # 初始化进程组 dist.init_process_group( backend='nccl', init_method=f'tcp://localhost:{os.environ["MASTER_PORT"]}', rank=rank, world_size=world_size) class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(1, 32, kernel_size=(5, 5)) ... transform = transforms.Compose([ transforms.ToTensor(), ]) train_dataset = datasets.MNIST('.', train=True, transform=transform) sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) loader = DataLoader(dataset=train_dataset, batch_size=64, sampler=sampler) model = SimpleCNN().to('cuda') ddp_model = DDP(model) optimizer = SGD(ddp_model.parameters(), lr=0.01) for epoch in range(num_epochs): loader.sampler.set_epoch(epoch) for data, target in loader: optimizer.zero_grad() output = ddp_model(data.cuda()) loss = F.nll_loss(output, target.cuda()) loss.backward() optimizer.step() if __name__ == "__main__": num_gpus = torch.cuda.device_count() mp.spawn(setup, args=(num_gpus,), nprocs=num_gpus, join=True) ```