1. 迁移前先想明白为什么大家都在从单卡迁到DDP先把话放前面如果你的模型一个卡能装下、训练速度能接受那真没必要折腾分布式。我见过不少同学参数就那么几百万非得上 DDP结果一堆通信开销比算力省下的时间还多纯属自找麻烦。但反过来模型一上亿参数、单卡显存吃紧、训练一轮要跑几个小时的时候多卡分布式就是刚需了。PyTorch 里最常见的多卡方案就是 DDPDistributedDataParallel而它最常见的迁移路径恰恰就是把手头已经调通了的单卡训练脚本改造成分布式版本。这篇文章我把整个迁移过程拆开讲从概念到代码到避坑尽量一套讲透保证你照着改能少走弯路。1.1 多卡方案摆在面前我为什么最后选了 DDPPyTorch 提供多卡训练的路线不止一条。最古老的是DataParallel新一些的是DistributedDataParallel再往上是 Horovod、DeepSpeed 这类第三方框架。对于从单卡脚本出发的人来说真正需要考虑的就是 DP 和 DDP 之间怎么选。DataParallel的用法确实简单包一层nn.DataParallel(model)就完事但它有个致命伤它是单进程多线程的所有卡的计算都塞在同一个进程里Python 的 GIL 加上频繁的 GPU 通信导致多卡利用率很难上去。而且它默认把输入在 batch 维度切分后分给各卡每张卡只拿到原始 batch 的一部分batch size 的语义都变了。更常见的问题是第一张卡主卡不仅负责汇总梯度还要把梯度广播回所有卡显存和带宽都卡在它身上经常出现0 号卡爆显存其他卡还有富余的尴尬局面。DDP 走的是完全不同的路线每个 GPU 一个独立进程各自持有模型副本前向反向独立算梯度用 Ring AllReduce 做一次全局同步最后每个进程更新出完全一致的参数。通信只在梯度同步那一下发生而且是分桶bucket进行效率比 DP 高一个量级。实测同样的 4 卡 ResNet-50 训练DP 可能只比单卡快 2 倍出头DDP 可以跑到 3.5 倍以上。所以结论很简单新项目、迁移老项目统一选 DDP。除非你只是想临时跑个实验、不想改任何代码那 DP 可以救急但别指望它当生产方案用。1.2 DDP 的梯度同步原理搞懂这个后面调参全通了DDP 名义上叫分布式数据并行核心逻辑拆开其实就四步每个进程独立初始化一份模型参数初始随机种子一致的话每份初始参数完全相同。各自从数据集的不同分片里取 batch前向传播、计算损失、反向传播算出梯度。反向传播过程中梯度被切成一个个 bucket每个 bucket 算完后立即通过 Ring AllReduce 在卡间求平均这也是 DDP 能异步通信、不阻塞计算的关键。每个进程拿到同步后的梯度各自调用 optimizer.step()因为初始参数和梯度都一致更新完参数仍然完全一致。这里有个常被忽略的细节梯度同步是求平均。也就是说如果你的原始单卡脚本没有对梯度做平均相关的处理迁移到 DDP 后等效的 batch size 会变成单卡 batch size × 卡数因为每张卡的梯度都除以了 world_size。这个对学习率的影响我后面专门讲很多人直接踩坑训练曲线一下子崩掉。理解了这套机制你就能明白为什么 DDP 要求每个进程模型必须一模一样也就能理解后面检查点保存、加载时为什么要小心处理 rank 的区分。原理永远是排错的第一把钥匙。2. 动手改代码之前先把这几个概念嚼烂迁移工作表面上是改代码实际上是补概念。我见过太多人把 DDP 的样板代码抄过去跑通了但一报错就抓瞎因为根本不知道每个参数是干嘛的。2.1 rank、world_size、local_rank 到底各管什么事这三个变量是所有分布式训练代码里最基础的东西必须在动手前弄明白。world_size参与训练的进程总数。如果你用 4 张卡每个卡一个进程那 world_size 就是 4。rank当前进程在全局的编号范围是 0 到 world_size - 1。这个编号决定了谁是主进程——一般来说 rank 0 负责保存 checkpoint、打印日志、做验证集评估。local_rank当前进程在本机的编号。单机多卡场景下 local_rank 和 rank 相等但一旦涉及多机训练local_rank 就必须单独区分。因为每台机器上都有 0、1、2、3 号卡如果不小心用了其他机器的 rank 去指定 CUDA 设备就会报设备不存在或者干脆访问了别人的显存。实际代码里最常见的对应关系就是torch.cuda.set_device(local_rank)然后用device torch.device(cuda, local_rank)来创建张量。注意不能直接写model.cuda()不带参数那样默认进 0 号卡4 个进程全挤在一块训练照样能跑但显存和效率都崩了。2.2 通信后端选对成功一半NCCL 还是 GLOOPyTorch 的init_process_group里有个backend参数最常用的两个值是nccl和gloo。NCCL 是英伟达专门为 GPU 通信优化的库实测在单机多卡、多机多卡场景下性能都远超 GLOORing AllReduce 的实现效率非常高。所以 GPU 训练直接无脑选 NCCL不需要犹豫。GLOO 则可以跑在 CPU 上、也可以走共享内存通信适合 CPU 分布式训练或者调试阶段用。有些同学在容器里跑 DDP 报找不到 NCCL多半是镜像没装 NCCL 或者版本不匹配这时候临时用 GLOO 可以先把流程调通但正式训练必须修好 NCCL否则速度没法看。对比项NCCLGLOO适用设备NVIDIA GPUCPU、GPU 均可单机多卡性能极佳一般多机扩展性好较差典型用途正式训练调试、CPU 环境另外还有一个init_method的问题默认用env://就行配合 torchrun 自动把环境变量注入到进程里基本不用手写。后面启动方式那节我会展开讲。2.3 DistributedSampler 不换你的每个 epoch 都在用同一批数据训练单卡脚本的 DataLoader 通常写成DataLoader(dataset, batch_size32, shuffleTrue)但在 DDP 下直接这么干会出大问题每个进程都独立 shuffle数据重叠严重不同卡可能在同一个 step 训练几乎相同的 batch梯度高度相关模型泛化能力直接下降而且每张卡都在全量数据上遍历一个 epoch 的实际步数还是全量数据量除以单卡 batch size训练效率也翻不了倍。正确的做法是把shuffleTrue去掉改用DistributedSamplerfrom torch.utils.data.distributed import DistributedSampler sampler DistributedSampler(dataset, num_replicasworld_size, rankrank, shuffleTrue) dataloader DataLoader(dataset, batch_sizebatch_size, samplersampler, num_workers4)DistributedSampler做的事情本质上是把完整数据集平均切成 world_size 份每个进程拿属于自己的那一份并且通过一个随机的 epoch 种子保证同一个 epoch 内每张卡的数据不重叠不同 epoch 之间每个进程拿到的分片又是重新打乱的。这里有个必须记住的细节DistributedSampler依赖set_epoch()来更新 shuffle 的随机种子。每轮 epoch 开始前都要调用一次for epoch in range(epochs): sampler.set_epoch(epoch) for batch in dataloader: train_step(batch)漏了这一步每轮 epoch 数据排列完全一样模型在小数据集上很快会过拟合而且你很难意识到是这个原因造成的。3. 从单卡到 DDP一份能直接抄的改造实录下面这部分我按真实项目的改造顺序来走把一段典型的单卡训练脚本逐步改成 DDP 版本。为了有代入感就拿一个常见的图像分类训练脚本举例结构是建模型 → 建 DataLoader → 循环训练 → 保存 checkpoint。3.1 第一步进程组初始化和设备绑定单卡脚本的开头一般是读参数、设随机种子、device torch.device(cuda if torch.cuda.is_available() else cpu)。改成 DDP 后这部分要替换为进程组初始化import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def setup(backendnccl): dist.init_process_group(backendbackend) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) return local_rank def cleanup(): dist.destroy_process_group()注意两点。第一init_process_group如果什么都不传它默认从环境变量里读MASTER_ADDR、MASTER_PORT、RANK、WORLD_SIZE而这些环境变量在 torchrun 启动时已经自动设置好了所以这行代码写起来异常简洁。第二torch.cuda.set_device(local_rank)这一步不能省它决定了当前进程的默认 GPU 是哪张卡后续所有.cuda()调用都会落到这个设备上。随机种子部分建议每张卡用rank做个偏移避免所有进程同步后仍然产生相同的随机序列虽然 DDP 对参数一致性要求严格的场景下要保证初始参数一样但训练过程中 Dropout 这类操作的随机性是独立的加偏移更合理。3.2 第二步模型包装成 DDP单卡代码里是model MyModel().to(device)DDP 下改成model MyModel().to(device) model DDP(model, device_ids[local_rank], output_devicelocal_rank)包装完之后使用方式几乎不变outputs model(inputs)照常调用model.module才能拿到原始模型。这句话意味着两处代码要跟着改保存 checkpoint 时model.state_dict()要写model.module.state_dict()否则 key 前面会多出module.前缀单卡加载时匹配不上。如果要打印模型结构、访问某个层做可视化也要先model.module拿到原始模型。device_ids参数在单机多卡时一定要显式传它在后续版本里如果没传可能会报错或者默认行为不符合预期。另外还有一个find_unused_parametersTrue的参数如果模型里有参数没参与 loss 计算比如某些中间层输出没有被使用DDP 的梯度同步会卡住加了这个参数才能跑通但它会牺牲一定性能能不加就不加。注意DDP 包装之后模型参数更新的职责仍然由你在外面写的优化器负责。优化器要基于model.parameters()来构建而不是model.module.parameters()。两种写法拿到的参数对象是同一批所以大多数情况下不报错但我见过有人混用导致 optimizer 更新了旧对象、模型权重纹丝不动的诡异情况统一用model.parameters()就好。3.3 第三步DataLoader 改造和 batch size 调整这一步最容易抄错。单卡代码dataloader DataLoader(train_set, batch_size64, shuffleTrue, num_workers8)DDP 代码sampler DistributedSampler(train_set, shuffleTrue) dataloader DataLoader(train_set, batch_size64, samplersampler, num_workers8)这里 batch_size 指的是每张卡的 batch size。也就是说4 卡训练时一次全局更新吃进去的真实样本数是64 × 4 256。有些项目的习惯是配置里写总 batch size然后代码里除以 world_size 得到单卡 batch size两种都行关键是团队内部约定清楚。训练循环里别忘了sampler.set_epoch(epoch)这个我在前面已经强调过了。另外验证集的 DataLoader 不需要 DistributedSampler保持普通 DataLoader 即可但要注意验证通常只在 rank 0 进程上做避免 4 个进程同时评估打日志刷屏也避免验证带来重复计算。3.4 第四步checkpoint 的保存与加载坑最多的地方单卡脚本里保存 checkpoint 往往就是torch.save({model: model.state_dict(), optimizer: optimizer.state_dict()}, ckpt.pth)DDP 这么写会出两个问题。第一4 个进程同时写同一个文件互相覆盖轻则文件损坏重则报错。第二model.state_dict()里多了一层module.前缀加载回单卡模型需要改 key。正确的保存方式是只在 rank 0 上执行if dist.get_rank() 0: torch.save({ model: model.module.state_dict(), optimizer: optimizer.state_dict(), epoch: epoch, }, ckpt.pth)加载的时候也很有讲究。DDP 内加载和单卡加载情况不一样checkpoint torch.load(ckpt.pth, map_locationdevice) # device 是当前进程的 cuda:local_rank model.module.load_state_dict(checkpoint[model])注意map_locationdevice这一步。如果直接不加 map_locationtorch.load 会默认把张量读到进程启动时的默认设备多卡环境下经常读到了 cuda:0加载后模型权重所在设备混乱跑前向时报expected all tensors on same device。这也是我从实际项目里踩出来的教训。还有一个进阶技巧如果训练中断要断点续跑需要把 optimizer 和 sampler 的 epoch 状态也一起恢复并且恢复后要正确设置sampler.set_epoch(epoch)否则数据顺序和中断前不衔接。DDP 场景下每个进程都各自加载同一份 checkpoint但因为各进程的数据分片不同续跑后仍然能保持正确的数据分配。3.5 第五步启动方式torchrun 是最省心的答案早期写 DDP 脚本的人喜欢用torch.multiprocessing.spawn手动拉起多进程代码长而且容易出错。PyTorch 1.10 之后的官方推荐是直接用torchruntorchrun --nproc_per_node4 --master_port29500 train_ddp.py--nproc_per_node表示每台机器上启动多少个进程通常就是 GPU 数量。--master_port是通信端口如果默认的 29500 被占用了换一个没被占的就行。torchrun 会帮我们做的事情包括设置RANK、LOCAL_RANK、WORLD_SIZE、MASTER_ADDR、MASTER_PORT环境变量然后把进程有效地分配到对应 GPU 上运行。也就是说你的脚本里不要再手写os.environ[CUDA_VISIBLE_DEVICES]这类逻辑靠LOCAL_RANK定位设备就好。如果是多机训练比如两台机器各 4 卡启动命令变成# 节点 0 torchrun --nnodes2 --nproc_per_node4 --master_addr192.168.1.10 --master_port29500 train_ddp.py # 节点 1 torchrun --nnodes2 --nproc_per_node4 --master_addr192.168.1.10 --master_port29500 train_ddp.py两台机器的命令几乎一样唯一的区别是让 torchrun 自己识别节点编号还是手动加--node_rank。最简单的办法是保证两台机器都能访问同一个共享存储然后在每台机器上顺序启动同一个命令torchrun 会自动分配 node rank。我个人的建议是哪怕你不想用 torchrun也尽量把代码写成兼容环境变量驱动的方式这样不管是手动spawn、Slurm 调度还是云平台容器都能无缝切换。4. 跑起来之后问题清单与排查技巧实录代码改完、启动成功只是第一步真正让人头大的是后面各种「能跑但不对劲」的情况。这里把我实际遇到过的、以及帮别人排查过的高频问题整理成清单每一条都是真金白银踩出来的。4.1 通信超时、端口占用和进程挂起不退出最典型的报错是RuntimeError: The server socket has failed to listen on any local network address原因几乎都是MASTER_PORT被占或者防火墙拦了通信端口。换个端口试试是第一步。另一个常见现象是训练卡在init_process_group一直不出来这种十有八九是某个节点的进程起不来导致所有节点都在等 handshake。排查方法简单粗暴把所有进程杀掉确认各节点能互相 ping 通、端口能连通再重新启动。还有一个隐蔽问题训练正常结束但脚本就是不退出像死锁一样卡在最后。这通常是某个子进程没有正常销毁常见于 DataLoader 开了太多 num_workers 或使用了自定义的 multiprocessing 逻辑。目前比较推荐的做法是训练完后显式调用cleanup()销毁进程组并把主流程放在if __name__ __main__里。4.2 BN 层在多卡下的坑同步 BN 究竟要不要用单卡模型里 BN 层的统计量均值和方差是在一个 batch 里算的。DDP 下每个进程只看到 1/world_size 的数据如果单卡 batch size 本身就小比如 8、16BN 统计量会非常不稳模型的收敛效果可能明显变差。解决办法是把普通 BN 换成SyncBatchNorm让所有卡一起计算 BN 的均值和方差。PyTorch 里提供了torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)一行代码完成转换不需要改模型结构。但要注意SyncBatchNorm 会增加通信开销图像分类这种任务通常不明显但目标检测、分割这类本身就是多卡打 batch 的项目要权衡一下全局同步的通信量别让小 batch 问题变成性能瓶颈。我的经验是如果单卡能支撑 batch size 32 以上一般不需要 SyncBatchNorm如果卡多且每卡 batch 小于 16果断转 SyncBatchNorm收敛质量优先。4.3 学习率不调训练曲线直接爆炸前面原理部分提过DDP 的梯度是全体卡求平均后各自更新的等效 batch size 变大了 4 倍。按照线性缩放规则学习率也应当近似放大sqrt(world_size)或world_size倍具体取哪个看你的优化器和数据分布。举个例子单卡 batch size 64 时你用的是 lr0.14 卡训练 batch size 变成 256那么 lr 可以尝试 0.2 到 0.4 的范围。但线性缩放不是无脑放大它有一个隐性前提训练总步数不变的话放大学习率会改变整个优化轨迹所以实践中更稳妥的做法是结合 warmup 一起调前几个 epoch 从小学习率逐步升到目标值。我因为懒曾经把单卡的 lr 直接搬到 4 卡 DDP 上结果 loss 曲线前几个 step 就开始震荡发散白白浪费了一天时间。从那之后我把改 DDP 必调学习率写进了自己的 checklist。4.4 一套代码同时兼容单卡和多卡日常开发中你不可能永远在多卡环境调试。更靠谱的做法是让脚本同时支持单卡模式和DDP 模式通过一个--distributed参数切换if args.distributed: dist.init_process_group(backendnccl) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) device torch.device(cuda, local_rank) model DDP(model, device_ids[local_rank]) else: device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device)DataLoader 那边同理distributed模式下才包 DistributedSampler否则保持原来的 shuffle 逻辑。这会让代码里多一些 if 分支但换来的是本地单卡调试、线上多卡训练完全走同一份代码省掉很多在两个版本间同步修改的灾难。4.5 多卡比单卡还慢问题出在哪DDP 加了几张卡结果速度不升反降。我排查过这类问题常见的元凶按照出现频率排序如下原因表现解决方案num_workers 太少数据加载跟不上 GPU 计算每卡单独调大 num_workers4~8 起步小 batch 高频通信梯度太小通信次数太多适当加大 batch size或用梯度累积模型太小计算量远低于通信开销DDP 只在大模型上收益明显每步有大量小张量bucket 机制失效尽量复用缓存、减少动态分配验证/日志在主进程串行每 epoch 卡顿明显验证时用同步锁减少打印机器间网络带宽不足多机训练明显掉速换 InfiniBand或压缩梯度通信这里最值得展开的是第一条。很多人从单卡搬过来之后num_workers 还保持原来的 4但多卡时每张卡同时需要数据进程之间还会争抢 CPU 资源数据供给立刻成为瓶颈。我通常的做法是num_workers按每卡 4~8 个来设并且用prefetch_factor4多做几批预取能让 GPU 的空闲时间明显减少。最后的经验之谈这套迁移流程我带过不少同学走过大多数人的节奏是周末看文档、周一改代码、周三踩坑、周五跑通。如果你今天就要动手我建议按这个顺序来先把 torchrun 的最小示例跑通确认多进程能起来、rank 能对上再改 DataLoader 和 sampler最后加 checkpoint 兼容逻辑。一次别贪多逐步验证比一次性大改完再排错要快得多。还有一个小建议值得单独提训练出现问题的时候先打印每个进程的 rank 和设备编号确认进程分配没有重叠再往下查。分布式训练最大的麻烦在于错误的症状和错误的原因往往隔着好几层而所有排查的第一步永远是先确认每个进程都运行在正确的位置上。