### 配置PyTorch的DistributedDataParallel 在PyTorch中设置和使用`DistributedDataParallel` (DDP) 进行多GPU或多节点模型训练涉及几个关键步骤。这些步骤确保了不同设备上的模型能够有效地协同工作，从而加快训练速度。 #### 初始化环境变量 为了使各个进程之间可以相互通信，在启动程序之前需设定一些必要的环境变量。这通常包括指定每台机器的地址以及端口等信息以便于建立连接[^1]。 ```bash export MASTER_ADDR=localhost export MASTER_PORT=12355 ``` 对于跨多个节点的情况，则需要替换 `MASTER_ADDR` 为主节点的实际IP地址，并保证所有参与方都能访问此地址。 #### 启动方法的选择 有多种方式来启动所需的多个进程： - **单机多卡**：可以通过简单的循环调用来创建所需数量的工作线程； - **多机集群**：推荐采用诸如 `torchrun`, `mpiexec`, 或者其他类似的工具来进行分布式作业提交[^4]。 这里给出基于 `torch.multiprocessing.spawn()` 的简单例子用于本地测试目的: ```python import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP import torch.multiprocessing as mp def setup(rank, world_size): # Initialize the process group with default settings. dist.init_process_group("nccl", rank=rank, world_size=world_size) def cleanup(): dist.destroy_process_group() def main(rank, world_size): setup(rank, world_size) device = f'cuda:{rank}' model = ... # Define your neural network here ddp_model = DDP(model.to(device), device_ids=[rank]) # Training loop goes here... cleanup() if __name__ == "__main__": n_gpus = torch.cuda.device_count() world_size = n_gpus mp.spawn(main, args=(world_size,), nprocs=n_gpus, join=True) ``` 这段代码展示了如何在一个进程中为每一可用显卡分配子任务，并通过 `setup()` 函数初始化通信组。注意这里的 `device_ids` 参数指定了哪个 GPU 应该负责当前实例的数据处理。 #### 数据集划分与采样器配置 当涉及到大规模数据集时，合理的样本分布至关重要。为此 PyTorch 提供了一个专门针对这种情况设计好的类——`DistributedSampler` 。它能自动完成对整个数据集合均匀切分的任务，使得每个 worker 只看到属于自己那份的部分。 ```python train_sampler = torch.utils.data.distributed.DistributedSampler( dataset=train_dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank()) loader = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler) ``` 上述片段说明了怎样构建带有自定义采样的加载器对象。其中 `num_replicas` 和 `rank` 分别对应总worker数目及其唯一编号。