pytorch 多卡训练

<think>好的，我现在需要帮助用户了解如何在PyTorch中实现多GPU训练。首先，用户提到的是使用Distributed Data Parallel（DDP），所以我应该集中在这个方法上。用户可能已经知道一些基础，但需要详细的步骤指南。 首先，我应该回顾PyTorch的DDP文档，确保自己理解正确。DDP通过多进程实现，每个进程控制一个GPU，进程间通信使用NCCL后端。需要初始化进程组，设置rank和world_size。然后，用DistributedSampler来分配数据，避免数据重复。模型需要用torch.nn.parallel.DistributedDataParallel包裹，确保梯度同步。 接下来，用户可能需要具体的代码示例。记得按照用户的要求，所有数学公式用$...$，独立公式用$$...$$，但这里可能用不上数学公式。代码块需要用正确的格式，比如Python的代码用三个反引号包裹。 用户提供的引用中有提到使用torch.distributed.launch，但现在PyTorch>=1.9推荐使用torch.distributed.run。所以需要更新这一点，避免用户使用过时的方法。还要注意单节点和多节点的命令区别，单节点使用--nproc_per_node，多节点需要指定--nnodes和--node_rank。 另外，数据加载部分要使用DistributedSampler，并且设置DataLoader的sampler参数，shuffle=False，因为Sampler会处理数据分片。Batch size需要按每个GPU调整，通常是总batch size除以GPU数量。 训练循环中要确保每个epoch前调用sampler.set_epoch(epoch)，这样数据划分会随机化，避免每个epoch数据顺序相同。还要注意清理进程，使用dist.destroy_process_group()。 用户可能遇到的常见问题包括端口冲突，可以设置不同的MASTER_PORT环境变量。另外，确保所有进程正确同步，避免死锁。性能优化方面，可能提到梯度累积，但DDP本身已经高效，主要在于正确配置。 最后，生成相关问题时要涵盖DDP与DP的区别、如何选择后端、多节点配置、性能调优、错误排查等，帮助用户进一步学习。</think>在PyTorch中实现多GPU分布式训练主要通过`DistributedDataParallel`（DDP）实现，以下是具体步骤： --- ### 1. 基础配置 ```python import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def setup(rank, world_size): dist.init_process_group( backend="nccl", # GPU推荐NCCL后端 init_method="env://", rank=rank, world_size=world_size ) torch.cuda.set_device(rank) ``` - **关键参数**： - `backend`：GPU集群使用`nccl`，CPU集群用`gloo` - `init_method`：指定进程发现方式（环境变量`env://`为默认） - `rank`：当前进程的全局序号（0为master节点） - `world_size`：总进程数（通常等于GPU总数） --- ### 2. 数据并行加载 ```python from torch.utils.data.distributed import DistributedSampler def prepare_dataloader(dataset, batch_size): sampler = DistributedSampler( dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank(), shuffle=True ) return DataLoader( dataset, batch_size=batch_size, sampler=sampler, num_workers=4, pin_memory=True ) ``` - **注意事项**： - 每个GPU获得数据的非重叠子集（自动数据分片）[^1] - `batch_size`应设为单卡batch大小（总batch_size = batch_size * num_gpus） --- ### 3. 模型封装 ```python model = MyModel().to(rank) model = DDP(model, device_ids=[rank]) optimizer = torch.optim.Adam(model.parameters()) ``` - DDP工作原理： - 前向传播：各GPU独立计算 - 反向传播：自动同步梯度（AllReduce算法）[^2] - 参数更新：各GPU保持一致性 --- ### 4. 启动训练脚本 #### 单机多卡（4卡）： ```bash python -m torch.distributed.run \ --nproc_per_node=4 \ --nnodes=1 \ --node_rank=0 \ train.py ``` #### 多机多卡（需配置SSH免密登录）： ```bash # 节点0 python -m torch.distributed.run \ --nproc_per_node=8 \ --nnodes=2 \ --node_rank=0 \ --master_addr="192.168.1.1" \ --master_port=1234 \ train.py # 节点1 python -m torch.distributed.run \ --nproc_per_node=8 \ --nnodes=2 \ --node_rank=1 \ --master_addr="192.168.1.1" \ --master_port=1234 \ train.py ``` --- ### 5. 完整训练模板 ```python def main(rank, world_size): setup(rank, world_size) dataset = MyDataset() dataloader = prepare_dataloader(dataset, 32) model = DDP(MyModel().to(rank), device_ids=[rank]) optimizer = torch.optim.SGD(model.parameters(), lr=0.01) for epoch in range(100): dataloader.sampler.set_epoch(epoch) # 保证shuffle随机性 for batch in dataloader: inputs, labels = batch outputs = model(inputs) loss = torch.nn.functional.cross_entropy(outputs, labels) loss.backward() optimizer.step() optimizer.zero_grad() dist.destroy_process_group() if __name__ == "__main__": world_size = torch.cuda.device_count() torch.multiprocessing.spawn(main, args=(world_size,), nprocs=world_size) ``` --- ### 常见问题解决 - **端口冲突**：通过`export MASTER_PORT=29500`指定不同端口 - **同步失败**：检查所有进程的`world_size`和`rank`配置一致性 - **显存不足**：减少`batch_size`或使用梯度累积（需手动实现） ---