dataparallel导致gpu负载不均的根本原因是其主从式设计：cuda:0承担梯度汇总、loss计算、输出拼接和参数更新，而其他gpu仅执行局部前向/反向，造成cuda:0显存高20%~30%、利用率持续95%+，其余卡常低于40%。
为什么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”。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
