### PyTorch `nn` 模块下 `DistributedDataParallel` (DDP) 的用法 #### 初始化 DDP 和环境设置 为了使用 `torch.nn.parallel.DistributedDataParallel`, 需要先初始化进程组并配置必要的环境变量。通常通过命令行参数传递本地排名 (`local_rank`) 来启动多个进程。 ```bash python -m torch.distributed.launch --nproc_per_node=8 train.py --local_rank=$LOCAL_RANK ``` 这段脚本会自动分配可用GPU上的训练任务[^1]。 #### 创建模型实例并与 DDP 绑定 当定义好神经网络结构之后，在主函数中创建模型对象，并将其封装到 `DistributedDataParallel` 类内： ```python import torch import torch.nn as nn from torch.nn.parallel import DistributedDataParallel as DDP model = MyModel().to(device) ddp_model = DDP(model, device_ids=[args.local_rank], output_device=args.local_rank) ``` 这里指定了设备ID列表以及输出目标设备，确保每个进程只操作指定的GPU资源。 #### 数据加载器配置 为了让不同rank之间的样本不重复且均匀分布，需采用特殊的采样策略——`DistributedSampler`: ```python train_sampler = torch.utils.data.distributed.DistributedSampler(dataset) train_loader = DataLoader( dataset, batch_size=batch_size, shuffle=(train_sampler is None), num_workers=num_workers, pin_memory=True, sampler=train_sampler ) ``` 此方法可以保证整个集群中的所有worker共同遍历完整的数据集一次而不重叠[^2]。 #### 前向传播与损失计算 由于DDP模式下各个GPU执行相同的操作流程，因此标准的前向传播和损失计算逻辑保持不变： ```python outputs = ddp_model(inputs) loss = criterion(outputs, targets) ``` 注意此时传入的数据应当已经由对应的sampler进行了预处理，使得每张卡上接收到的是不同的batch子集。 #### 反向传播及优化步骤 同样地，反向传播过程也不受影响；只需调用`.backward()`即可触发梯度累积机制，随后利用Optimizer完成参数更新动作： ```python optimizer.zero_grad() loss.backward() optimizer.step() ``` 值得注意的是，默认情况下DDP会在每次迭代结束时同步各节点间的梯度信息，从而实现全局一致性的权重调整。