根本原因是nn.dataparallel将batch按序分配但仅在gpu0汇总loss和梯度,导致gpu0显存与计算压力倾斜;ddp通过本地计算+all-reduce实现真并行,负载均衡。

为什么 nn.DataParallel 会让某块 GPU 显存爆满而其他空闲?
根本原因不是模型本身,而是 nn.DataParallel 的默认数据分发逻辑:它把一个 batch 拆成子 batch,按顺序分配给 GPU0、GPU1……GPUn,但**只在 GPU0 上做最终的 loss 计算和梯度汇总**。这意味着 GPU0 不仅要跑自己的前向/反向,还要额外承担 gather 操作(比如把所有 GPU 的 logits 拼成一个 tensor)、loss 计算、以及反向传播时的梯度累加——显存和计算压力天然倾斜。
常见错误现象:RuntimeError: CUDA out of memory 只报在 cuda:0,nvidia-smi 显示 GPU0 显存占用远高于其他卡,GPU 利用率却可能全都不高。
- 使用场景:batch size 较大、模型输出维度高(如分类数多)、或自定义 loss 中有 gather/concat 操作时,GPU0 压力会指数级放大
-
nn.DataParallel不支持跨卡优化器状态同步,也不能用torch.compile,兼容性差 - 参数差异:它只接受单个 device_ids 列表,不暴露分片粒度控制;没有内置的负载感知调度
怎么快速验证是不是数据分配策略导致的负载不均?
不用改模型,先做两件事定位问题:
- 运行时用
nvidia-smi -l 1观察各卡显存峰值和持续利用率,重点看 GPU0 是否长期高位 - 在训练循环里加一行日志:
print(f"Batch on {x.device}, shape: {x.shape}"),确认输入 tensor 确实被分到了不同设备(注意:x是每个 GPU 上的 local batch,不是原始 batch) - 临时把模型换成最简线性层 +
nn.CrossEntropyLoss(),如果负载仍不均,基本排除模型结构问题,锁定是并行机制本身
这个阶段别急着换方案——很多所谓“不均”其实是 batch size 没对齐 GPU 数量(比如 32 卡用 batch=65),导致最后一张卡分到 1 个样本,其余分到 2 个,表面不均实则合理。
让 macOS 自动化像说话一样简单 —— 用代码控制你的 Mac,无需学习 AppleScript 系统要求 macOS 14+ (Sonoma) Python 3.9+ 需要开启:辅助功能 + 屏幕录制权限
用 nn.DistributedDataParallel 替代前必须处理的三个坑
nn.DistributedDataParallel(DDP)不是开箱即用的“更高级 DataParallel”,它要求显式管理进程、数据和同步点。跳过这三步,大概率卡死或报 RuntimeError: Expected all tensors to be on the same device。
- 必须手动启动多进程:不能靠
python train.py直接跑,得用torch.distributed.run或mp.spawn,例如python -m torch.distributed.run --nproc_per_node=4 train.py - 数据集必须用
DistributedSampler,且shuffle=True时需在每个 epoch 调用sampler.set_epoch(epoch),否则各进程看到的数据子集重复 - 模型初始化后必须调用
model = nn.DistributedDataParallel(model, device_ids=[local_rank]),其中local_rank来自torch.distributed.get_rank(),不能硬编码
DDP 下每张卡只跑完整 batch 的一个子集,loss 计算、梯度更新都在本地完成,最后才通过 all-reduce 同步梯度——这才是真并行,负载自然均衡。
检查 DataLoader 的 num_workers 和 pin_memory 是否拖了后腿
即使用了 DDP,如果数据加载没调好,GPU 会大量时间等数据,看起来像“某卡卡住”,其实是 CPU 瓶颈伪装成 GPU 负载不均。
-
num_workers > 0时,每个 worker 进程会复制一份模型和数据集对象,若数据集__init__里加载了大文件或全局变量,会导致内存暴涨甚至 OOM -
pin_memory=True只对 GPU 训练有效,但它会让 DataLoader 把 batch 预分配在 pinned memory 中,加速 host-to-device 传输;但若系统内存不足,反而引发 swap,拖慢整体 - 验证方法:把
num_workers设为 0,观察 GPU 利用率是否从 10% 升到 60%+;再逐步加到min(32, os.cpu_count())并监控内存增长
真正复杂的点在于:DDP 的每个进程都有独立的 DataLoader 实例,所以 num_workers 是 per-process 的——4 卡 DDP 开 8 个 workers,实际会起 32 个子进程,CPU 和内存压力容易被低估。
Python免费学习笔记(深入):立即使用
在学习笔记中,你将探索 Python 的核心概念和高级技巧!










