我一开始真没觉得PyTorch多GPU分布式训练是个非得动代码的事。单卡训练脚本跑得好好的无非就是慢一点。直到batch size顶到显存上限、一个epoch要27分钟、完整实验要一周的时候我才意识到不迁不行。第一版偷懒用了nn.DataParallel结果四张卡加起来只比单卡快一点点GPU利用率还在那儿一卡有难三卡围观。后来老老实实改成DDPDistributedDataParallel训练时间直接从按天算变成按小时算。这篇就把我这次从单卡脚本迁移到DDP的完整过程、改过的每一处代码、踩过的坑都记录下来。DDP这套东西网上教程很多但大多数只贴代码片段很少讲“为什么这么改”和“出问题时怎么查”。我按自己的迁移链路写从选型、改造、踩坑、调优到兼容单卡尽量讲得实在一点给正打算迁多卡的人一个可以参考的路线。1. 迁移前先算清楚账DDP解决的是哪类痛点1.1 单卡训练的瓶颈到底卡在显存还是算力我一开始的想法比较简单训练慢加卡不就快了但真的动手之前得先搞明白瓶颈在哪。如果只是算力不够加卡确实能缩短墙钟时间。比如单卡一个epoch 27分钟四卡理论上能压到七八分钟。但如果瓶颈在显存——模型太大、batch size提不上去那加卡并不能让你单卡省显存。DDP不会让一张卡装下一个原本装不下的算子它只是把一个大batch拆到多张卡上并行算。另一个容易被忽略的问题是一旦用了多卡数据加载、CPU预处理、验证循环、checkpoint写入这些外围逻辑都可能变成新瓶颈。我第一次迁完DDP4卡训练速度反而更慢了查了半天发现是num_workers没调磁盘IO被打满每张卡都在等数据。所以迁移前先确认你的单卡脚本已经把数据加载和预处理优化得差不多了吗如果没有多卡只会放大这个问题。1.2 为什么DataParallel看着省事实际却跑不快很多人和我一样第一反应是换nn.DataParallel因为它改一行就行。但DP和DDP的底层逻辑完全不同对比项DataParallelDistributedDataParallel进程模型单进程多线程受GIL限制多进程每个进程独立解释器卡的使用每次forward时把输入切分到多卡再收集结果每进程一张卡各算各的梯度阶段通信梯度同步每次forward后由主卡做一次全量梯度收集梯度算完通过NCCL做all-reduce与计算重叠扩展能力只能单机单机多卡、多机多卡都能用代码成本包一层就行需要初始化进程组、改数据samplerDP慢的核心原因有两个一是单进程多线程在PyTorch的Python侧会受到GIL影响并行度打折扣二是它每个step都得把输入和输出在主卡与其他卡之间来回搬运小batch还好batch一大这部分内存拷贝和同步开销非常明显。DDP则是每个进程独立持有模型副本forward和backward都在本地完成只在反向后把梯度做一次all-reduce配合NCCL的ring-allreduce实现通信效率高很多。所以结论很明确如果只是想在单机多卡上跑实验直接上DDP别在DP上浪费时间。后面所有内容都以DDP为例。1.3 DDP“解决不了”的三类问题得提前有预期多卡不是万能药我踩过一轮之后理出三类DDP解决不了的事提前说出来免得后面期望落空。第一单卡显存不够时DDP救不了你。模型本体的参数、激活值还是得装进单卡显存。想跑更大模型得上梯度检查点、模型并行、CPU offload或者干脆换更大显存的卡。第二小模型或者单step计算量太小时多卡可能更慢。DDP每次backward都要做一次梯度全局同步如果模型很小计算几百毫秒、通信也要几百毫秒那通信开销直接吃掉并行收益。我有一个极小的分类网络4卡跑出来和单卡几乎没差就是这个原因。第三数据加载、验证、日志、checkpoint这些外围环节多卡会把它们放大成新瓶颈。分布式训练不是只把model包一层就完事数据管线和验证逻辑都得跟着改。2. 单卡脚本到DDP的改动量初始化、模型、数据、checkpoint一次说清2.1 先定启动方式torchrun是绕不开的第一步单卡脚本通常是python train.py改成DDP后推荐用torchrun启动torchrun --nproc_per_node4 train.py--nproc_per_node4表示每台节点起4个进程每个进程负责一张卡。如果只想先用单卡调试DDP路径写--nproc_per_node1即可。torchrun会自动把RANK、LOCAL_RANK、WORLD_SIZE这些环境变量填好。我在代码里统一取import os import torch import torch.distributed as dist local_rank int(os.environ[LOCAL_RANK]) rank int(os.environ[RANK]) world_size int(os.environ[WORLD_SIZE])RANK是全局进程编号LOCAL_RANK是当前机器上的进程编号单机时两者一致WORLD_SIZE是总进程数也就是总卡数。注意老教程里的python -m torch.distributed.launch已经被官方标记为废弃新项目直接用torchrun。2.2 初始化进程组必须在任何CUDA操作之前在main函数最前面加上这段dist.init_process_group(backendnccl, init_methodenv://) torch.cuda.set_device(local_rank) device torch.device(cuda, local_rank)backendnccl是NCCLCUDA环境下就选它通信效率最高。init_methodenv://表示从环境变量读取分布式配置正好配合torchrun。这里有个我见过非常多的坑有人不调torch.cuda.set_device(local_rank)代码里直接用model.cuda()。这样所有进程都会默认落到cuda:0轻则显存爆炸重则直接报错。所以必须在初始化后马上set_device后续模型、数据、输出全部用device这个变量。另外多进程跑起来之后每张卡的数据分布不同随机种子也必须不同。我一般在初始化后做seed 42 torch.manual_seed(seed rank)否则多个进程用同一个随机种子初始化dropout、初始化参数训练结果会异常地“一致”实际上是错误的。2.3 模型、数据加载器、优化器的三处核心改动模型包装是最直观的一处model UNet(in_channels3, num_classes21).to(device) model torch.nn.parallel.DistributedDataParallel( model, device_ids[local_rank], output_devicelocal_rank )device_ids[local_rank]必须写让DDP知道这个进程对应哪张卡。output_device一般和device_ids一致即可。数据加载器这部分容易被忽略但它决定了多卡训练是否真的“各算各的”from torch.utils.data.distributed import DistributedSampler dataset MyDataset(...) sampler DistributedSampler( dataset, num_replicasworld_size, rankrank, shuffleTrue, seed42 ) dataloader DataLoader( dataset, batch_sizebatch_per_gpu, samplersampler, num_workers4, pin_memoryTrue, drop_lastTrue )关键点用了DistributedSampler之后DataLoader里就不能再写shuffleTrue因为每个epoch是否shuffle、怎么shuffle已经由sampler决定。如果两边都开shuffle多个进程看到的数据会完全乱掉训练一个epoch等于重复采样验证也说不清到底看了哪些样本。每个epoch开始时还要显式调用一下sampler.set_epoch(epoch)for epoch in range(start_epoch, total_epochs): sampler.set_epoch(epoch)不调用会怎样每个epoch所有进程拿到的数据顺序都一样或者说shuffle只发生一次训练效果会受影响。优化器本身不用特殊处理它还在DDP外面optimizer torch.optim.AdamW(model.parameters(), lr1e-4)DDP会同步梯度并不会接管优化器状态。关于batch size怎么设我给自己定了一条规则先看单卡显存能承受多大batch再决定整体策略。训练目标每卡batch设置global batch学习率调整保持原来的global batch不变原batch size除以world_size与原脚本一致基本不用动保持每卡batch不变扩大总batch沿用原batch size翻world_size倍按线性缩放规则适当调大我这次的任务是图像分割单卡最大batch只有4所以选择了扩大global batch每卡还是4四卡总batch变成16。学习率从1e-4调到2e-4只加了点warmup效果和单卡持平。2.4 checkpoint保存和加载rank 0说了算别的卡别乱动多卡环境下checkpoint保存最怕两件事一是每张卡各存一份互相覆盖二是保存出来的模型在单卡上加载不了。我的做法是只在rank 0上保存if dist.get_rank() 0: torch.save({ model: model.module.state_dict(), optimizer: optimizer.state_dict(), epoch: epoch 1, scaler: scaler.state_dict() if scaler else None, }, checkpoint_path)注意我用的是model.module.state_dict()而不是model.state_dict()。DistributedDataParallel包装后的模型它的state_dict里所有key会多一个module.前缀。保存时不加处理后面自己加载或者给单卡用都会多一层麻烦。统一保存model.module.state_dict()这样单卡和多卡加载时都干净。加载是另一个容易被坑的地方。我推荐的做法是先用裸模型加载参数再包DDPmodel UNet(...).to(device) checkpoint torch.load(checkpoint_path, map_locationcpu) model.load_state_dict(checkpoint[model]) model torch.nn.parallel.DistributedDataParallel(model, device_ids[local_rank])如果是在多卡训练中断后恢复还需要手动把epoch和sampler状态对齐每个进程都以checkpoint里的epoch为起点加载optimizer然后sampler.set_epoch(epoch)。否则不同进程可能从不同epoch的数据开始batch顺序错位。2.5 训练循环里的日志与指标只打一份但要聚合直接把单卡脚本里的print(fepoch {epoch} loss {loss.item()})原样搬过来多卡跑起来会输出四份日志而且每份都只是当前进程那部分数据的loss不能代表全局。我习惯这样处理日志只在rank 0打印指标用all_reduce做全局聚合。if dist.get_rank() 0: print(fepoch {epoch} iter {i} loss {loss_avg:.4f})聚合loss可以写个通用函数def reduce_mean(tensor, world_size): rt tensor.clone() dist.all_reduce(rt, opdist.ReduceOp.SUM) return rt / world_size调用时传入当前进程的loss返回全局平均loss。这个函数不只用于loss验证集的acc、miou这类标量指标同样适用。要注意的是all_reduce本身是一个同步操作如果某个进程没走到这句其他进程会一直等所以聚合调用点必须保证所有rank都执行。3. 迁移过程避坑实录我实际复现过的问题与定位思路3.1 启动即报“Address already in use”的真相第一次用torchrun跑完一次训练马上再启动同一个命令报错“Address already in use”。我当时以为是随机选的端口冲突后来发现是上一次训练进程没退干净。torchrun默认会在本机找一个空闲端口当master_port如果上次进程还没死占着端口不放第二次启动自然失败。排查方法很简单ps -ef | grep train.py看到残留进程直接杀掉或者干脆每次启动前用pkill清理pkill -f train.py如果多个人共用一台机器端口冲突会更频繁可以在启动时指定一个不常用的端口torchrun --nproc_per_node4 --master_port29501 train.py还有一次遇到进程完全hang住没有任何报错加上NCCL_DEBUGINFO才看到是NCCL初始化时某个进程连不上peer。这种情况通常和上次训练留下的共享内存段或者P2P通信异常有关先清理进程必要时export NCCL_P2P_DISABLE1试试但这是下策治标不治本根因一般还是环境问题。3.2 训练完整体loss对不上但代码没报错sampler的锅有一次我把脚本改成DDP后loss曲线很好看收敛得比单卡还快但验证集效果一塌糊涂。排查到数据管线才发现我虽然加了DistributedSampler但DataLoader里原来的shuffleTrue没删。两个shuffle叠加sampler先把自己那份数据洗了一遍DataLoader又洗了一遍多个进程实际取到的是错乱数据。这个问题的隐蔽性在于它不报错loss还在降不仔细对比验证集根本发现不了。所以我后来加了一行自检每个进程打印自己拿到的数据集长度和第一个batch的索引范围肉眼确认一下是不是约等于len(dataset) / world_size。如果每个进程打印出来的都是完整长度说明sampler没生效就得回头查代码。3.3 DDP训练卡死部分参数不被更新的模型会让人崩溃我的分割网络是双分支结构其中一个分支只在有标注时才会计算loss。单卡上完全正常但迁到DDP后训练到一半会hang住没有任何报错日志停在某个iterGPU利用率掉到0。后来我确认了原因DDP在backward阶段需要所有进程同步梯度但我的双分支模型里有些参数在当前batch中并没有参与前向计算也就是没有梯度产生。不同进程拿到的数据不同有的进程这个分支有loss有的没有于是DDP等待梯度同步时永远凑不齐。解决方案是给DDP加一个参数model torch.nn.parallel.DistributedDataParallel( model, device_ids[local_rank], find_unused_parametersTrue )find_unused_parametersTrue会让DDP先做一次额外的遍历找出哪些参数没参与loss计算然后跳过它们的梯度同步。代价是每个step都会多一点额外开销训练会慢一些。更优的解法是从模型结构上避免这种不稳定分支但实际业务里结构往往不能随便改所以这个参数是必须记住的保底方案。3.4 CUDA device-side assert triggered多卡把偶发坏样本暴露了这个问题最玄学同样一份代码单卡怎么跑都没事四卡DDP跑到第137个iter直接崩报错CUDA error: device-side assert triggered而且每次崩的iter还不一样。我用CUDA_LAUNCH_BLOCKING1重新跑定位到是Embedding或者CrossEntropyLoss内部触发。最后查出来是数据集中有极少数样本的label超出了类别数。单卡脚本里因为数据顺序固定这些坏样本恰好排在训练集末尾而我的实验只跑了前面的epoch从来没取到过换到多卡后DistributedSampler重新切分、shuffle坏样本被提前采样到了。这个坑给了一个教训多卡训练会把单卡下被“运气”掩盖的坏样本暴露出来。迁移DDP前最好先对数据集做一次全量扫描校验label范围、尺寸、是否为空。排查这类问题除了CUDA_LAUNCH_BLOCKING1还可以在dataloader里加一个前向前的断言打印出当前样本的index快速定位问题样本。3.5 多进程print输出乱序不是环境问题是定位问题多卡跑的另一个小麻烦是日志全混在一起。多个进程同时往stdout写内容行与行之间互相穿插基本没法看。我后来定了一个规范def log(msg, rank0): if dist.get_rank() rank: print(f[{time.time():.3f}] {msg})所有关键节点都带上rank前缀虽然不能完全消除交错但至少能看出是哪个进程在什么时候打的。真正调试时我会只保留rank 0的日志其他进程的日志暂时注释掉否则信息量太大反而影响判断。4. 迁移后的性能调优通信、batch大小与混合精度怎么配合4.1 四卡加速比不到3倍问题多半出在通信我的U-Net第一次在4卡上跑起来测了一下每秒处理的样本数只有单卡的2.3倍和理想4倍差了不少。用torch.profiler抓了一个step的timeline发现ncclKernel_AllReduce占了接近30%的时间。DDP每次backward后都要同步所有梯度通信量跟模型参数个数成正比跟batch size没直接关系。模型越大、batch越小通信开销占比越高。如果加速比不理想首先是确认模型够不够大、单step计算量够不够多。我这边把每卡batch从4提到8之后用了AMP全局batch翻倍通信占比明显下降加速比到了3.2倍。如果batch提不上去还有一个办法梯度累积加no_sync。正常情况下DDP每个step都做一次all-reduce但梯度累积时本来就要攒多个step才更新一次中间这些step的通信可以被省掉optimizer.zero_grad() accum_steps 4 for i, (x, y) in enumerate(dataloader): with model.no_sync(): loss criterion(model(x), y) loss.backward() if (i 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()注意这里model.no_sync()只包住了反向传播真正要更新时才走DDP的梯度同步。用这种方式通信次数直接除以accum_steps代价是语义上等于把accum_steps x per_gpu_batch合并成一个更大的batch学习率、BN统计都要跟着重新考虑。BN方面DDP默认不会跨卡同步BN统计量no_sync时更不会。如果模型里有BN且per-gpu batch比较小强烈建议用SyncBatchNormmodel torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)在我的分割任务里per-gpu batch只有4BN统计漂移明显换成SyncBatchNorm后验证miou提升了约1.2个点每step代价是增加约5%的通信时间总体来说划算。4.2 AMP混合精度不是简单加两行就完事DDPAMP是我常用的组合。在显存吃紧的情况下AMP几乎白送30%~50%的吞吐提升。标准写法scaler torch.cuda.amp.GradScaler() for x, y in dataloader: optimizer.zero_grad() with torch.cuda.amp.autocast(): loss criterion(model(x), y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这里有一个很多人问过的问题DDP会同步梯度AMP的loss scaler是每个进程独立的梯度会不会对不上我的理解是loss在backward前已经被scale过所以梯度也是scaled后的值DDP同步的是这些scaled梯度之后每个进程再各自做unscale_。只要每个进程的scaler初始状态一致它们本来就一样且随机种子按rank区分但不影响scaler更新逻辑整个流程是自洽的。我实测下来没有因为AMPDDP出现梯度不一致。真正要注意的是checkpoint里要保存scaler.state_dict()并且在恢复训练时加载。否则训练到一半挂掉再恢复scaler的信息丢了得重新从初始scale开始前半程的scale自适应基本白做。还有就是如果发现loss变成NaN别急着归咎于AMP。先在单卡DDP模式下跑一遍排除数据问题再用torch.cuda.amp.autocast()只在前向里对部分模块开启逐段排查。经验上NaN多数来自fp16精度下的梯度下溢可以适当调大GradScaler的init_scale或者把min_loss_scale调大一点避免scale降得太低。4.3 梯度累积、全局batch、学习率的关系要一起调调优时最怕“只改一个参数”。DDP把batch翻倍了学习率不动可能不收敛或者收敛慢梯度累积又改变了更新频率学习率还按原来的step schedule走也可能会出事。我给自己定了一套规则先确定effective batch sizeper_gpu_batch x world_size x accum_steps。如果effective batch和单卡时的batch相当学习率直接沿用单卡的。如果effective batch变大了学习率按线性缩放新的lr大概等于原lr乘以(new_batch / old_batch)但上限通常不超过2倍同时需要加长warmup。最简单的验证方法用相同种子分别跑单卡和DDP对比前几个epoch的loss曲线。如果DDP的loss曲线和单卡基本重合配置就是对的如果loss掉得太快或者完全不动就优先怀疑学习率和batch策略。5. 从多卡回到单卡脚本的兼容设计与多机扩展5.1 用配置控制训练模式少写分支迁移完DDP后我还需要保留单卡训练的能力因为调试小问题时单卡更快、更容易复现。为此我把脚本改成了可配置的。单卡时也可以直接走DDP流程只要torchrun --nproc_per_node1 train.pyworld_size为1DDP退化成单进程所有逻辑依然成立。这样代码里不需要到处写if use_ddp分支少心智负担小。唯一的麻烦是checkpoint的保存。我写了一个通用函数def unwrap_model_state(model): if isinstance(model, torch.nn.parallel.DistributedDataParallel): return model.module.state_dict() return model.state_dict()保存和加载都用它保证单卡、多卡之间checkpoint通用。5.2 加载旧权重时的module前缀问题同事给过一个在4卡上训出来的权重load到单卡模型里直接报“missing key: xxx”和“unexpected key: module.xxx”。这种情况太常见了我写了一个剥离前缀的小工具def strip_module_prefix(state_dict): new_state_dict {} for k, v in state_dict.items(): if k.startswith(module.): new_state_dict[k[len(module.):]] v else: new_state_dict[k] v return new_state_dict反向的添加前缀函数也留着万一单卡权重要喂给多卡模型也能用。虽然用load_state_dict(strictFalse)能跳过这些报错但这样会漏加载参数训练出来结果不对还不好查。宁可多写几行把key对齐。5.3 多机多卡比单机多出来的几步在多机环境下每台机器都要执行同样的启动命令但需要指定一个master节点# 在每台机器上执行node0换成第一台机器的IP export MASTER_ADDR192.168.1.10 export MASTER_PORT29500 torchrun --nproc_per_node8 train.pyRANK和LOCAL_RANK由torchrun自动填但所有节点的代码要能读到同一份数据集和配置。我一般把数据集放在共享存储上路径保持一致。多机比单机更容易受网络影响。如果节点之间走万兆以太网而没用RDMANCCL的all-reduce延迟会高不少。这种情况下优先把batch调大让通信占比下降其次考虑减少通信次数比如梯度累积。还有一个容易被忽略的点各节点GPU型号尽量一致。不同型号的卡算力差异会让整个训练被最慢的卡拖住严重时甚至出现NCCL超时报错。5.4 我推荐的迁移顺序别一上来就四卡如果让我给第一次接触DDP的人一个建议就是别直接上四卡。我自己第一次直接迁到4卡报错排查成本非常高因为你不知道是代码问题、环境问题还是数据问题。我后来固定用这套顺序确保单卡脚本能完整跑通loss正常下降。用torchrun --nproc_per_node1 train.py跑一遍DDP路径确认单进程模式下所有改造都正确。上两卡跑打印rank、world_size、sampler长度三件套确认数据切分正确对比单卡loss曲线。再加AMP、checkpoint保存与恢复。最后再上四卡看加速比做通信和batch调优。单机稳定后再考虑多机。这套顺序帮我少踩了很多无谓的坑。DDP的报错链路本来就长一次只引入一个新变量定位问题会容易得多。多GPU这件事本质上没有什么魔法它就是把“单卡跑一个任务”拆成“多进程各自跑一个小任务再在关键节点同步”。把启动方式、数据切分、checkpoint、指标聚合这几件事理顺了迁移也就完成了一大半。