PyTorch多GPU负载不均解决方法

PyTorch中DataParallel因主从式单进程设计导致cuda:0严重过载（显存高20%~30%、利用率常年95%+），而其他GPU闲置，根源在于其无法绕过Python GIL和主线程瓶颈；真正高效的多GPU训练需转向DistributedDataParallel——通过多进程隔离、NCCL通信、DistributedSampler数据分片及严格的设备对齐（如tensor自动迁移、rank=0专属日志与保存），才能实现各卡计算负载均衡与线性加速，否则极易陷入“伪多卡”陷阱：表面启动四卡，实则性能被cuda:0死锁。

为什么DataParallel会导致GPU负载不均

DataParallel在单进程内用主GPU（通常是cuda:0）做前向/反向的汇总和参数同步，其余GPU只负责计算子任务。这导致cuda:0不仅要跑自己的batch，还要收发梯度、拼接输出、更新参数——I/O和计算双重压力远超其他卡。常见现象是nvidia-smi里cuda:0显存占用高20%~30%，GPU-Util持续95%+，而cuda:1~cuda:3经常掉到40%以下。

根本原因是DataParallel不是真正的并行：它把一个batch按device数量切分，但所有子模块仍运行在同一个Python进程里，无法绕过GIL和主线程调度瓶颈。

改用DistributedDataParallel的关键三步

必须启动多个独立进程，每个进程绑定唯一GPU，并由torch.distributed协调同步。漏掉任意一步都会报错或退化为单卡。

• 启动方式必须用torch.distributed.launch或torchrun（推荐后者，PyTorch 1.9+默认）：

  torchrun --nproc_per_node=4 train.py

• 每个进程内要初始化进程组：torch.distributed.init_process_group(backend="nccl")，且rank和world_size需从环境变量读取（os.environ["LOCAL_RANK"]），不能硬编码
• 模型包装必须在torch.nn.parallel.DistributedDataParallel中，且device_ids只传单个GPU ID（如[int(os.environ["LOCAL_RANK"])]），不能像DataParallel那样传列表

DistributedDataParallel训练时数据加载要注意什么

DistributedSampler不是可选优化项，而是必须项。它确保每个进程只看到全局数据集的一个互斥子集，否则所有进程会重复处理全部数据，loss虚低、收敛异常。

• 实例化DataLoader时，sampler参数必须设为DistributedSampler(dataset, shuffle=True)
• 训练循环中，每个epoch开始前要调用sampler.set_epoch(epoch)，否则shuffle失效
• DataLoader的num_workers建议设为0或≤2，多worker在分布式下易触发文件句柄泄漏或死锁，尤其在NFS路径上

常见报错与绕过陷阱

最典型的是RuntimeError: Expected all tensors to be on the same device，本质是某个tensor没随LOCAL_RANK移到对应GPU。比如日志打印、loss计算、metric累积时用了.cpu()或.item()后又试图回传，或者验证时忘了用torch.no_grad()和model.eval()。

• 所有tensor操作前加.to(device)，其中device = torch.device(f"cuda:{os.environ['LOCAL_RANK']}")
• 避免在forward里写print(x.shape)——不同卡输出混在一起；改用if rank == 0: print(...)
• torch.save只在rank == 0时执行，否则多进程同时写同一文件会损坏
• 使用torch.compile时，必须在DistributedDataParallel包装之后再编译，顺序颠倒会丢失梯度同步逻辑

多卡训练真正难的不是启动，而是让每个进程干净地隔离设备、内存、I/O和随机种子——稍有交叉，就变成“看起来是多卡，实际卡在cuda:0”。

好了，本文到此结束，带大家了解了《PyTorch多GPU负载不均解决方法》，希望本文对你有所帮助！关注golang学习网公众号，给大家分享更多文章知识！