PyTorch的分布式数据并行（DistributedDataParallel, DDP）模式使得在多GPU、甚至多节点环境下高效地训练深度学习模型成为可能。当涉及到多机多卡环境下的DDP配置时，有几个关键点需要关注： ### 1. 环境准备 首先确保所有机器都安装了相同的 PyTorch 版本，并能正常访问共享文件系统用于存储检查点等。 ### 2. 初始化进程组 为了实现跨设备通信，在程序开始处需通过 `init_process_group` 函数初始化一个全局唯一的进程组: ```python import torch.distributed as dist def init_ddp(rank, world_size): # 使用 nccl 后端对于 GPU 训练性能更佳 backend = 'nccl' if rank == 0: print("Initializing Process Group...") # 设置 master 地址及端口供其他 worker 连接 os.environ['MASTER_ADDR'] = 'localhost' # 或者是集群中某台机器 IP os.environ['MASTER_PORT'] = '8738' dist.init_process_group(backend=backend, rank=rank, world_size=world_size) ``` 注意替换 `'localhost'` 和 `'8738'` 到实际可用的服务地址和未占用端口号；如果是单机多卡可以保持默认值不变。 ### 3. 模型封装与优化器创建 将模型实例传递给 `torch.nn.parallel.DistributedDataParallel()` 构造函数之前，先将其移动到正确的设备上 (`cuda:local_rank`) 并设置 `find_unused_parameters=True` 参数来处理某些特殊情况下的梯度同步问题. 接着像往常一样定义损失函数和选择合适的优化算法即可完成基本准备工作。 ```python model.to(local_rank) ddp_model = DistributedDataParallel(model, device_ids=[local_rank], find_unused_parameters=True) optimizer = optim.Adam(ddp_model.parameters(), lr=args.lr) criterion = nn.CrossEntropyLoss() ``` 这里假设每张显卡都有独立编号作为其对应的本地 rank (即 local_rank)，它表示当前进程中使用的特定 GPU ID。 ### 4. 数据加载 为了让每个 worker 获得不重复的数据批次，在 DataLoader 中添加参数 `sampler=DistributedSampler(dataset)` 来代替原来的 shuffle 属性。这会自动保证整个 epoch 内样本遍历一次且不会遗漏或重复。 ```python train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=batch_size, sampler=train_sampler, num_workers=num_workers, pin_memory=True) ``` ### 5. 主循环与其他注意事项 最后就是编写主训练循环逻辑了——记得只让 main process 做日志记录等工作以免造成冗余输出影响效率。 此外还需要适当调整学习率调度策略以适应更大规模batch size带来的变化等因素。 以上就是在 PyTorch 的 DDP 下进行多机多 GPU 训练的基本步骤。需要注意的是实际情况可能会更为复杂，因此建议参考官方文档获取最新信息和支持。