外观
PyTorch nn.DataParallel 多GPU训练详解
内容整理自学习笔记，仅供面试备考参考；不构成录用、培训或考试承诺。
1. PyTorch中GPU操作的默认行为
PyTorch的GPU操作默认是异步的——需要使用GPU的函数会进入特定设备队列等待执行，PyTorch支持并行计算，但并行效果对调用者不可见。
PyTorch的多GPU并行计算是数据级并行，相当于开了多个进程，各自独立运行后再整合。
2. nn.DataParallel函数介绍
python
torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)| 参数 | 说明 |
|---|---|
| module | 要并行训练的模型，每个GPU上都会拷贝一份完整模型 |
| device_ids | 参与训练的GPU列表，如 [0, 1, 2] |
| output_device | 指定输出GPU，默认为 device_ids[0]（第一块卡） |
| dim | 数据分散维度，默认为0（按batchsize划分） |
使用示例：
python
device_ids = [0, 1]
net = torch.nn.DataParallel(net, device_ids=device_ids)输出结果中模型被DataParallel包裹，自动在指定GPU上并行计算并合并结果。
3. nn.DataParallel的处理逻辑（参数服务器模式）
参数服务器框架中：
| 角色 | 说明 |
|---|---|
| Worker | 计算GPU，负责前向/反向计算 |
| Server | 梯度聚合GPU，负责AllReduce操作 |
处理流程：
- 若干块计算GPU（如GPU0~GPU2）和1块梯度收集GPU
- 每块计算GPU上拷贝一份完整模型参数
- 将一份数据X均匀分给不同计算GPU
- 每块计算GPU做一轮FWD和BWD，算得一份梯度G
- 每块计算GPU将梯度push给梯度收集GPU，做聚合操作（梯度累加）
- 梯度收集GPU聚合完毕后，计算GPUpull完整梯度，更新模型参数W
- 聚合再下发梯度的操作称为AllReduce
注意： 1个Worker或Server下可以不止1块GPU；Server可只做梯度聚合，也可梯度聚合+全量参数更新一起做。
4. nn.DataParallel常见问题
4.1 多GPU计算为何反而更慢？
DataParallel将每个batch数据平均分配到不同device上并行计算，再合并结果。GPU之间的开关和通讯过程占了大部分时间开销。
诊断： 使用
watch -n 1 nvidia-smi查看GPU占用率，若各GPU占用率均低于50%，说明多GPU计算比单GPU更慢。
4.2 如何保存和加载多GPU训练模型？
DataParallel实际上是一个nn.Module，保存时需调用 net.module：
python
net = torch.nn.Linear(10, 1)
net = torch.nn.DataParallel(net, device_ids=[0, 3])
# 保存
torch.save(net.module.state_dict(), './networks/multiGPU.h5')
# 加载
new_net = torch.nn.Linear(10, 1)
new_net.load_state_dict(torch.load("./networks/multiGPU.h5"))4.3 为什么第一块卡显存占用更多？
output_device 默认为 device_ids[0]，即第一块卡。输入数据是并行分配的，但output loss只在第一块GPU上相加计算，导致第一块GPU负载远大于其他卡。
4.4 多卡训练时的warning问题
UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars;
will instead unsqueeze and return a vector.原因： 多GPU计算时返回多个loss，DP按dim0分块数据。每个GPU独立计算loss，返回向量维度等于GPU数量。
解决方法： 使用 size_average=False, reduce=True 参数，或使用gather方式求loss平均，确保不同GPU数量的最终平均loss一致。
4.5 device_ids[0]被占用的解决方法
若0号卡被占用，需要先指定可见的devices：
python
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"
device_ids = [0, 1] # 逻辑0对应物理2号卡，逻辑1对应物理3号卡
net = torch.nn.DataParallel(net, device_ids=device_ids)注意： 这两行代码须放在import package之后、训练代码之前。
5. nn.DataParallel参数更新方式
python
optimizer = torch.optim.SGD(net.parameters(), lr=lr)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)完整更新流程：
- DataLoader把数据通过多个worker读到主进程内存
- 通过tensor的split语义，将batch切分成多个小batch，送往不同cuda设备
- 不同cuda设备完成前向计算，输出被gather到主cuda设备，loss在此计算
- loss被scatter到每个cuda设备，每块设备通过BP计算梯度
- 每块设备的梯度被reduce到主cuda设备，权重在主设备更新
- 下次迭代前，主cuda设备将参数broadcast到其他cuda设备，完成同步
通信原语： broadcast（主进程分发相同数据）、scatter（分发不同部分）、gather（收集数据）、reduce（收集+操作如SUM）。加"all"前缀变为多对多。
6. nn.DataParallel优点
- 使用简单、易于理解
- 只需在模型前加
nn.DataParallel即可 - 能充分利用多个GPU进行训练
7. nn.DataParallel缺点
| 缺点 | 说明 |
|---|---|
| 内存占用 | 将整个模型复制到每个GPU，大模型可能内存不足 |
| 数据通信 | 需汇总每个GPU梯度，大量数据通信导致训练速度下降 |
| 不支持分布式 | 要求所有GPU在同一节点上 |
| 不支持Apex | 不能与Apex混合精度训练配合使用 |
8. nn.DataParallel实战代码
python
import torch
import torch.nn as nn
import torch.optim as optim
gpus = [0, 1, 2, 3]
torch.cuda.set_device('cuda:{}'.format(gpus[0]))
train_dataset = ... # build train dataset
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=...)
# 核心步骤：先model.to(device)，再DataParallel包裹
model = ... # get model
model = nn.DataParallel(model.to(device), device_ids=gpus, output_device=gpus[0])
optimizer = optim.SGD(model.parameters())
for epoch in range(100):
for batch_idx, (images, target) in enumerate(train_loader):
images = images.cuda(non_blocking=True)
target = target.cuda(non_blocking=True)
output = model(images)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()关键注意点：
- model加载到GPU的顺序：先get model → model.to(device) → 再用DP
- 数据加载到GPU的时机：在iter的forward之前完成
to(device)- batch_size应能被GPU数量整除，否则需改写data_parallel添加chunk_size
- model须先加载到
device_ids[0]上再包裹DP