在Python中使用PyTorch进行多GPU训练时，常常出现有的显卡占用率接近满血，有的却长期低位运行，这种多GPU负载不均会拖慢整体效率。要从数据分配策略和模型并行两个方向去排查。
一、检查数据分配策略
数据分配不均是最常见的原因。使用nn.DataParallel时，默认在主卡上分发数据，若batch_size不能被显卡数整除，末尾批次会让部分卡空转。
1. 观察每张卡的输入数量
可以在训练循环里打印每张卡实际拿到的样本数：
import torch
import torch.nn as nn
# 模拟四卡环境
gpus = [0, 1, 2, 3]
batch = torch.arange(10) # 共10条样本
# 手动按卡数切分观察
chunks = list(batch.chunk(len(gpus)))
for i, c in enumerate(chunks):
print("gpu", gpus[i], "样本数", len(c))
# 输出可见最后一张卡只有2条，其余3条，造成轻微不均
2. 使用DistributedDataParallel替代
DistributedDataParallel让每个进程独立取数据，配合DistributedSampler能更均匀划分。
from torch.utils.data import DataLoader, DistributedSampler sampler = DistributedSampler(dataset, num_replicas=4, rank=local_rank) loader = DataLoader(dataset, batch_size=8, sampler=sampler) # 每张卡拿到约等量数据，避免主卡瓶颈
二、检查模型并行策略
当模型太大无法单卡放下，会采用模型并行，把层拆分到不同GPU。若某层计算量过重，就会出现负载倾斜。
1. 手动模型并行的简单示例
class ManualParallel(nn.Module):
def __init__(self):
super().__init__()
# 第一层放gpu0
self.layer1 = nn.Linear(1024, 2048).to('cuda:0')
# 第二层放gpu1
self.layer2 = nn.Linear(2048, 10).to('cuda:1')
def forward(self, x):
x = x.to('cuda:0')
x = self.layer1(x)
x = x.to('cuda:1')
return self.layer2(x)
# 若layer1参数量远大于layer2，gpu0会长期忙碌
2. 用torch.cuda.utilization观察
训练中另开脚本周期性读取利用率：
import torch
for i in range(5):
for d in range(2):
print(d, torch.cuda.utilization(d))
import time
time.sleep(2)
三、常见排查清单
- 确认batch_size能被GPU数量整除
- 优先使用DistributedDataParallel而非DataParallel
- 模型并行时平衡各设备层的计算与显存
- 检查数据预处理是否只在主进程造成阻塞
| 现象 | 可能原因 | 排查点 |
|---|---|---|
| 0卡高占用其他低 | DataParallel主卡分发 | 切换DDP |
| 某卡周期空载 | batch除不尽 | 调整batch_size |
| 模型并行一卡满一卡闲 | 层切分不均 | 重排模型层 |
通过上述数据分配与模型并行的检查，基本可以定位Python下PyTorch多GPU负载不均的来源，进而做针对性优化。