把模型从单卡搬到多卡最怕的不是代码写不出来而是训练跑起来之后效率反而比单卡还差甚至直接卡死不动。PyTorch分布式训练里最常用的主力工具就是torch.nn.parallel.DistributedDataParallel也就是大家口中的DDP。无论你是第一次尝试把模型扩展到两张卡还是准备在一台8卡机器上跑较大规模并行DDP基本都是绕不开的主路径。这篇是分布式系列的第三篇我会把DDP的原理、启动方式、采样规则、checkpoint策略、常见坑和调参路径一次讲透。适合已经用PyTorch写过完整训练循环、但对多卡运行还停留在“听说”阶段的读者也适合那些跑通了DDP但经常遇到卡死、梯度过不去、checkpoint加载出错的兄弟。1. 单卡到多卡分布式训练解决的问题到底是什么1.1 显存墙与算力墙单卡训练的天花板很多初学者以为分布式训练的核心目的就是“显存不够了多卡拼一下”。这个说法只对了一半。显存确实是第一道墙一张消费级显卡24GB已经是比较常见的容量大模型或者大batch训练时光是一个batch的数据加激活值就可能把显存吃穿更别提模型本身还占着一大块。但就算你的模型小到能塞进单卡训练时间依然是硬伤。一个需要跑十几个小时甚至几天的训练任务单卡只能慢慢等而多卡并行可以把同样规模的数据切成多份让每张卡同时干活吞吐量翻上去。所以说DDP解决的是两件事第一是装不下第二是等不起。装不下的时候DDP靠数据并行把不同batch数据分到不同卡上模型本身在每张卡都有一份完整副本所以模型显存这部分没有省下来省的是单卡上同时处理的数据量和中间激活内存。等不起的时候DDP靠多卡同时计算梯度、同时更新参数训练总时间近似与卡数成反比。1.2 DDP要解决的三类核心问题模型同步、数据切分、进程协调DDP真正要处理的不是“把数据分下去”这么简单。数据切分只是第一步后面还跟着三件麻烦事模型同步每张卡上都有一份独立的模型参数。训练刚开始时它们是一样的但反向传播计算出的梯度依赖于各自的数据所以各卡的梯度天然不同。想要让模型继续保持“多卡算出来的更新方向一致”就必须把梯度在所有卡之间做一次集体通信让每张卡拿到全局平均后的梯度再用这个梯度去更新自己的参数副本。数据切分多卡训练时每个进程不能拿到完整的batch也不能随便撞车否则整体训练就相当于几倍batch混在一起还互相重复。DDP配合DistributedSampler把整个数据集切分成互不重叠的若干份分别交给不同进程。进程协调多进程不是多个线程它们不共享内存也没有全局锁。进程之间需要约定谁是master、谁是worker、谁先启动谁等待还要初始化通信后端。这些协调工作都靠torch.distributed里的进程组完成。打个比方单卡训练是一个人从头到尾做题数据并行训练是一组人分别做不同的题目但每做完一部分就要停下来对答案如果其中一个人的答案和其他人差距太大所有人都要按平均后的答案来调整自己的思路。DDP就是那个组织“对答案”的角色。1.3 为什么选DDP而不是老牌DataParallelPyTorch里实现数据并行的工具其实有两个方向。老一点的torch.nn.DataParallelDP用一个进程管理多张卡输入batch被切成几块分发给各卡前向结果回收到主卡主卡计算loss后反向传播然后把更新后的参数广播给各卡。这个方案最大的问题是主卡承担了太多活儿分发数据、接收输出、计算梯度、广播参数全都压在主卡上所以经常出现主卡显存比其他卡高一大截的情况而且受限于GIL和进程内通信扩展性很差。DDP则完全不同。它启动多个独立进程每个进程持有自己的模型副本和优化器互相之间通过集合通信同步梯度没有主从之分。通信的内容只有梯度而不是每次迭代都要传模型参数。梯度通常比模型参数小不了多少但因为每个进程都在本地完成优化器更新省去了模型参数的反复拷贝加上Ring-AllReduce这类通信算法把压力平摊到所有节点DDP在单机多卡和跨机场景下的吞吐都明显优于DP。下面的对比可以快速看出差异对比维度DataParallelDistributedDataParallel进程模型单进程多线程多进程参数更新位置只在主卡更新再广播各进程本地更新通信内容每轮广播参数、收集梯度只同步梯度负载均衡主卡显存和带宽压力大基本均匀推荐程度简单但扩展性差生产环境通用所以现在PyTorch官方的默认建议也是能上DDP就上DDPDP只适合快速验证或超小规模场景。这也是这个系列花了整篇来讲DDP的原因。2. DDP的同步机制进程组、集合通信与梯度all-reduce2.1 init_process_group背后的三件套backend、init_method、world_sizeDDP看起来只是一个包装类把它套在原模型外面就能用。但套之前你必须先初始化分布式环境。这一步几乎所有的坑都集中在这里理解它比背代码重要得多。torch.distributed.init_process_group()的完整参数很多核心其实只有三个backend、init_method、world_size/rank。backend通信后端。常见的有nccl、gloo、mpi。GPU训练几乎都用nccl它是为英伟达GPU的高带宽通信设计的性能最好CPU训练或者GPU之间网络环境比较特殊时用gloo它是纯TCP实现稳定性高但慢不少。mpi需要额外编译支持我没怎么在实际项目里用过一般不建议碰。init_method进程之间如何互相发现。最常用的有两种env://和tcp://。env://表示从环境变量读取MASTER_ADDR、MASTER_PORT、RANK、WORLD_SIZEtcp://ip:port则是显式指定一个rank0节点的地址其他进程都会去连这个地址。world_size/rankworld_size是所有参与训练的进程总数不是GPU总数也不是节点数。一个进程对应一张卡这种约定只是常见做法技术上world_size可以大于GPU数比如同一张卡上跑多个进程但现实里没人这么干。rank是当前进程的编号从0到world_size-1。还有个相关的local_rank表示当前进程在所在机器内的编号。多机训练时world_size是所有机器上的进程总数local_rank只在单机内看。需要注意的是nccl后端初始化时本身就会建立连接所以init_process_group这行代码往往比想象的耗时。如果后端初始化阶段就卡住优先怀疑网络连通性而不是代码写错了。2.2 Ring-AllReduce如何把梯度同步时间降下来进程组建好以后DDP在做的是对梯度执行all-reduce操作。朴素的做法很好理解每个进程把自己的梯度发给某个中心节点中心节点把所有梯度加起来求平均再把结果发给所有人。问题在于进程多了以后中心节点会变成巨大的瓶颈而且每次通信量等于总数据量乘以进程数网络很快被打爆。Ring-AllReduce的思路反过来了。它把所有进程想象成一个环每个进程只和相邻的两个进程通信。通信分成两个阶段reduce-scatter阶段把梯度张量分成N份每个进程只把自己对应位置那一份作为接收目标。过程中每个进程把自己手上的数据块传给下一个进程同时接收上一个进程的数据块并累加。这一阶段结束时每个进程都得到了一份完整的“部分求和结果”但每个人手里的块不一样。all-gather阶段把剩余缺失的数据块通过同样的环形接力补齐。每个进程把自己已经算好的块继续传给下一个同时接收来自上一个进程的块。最终每个进程都拥有完整的平均梯度。整个过程里每个进程的总通信量大约是2 × (N-1) / N × 数据总量。当N很大时接近2倍数据量与进程数几乎无关所以扩展性比中心化方案好得多。用生活化的比喻一群人围成一圈每个人只把自己那份答案传给旁边的人转两圈之后所有人都拿到了完整答案而不是每个人跑去找一个汇总员对答案。理解了Ring-AllReduce也就理解了为什么DDP在8卡、16卡甚至更大规模下还能保持不错的扩展性。DDP的性能问题很少出在算法本身而是出在通信被计算打断、或者数据没喂够导致计算和通信没有重叠。2.3 反向传播钩子、梯度桶与更新一致性DDP在初始化时会遍历模型的所有参数给每个参数注册一个自动求导钩子。这个钩子的作用是当某个参数的反向传播梯度计算完成后立刻把这个梯度写入一个“缓冲区”而不是马上触发通信。PyTorch之所以这么做是因为每个参数的梯度是分批到达的如果每来一个梯度就通信一次通信次数等于参数量级那训练基本就废了。DDP的解决方案是梯度桶bucket。默认情况下参数按照模型定义的顺序分组成若干个桶每个桶大约25MBbucket_cap_mb参数控制。当一个桶里的所有梯度都到齐时DDP才对这个桶内的梯度做一次all-reduce。通过这种方式梯度同步次数从参数量级降到了桶数量级通信效率大幅提升。这里有个非常重要的细节DDP的梯度同步发生在反向传播结束的时候而不是优化器step()的时候。这是很多从单卡切换到DDP的人踩坑的地方。在单卡代码里我们习惯写loss.backward()然后optimizer.step()再optimizer.zero_grad()因为反正只有一个进程梯度就放在那里慢慢算。但DDP因为有钩子存在backward()返回之前集群范围内的梯度同步就已经发生了。所以如果你想做多步梯度累积每一小步的backward都会触发一次all-reduce通信这和单卡的行为完全不同。这个坑后面专门讲。各个进程用同步好的梯度各自执行optimizer.step()因为初始参数一样、同步后的梯度也一样、优化器逻辑也一样所以更新后的参数仍然保持一致。看起来好像每张卡都在“自己练自己的”实际上每一步都是整个集群协同后的结果。3. 手写一个最小可运行的DDP训练脚本3.1 初始化阶段代码从env变量到进程组理论讲完直接上一段能跑的骨架。这是最常见的基于环境变量的初始化方式import os import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def init_distributed(): dist.init_process_group(backendnccl, init_methodenv://) torch.cuda.set_device(int(os.environ[LOCAL_RANK]))这段代码里注意几个点init_methodenv://表示从环境变量读取配置。使用torchrun或后文说的启动脚本时环境变量会被自动设置好不需要手动读RANK、WORLD_SIZE再传进去。torch.cuda.set_device(local_rank)非常重要。多进程环境下每个进程对应一张卡如果不把当前进程的默认CUDA设备设置成自己的local_rank后续所有张量都可能被创建到cuda:0上轻则浪费显存重则直接OOM。初始化代码要在模型创建之前执行因为通信后端需要提前建好连接。如果你的机器是单机8卡那WORLD_SIZE就是8LOCAL_RANK从0到7。多机的话每台机器上的LOCAL_RANK都是0到N-1但全局RANK不同。3.2 启动器怎么选torchrun与launch的差异早期PyTorch官方教程里大量使用python -m torch.distributed.launch但新版本已经把推荐方式切换成了torchrun。两个工具做的事情差不多根据nproc_per_node在每个节点上拉起N个进程自动设置RANK、WORLD_SIZE、LOCAL_RANK、MASTER_ADDR、MASTER_PORT等环境变量。# 单机8卡 torchrun --nnodes1 --nproc_per_node8 --rdzv-endpointlocalhost:29500 train.py # 多机场景假设有两台机器每台8卡 torchrun --nnodes2 --nproc_per_node8 \ --rdzv-endpointmaster节点ip:29500 \ --rdzv_backendc10d \ train.pytorchrun相比老的launch有几个我很看重的改进自动处理了节点的重新调度和故障重启训练中途有进程崩了不至于整个作业彻底归零。环境变量规范统一脚本代码不用关心是被哪个launcher拉起来的。多机场景下地址发现更灵活用一个共享的redis或者c10d后端就能让不同节点上的进程互相找到对方。在脚本里区分进程身份只需要local_rank int(os.environ[LOCAL_RANK]) global_rank int(os.environ[RANK]) world_size int(os.environ[WORLD_SIZE])如果你看到老代码里用--local_rank传参新版本里已经不推荐了尽量换成环境变量读取。3.3 完整训练循环能被直接改造成自己项目的骨架下面这个骨架可以直接抄到自己的脚本里进行改造import os import torch import torch.nn as nn import torch.optim as optim import torch.distributed as dist from torch.utils.data import DataLoader, DistributedSampler from torch.nn.parallel import DistributedDataParallel as DDP def build_model(): model nn.Sequential( nn.Linear(1024, 2048), nn.ReLU(), nn.Linear(2048, 10), ) return model.cuda() def train(): dist.init_process_group(backendnccl, init_methodenv://) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) model build_model() ddp_model DDP(model, device_ids[local_rank], output_devicelocal_rank) dataset YourDataset() # 替换成自己的数据集类 sampler DistributedSampler(dataset, num_replicasdist.get_world_size(), rankdist.get_rank()) loader DataLoader(dataset, batch_size32, samplersampler, shuffleFalse, num_workers4, drop_lastTrue) optimizer optim.SGD(ddp_model.parameters(), lr0.01, momentum0.9) loss_fn nn.CrossEntropyLoss() for epoch in range(20): sampler.set_epoch(epoch) for x, y in loader: x, y x.cuda(), y.cuda() optimizer.zero_grad() out ddp_model(x) loss loss_fn(out, y) loss.backward() optimizer.step() if dist.get_rank() 0: print(fepoch {epoch} done) dist.destroy_process_group() if __name__ __main__: train()有几个点值得特别解释device_ids[local_rank]是明确告诉DDP这个进程对应的GPU编号。单卡时也可以不传但多卡时最好显式指定。DistributedSampler传入后DataLoader的shuffle必须设为False因为随机打乱的工作由sampler负责。如果不设False每个进程会先内部shuffle再把索引交给sampler不仅浪费还会导致数据划分不符合预期。drop_lastTrue在多机多卡下很有必要否则最后一批不满时各进程数据量可能不一致触发等待问题。日志打印只让rank0做否则8个进程各打一遍日志直接糊墙。这段代码跑起来说明DDP的环境、启动和训练循环已经通了。后面要做的就是在真实项目里处理各种细节问题。4. Sampler、batch与混合精度多进程下逃不掉的三个细节4.1 DistributedSampler的set_epoch与shuffle关系很多人在DDP脚本里写了sampler DistributedSampler(...)但忘记在epoch循环里调用sampler.set_epoch(epoch)。后果是什么每个epoch的数据划分完全一样。单卡训练里DataLoader的shuffle每次epoch都不同但DistributedSampler自己管理一份索引如果外部不告诉它“又过了一个epoch”它就认为你还在重复使用同一个随机划分模型相当于每个epoch都用同样的顺序看到同样的样本随机性被严重削弱。源码里的逻辑是shuffleTrue时set_epoch(epoch)会以epoch为随机种子重新打乱整体索引再按rank分片如果不调用set_epoch所有epoch的索引划分一模一样。所以正确的写法是for epoch in range(epochs): sampler.set_epoch(epoch) for batch in loader: ...这个细节虽然简单却直接影响收敛效果而且跑在小数据集上尤其明显过拟合现象会来得更早。4.2 总batch翻倍之后学习率到底要不要动DDP最常见的使用方式是把原本单卡的batch拆到每张卡上单独算最后梯度聚合相当于更大的batch。比如单卡batch size是328卡DDP下如果每张卡仍然是32那么全局有效batch就是256。这是DDP带来的第一个参数变化全局batch size变大了。线性缩放规则是一个比较实用的起点学习率大约按卡数翻倍。原先单卡用0.018卡时用0.08。但这不是绝对的。对于SGD这类优化器梯度本身是batch内样本梯度的平均batch变大时梯度噪声变小可以适当增大步长但增大太多又会不稳定所以实际调参一般是把学习率提高到原来的1倍到4倍之间看loss曲线的震荡程度。也可以用warmup来缓解高学习率初期的不稳定。如果不想调整学习率一个折中方案是每张卡上的本地batch减小到原来的1/N让全局batch仍然等于单卡时的batch。代价是每张卡的batch很可能太小显存利用率低、计算效率差。正常实践中还是建议维持较大的本地batch让每张卡吃饱再配合调学习率。4.3 梯度累积、drop_last与AMP在DDP下的交互梯度累积是个大坑。它原本是为了解决显存不够、batch又不想减小的问题连续几次backward()但只在累积一定步数后调用一次optimizer.step()等效于更大的batch。单卡上这个逻辑很自然因为梯度在backward()后就躺在那里不动。DDP下问题来了前面说过DDP的gradient hook在backward()返回前就会触发all-reduce所以哪怕你只是想累计三个mini-batch实际上每三个mini-batch也会同步三次梯度通信量是普通训练的3倍。正确的做法是用model.no_sync()在累积的前几步不触发同步for idx, (x, y) in enumerate(loader): with ddp_model.no_sync(): loss loss_fn(ddp_model(x), y) loss.backward() if (idx 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()注意no_sync()上下文管理器只影响这个block内的backward退出来以后再接管的backward又会正常同步。如果你的累积步数不是固定的需要在最后一步确保同步发生。这种做法能大幅削减通信次数但代价是梯度在最后一步all-reduce时个别进程可能因为前几步数据分布差异导致梯度方差略大实际影响一般可接受。drop_last的存在影响了最终有效样本数。世界大小是8、数据集大小是1000时每个sampler大概分配到125条样本如果每张卡读不满一个batch最后一批对齐不上。不设置drop_last时各进程的batch数可能不同出现一个进程还在算、其他进程等待的情况训练时间被拉长且行为不齐。所以多数生产代码都直接drop_lastTrue。混合精度和DDP的交互多见于显存优化场景。PyTorch的torch.cuda.amp.autocast配合GradScaler时梯度的scale和unscale在每个进程本地完成不影响all-reduce的数值一致性所以基本能直接组合使用。唯一要注意的依然是checkpointscaler.state_dict()和optimizer.state_dict()都要正确保存恢复否则后续的loss scale和优化器状态就对不上。5. checkpoint与训练恢复多进程状态下怎么保存才不丢5.1 模型参数保存只在rank0存还是每卡各存一份先明确一个事实DDP进程之间同步的是梯度每个进程用自己的优化器更新参数。因为初始参数一致、每次梯度一致、优化器逻辑一致所以训练过程中各进程的模型参数理论上是一致的。这意味着保存模型权重时不需要每张卡都存一份完整权重通常只在rank0上保存一份就够了if dist.get_rank() 0: torch.save({model: ddp_model.module.state_dict()}, ckpt.pt)不要直接把ddp_model.state_dict()拿去存。DDP包装后的state_dict()会把参数名加上module.前缀如果图省事存了这个加载时又忘了去前缀后面会遇到一堆key对不上的报错。稳妥做法是存ddp_model.module.state_dict()得到的就是原模型的权重字典。但如果要恢复训练问题就变得稍微复杂一点。因为每个进程的优化器状态不完全一样——数据顺序不同局部梯度不同Adam这种带一阶二阶动量累积的优化器在每张卡上累积的统计值会有细微差别。虽然权重一致但优化器状态不一致意味着从中间断点恢复时如果不处理这个差异后续训练表现可能会偏离预期。实操中常见两种策略简单策略只在rank0保存optimizer状态恢复时所有进程都加载rank0的那份。这样做优点是代码简单、文件只有一个缺点是后续各进程的优化器状态被人为统一了相当于发生了一次“强制同步”对训练收敛的影响通常不大但理论上不是最严谨的。严谨策略每个进程保存自己那份优化器状态到独立文件恢复时各进程加载各自对应的文件。比如rank0保存ckpt.rank0.pt、rank1保存ckpt.rank1.pt。这样每个优化器都恢复到自己原本的状态续训行为与中断前一致。我自己的实践是单机多卡且任务比较短时用简单策略就够了长训练任务、大规模分布式训练优先用严谨策略。5.2 恢复训练的完整流程与map_location处理分布式场景加载checkpoint时最容易踩的坑就是设备映射。单卡代码里你经常看到torch.load(ckpt.pt)直接把模型load到cuda:0DDP下如果你不看LOCAL_RANK8个进程可能全部去读cuda:0的显存不仅慢而且把0号卡的显存撑爆。正确的加载写法需要把checkpoint里的张量映射到当前进程对应的卡ckpt torch.load(ckpt.pt, map_location{cuda:0: fcuda:{local_rank}}) model.load_state_dict(ckpt[model]) optimizer.load_state_dict(ckpt[optimizer])map_location{cuda:0: fcuda:{local_rank}}的作用是如果checkpoint里的张量原本在cuda:0加载时自动搬到第local_rank张卡。如果文件里存的是CPU张量直接map_locationfcuda:{local_rank}即可。加载顺序也有讲究。建议先初始化进程组、再把模型装进DDP、然后加载权重而不是先load再包装。因为load_state_dict需要操作原模型的参数key包装前后key不一样容易出错。推荐顺序dist.init_process_group(...) model build_model() ddp_model DDP(model, device_ids[local_rank]) ckpt torch.load(...) ddp_model.module.load_state_dict(ckpt[model]) optimizer.load_state_dict(ckpt[optimizer])5.3 混合精度scaler与EMA状态如何一起恢复如果你的训练用了torch.cuda.amp.GradScaler那么checkpoint里至少应该包含以下内容ckpt { model: ddp_model.module.state_dict(), optimizer: optimizer.state_dict(), scaler: scaler.state_dict(), epoch: epoch, global_step: global_step, }恢复时记得把scaler.load_state_dict(ckpt[scaler])补上。很多人恢复完模型和优化器就把scaler忘了结果后面loss scale从默认值重新开始可能出现梯度下溢或溢出导致训练突然异常。EMA指数移动平均参数在DDP下也要小心。如果在模型外围单独维护了一组EMA参数EMA的更新发生在optimizer.step之后每张卡的本地模型初始一致、更新一致所以EMA参数在各进程间也是一致的。恢复时只要从rank0的checkpoint恢复即可。但要注意EMA参数不要放进DDP的模型参数里被梯度更新否则就失去了EMA的意义。通常做法是单独维护一个dict保存param_name: ema_value在需要评估或生成时替换模型参数或者在每轮结束时用model.module.load_state_dict(ema_dict)这种形式灌回模型。6. 实测中我踩过的坑与性能调优清单6.1 训练吞吐上不去的三个常见原因DDP不是加了就加速很多情况下跑起来比单卡还慢原因往往不在DDP本身而在数据加载和计算节奏上。第一个常见原因是DataLoader成了瓶颈。多卡并行时每张卡都在等数据如果数据增强和磁盘读取都在CPU上num_workers设得又小模型算得再快也没用。用nvidia-smi看GPU利用率如果卡上的利用率长期低于80%先查数据管道。建议把num_workers至少调到CPU线程数的一半同时把pin_memoryTrue打开。第二个原因是梯度同步打断了计算流水。DDP的all-reduce理论上应该和下一轮的forward重叠但如果你模型小、batch也小计算时间比通信时间还短通信就成了瓶颈。这时可以考虑把bucket_cap_mb调大一点减少通信次数也能减少同步粒度或者增大本地batch让计算时间覆盖通信时间。第三是多进程之间负载不均衡。比如某些进程被调度到更慢的GPU上、某张卡显存被打满导致swapping或者数据集本身划分后长度不均。DDP是同步训练最慢的进程决定整体速度。解决思路主要是检查GPU型号是否一致、给每张卡分配合适的数据量、必要时用find_unused_parameters和static_graph减少额外检查开销。6.2 从卡死到超时典型问题排查链路DDP最让人烦的问题就是卡死。遇到过几次之后我发现排查思路是有固定套路的直接按顺序检查看日志输出位置如果init_process_group之后第一轮forward都没跑大概率是通信初始化问题。先确认所有节点的MASTER_ADDR、MASTER_PORT、WORLD_SIZE是否一致防火墙有没有放开端口。看进程数torchrun拉起进程数要与WORLD_SIZE严格一致。如果手动启动时漏算一个进程就会出现某进程一直等着第N个rank加入表现就是毫无报错地卡住。看报错时间点如果卡在第一次backward大概率是梯度同步链路有问题。这时候打开环境变量NCCL_DEBUGINFO再看日志里all-reduce的进度能定位是哪两台机器之间通信超时。看超时时间init_process_group里的timeout参数默认是1800秒如果连接异常等待半小时才报超时相当痛苦。调试阶段可以设短一点比如timeouttimedelta(minutes5)让错误更快暴露。如果NCCL进程反复报各种连接错误先看不同机器上的网卡是否互通、是否处于同一网络段。跨机训练时最好有稳定的内网连接不要依赖公网IP。6.3 几个立刻能用的性能参数设置这里给一个比较稳的参数起点是我在单机4卡和单机8卡上都验证过的组合DDP( model, device_ids[local_rank], bucket_cap_mb50, static_graphTrue, )bucket_cap_mb50把梯度桶从默认25MB调大到50MB减少通信次数。模型特别大时收益明显模型很小或者batch很小时过大的桶反而会让第一轮通信等太久建议根据实测来回调整。static_graphTrue当模型结构在训练过程中完全不变大多数CV、NLP任务都满足时开启这个参数可以跳过一些动态图相关的重检查逻辑小幅提升性能。数据管道层面建议这样设置DataLoaderDataLoader( dataset, batch_sizebatch_size, samplersampler, num_workers8, pin_memoryTrue, persistent_workersTrue, prefetch_factor4, drop_lastTrue, )persistent_workersTrue让dataloader的worker子进程在epoch之间不退出减少反复创建进程的开销prefetch_factor4提前预取更多样本让GPU少等CPU。注意这些参数要匹配实际内存prefetch_factor太大内存会涨得厉害。另外如果你在多机场景下用NCCL建议把环境变量NCCL_DEBUGINFO只在排查问题时打开平时关掉否则日志量巨大可以用NCCL_DEBUGWARN保留关键告警。最后说一个我反复强调的经验DDP下的“本地batch变大”和“学习率调整”一定要分开验证。先把通信和训练跑通再一次性调优这些参数。实际操作中我更喜欢先固定一个小batch把DDP流程完全跑通确认loss能稳定下降、checkpoint能恢复然后再逐步增大batch、调学习率。否则一上来把所有因素混在一起出了问题很难定位是DDP配置的错还是超参变化导致的波动。从原理到代码再到踩坑DDP的运作逻辑其实没有想象中那么黑盒。掌握进程组、梯度同步方式、sampler规则和checkpoint策略这四件事DDP基本就能在你的项目里站住脚了。如果你正在从单卡往多卡迁移不妨先把这篇里的最小骨架跑通再往里面填自己的模型和数据遇到问题按最后那节的排查链路一步步走会比直接拿完整训练代码硬调试省力得多。
把模型从单卡搬到多卡最怕的不是代码写不出来而是训练跑起来之后效率反而比单卡还差甚至直接卡死不动。PyTorch分布式训练里最常用的主力工具就是torch.nn.parallel.DistributedDataParallel也就是大家口中的DDP。无论你是第一次尝试把模型扩展到两张卡还是准备在一台8卡机器上跑较大规模并行DDP基本都是绕不开的主路径。这篇是分布式系列的第三篇我会把DDP的原理、启动方式、采样规则、checkpoint策略、常见坑和调参路径一次讲透。适合已经用PyTorch写过完整训练循环、但对多卡运行还停留在“听说”阶段的读者也适合那些跑通了DDP但经常遇到卡死、梯度过不去、checkpoint加载出错的兄弟。1. 单卡到多卡分布式训练解决的问题到底是什么1.1 显存墙与算力墙单卡训练的天花板很多初学者以为分布式训练的核心目的就是“显存不够了多卡拼一下”。这个说法只对了一半。显存确实是第一道墙一张消费级显卡24GB已经是比较常见的容量大模型或者大batch训练时光是一个batch的数据加激活值就可能把显存吃穿更别提模型本身还占着一大块。但就算你的模型小到能塞进单卡训练时间依然是硬伤。一个需要跑十几个小时甚至几天的训练任务单卡只能慢慢等而多卡并行可以把同样规模的数据切成多份让每张卡同时干活吞吐量翻上去。所以说DDP解决的是两件事第一是装不下第二是等不起。装不下的时候DDP靠数据并行把不同batch数据分到不同卡上模型本身在每张卡都有一份完整副本所以模型显存这部分没有省下来省的是单卡上同时处理的数据量和中间激活内存。等不起的时候DDP靠多卡同时计算梯度、同时更新参数训练总时间近似与卡数成反比。1.2 DDP要解决的三类核心问题模型同步、数据切分、进程协调DDP真正要处理的不是“把数据分下去”这么简单。数据切分只是第一步后面还跟着三件麻烦事模型同步每张卡上都有一份独立的模型参数。训练刚开始时它们是一样的但反向传播计算出的梯度依赖于各自的数据所以各卡的梯度天然不同。想要让模型继续保持“多卡算出来的更新方向一致”就必须把梯度在所有卡之间做一次集体通信让每张卡拿到全局平均后的梯度再用这个梯度去更新自己的参数副本。数据切分多卡训练时每个进程不能拿到完整的batch也不能随便撞车否则整体训练就相当于几倍batch混在一起还互相重复。DDP配合DistributedSampler把整个数据集切分成互不重叠的若干份分别交给不同进程。进程协调多进程不是多个线程它们不共享内存也没有全局锁。进程之间需要约定谁是master、谁是worker、谁先启动谁等待还要初始化通信后端。这些协调工作都靠torch.distributed里的进程组完成。打个比方单卡训练是一个人从头到尾做题数据并行训练是一组人分别做不同的题目但每做完一部分就要停下来对答案如果其中一个人的答案和其他人差距太大所有人都要按平均后的答案来调整自己的思路。DDP就是那个组织“对答案”的角色。1.3 为什么选DDP而不是老牌DataParallelPyTorch里实现数据并行的工具其实有两个方向。老一点的torch.nn.DataParallelDP用一个进程管理多张卡输入batch被切成几块分发给各卡前向结果回收到主卡主卡计算loss后反向传播然后把更新后的参数广播给各卡。这个方案最大的问题是主卡承担了太多活儿分发数据、接收输出、计算梯度、广播参数全都压在主卡上所以经常出现主卡显存比其他卡高一大截的情况而且受限于GIL和进程内通信扩展性很差。DDP则完全不同。它启动多个独立进程每个进程持有自己的模型副本和优化器互相之间通过集合通信同步梯度没有主从之分。通信的内容只有梯度而不是每次迭代都要传模型参数。梯度通常比模型参数小不了多少但因为每个进程都在本地完成优化器更新省去了模型参数的反复拷贝加上Ring-AllReduce这类通信算法把压力平摊到所有节点DDP在单机多卡和跨机场景下的吞吐都明显优于DP。下面的对比可以快速看出差异对比维度DataParallelDistributedDataParallel进程模型单进程多线程多进程参数更新位置只在主卡更新再广播各进程本地更新通信内容每轮广播参数、收集梯度只同步梯度负载均衡主卡显存和带宽压力大基本均匀推荐程度简单但扩展性差生产环境通用所以现在PyTorch官方的默认建议也是能上DDP就上DDPDP只适合快速验证或超小规模场景。这也是这个系列花了整篇来讲DDP的原因。2. DDP的同步机制进程组、集合通信与梯度all-reduce2.1 init_process_group背后的三件套backend、init_method、world_sizeDDP看起来只是一个包装类把它套在原模型外面就能用。但套之前你必须先初始化分布式环境。这一步几乎所有的坑都集中在这里理解它比背代码重要得多。torch.distributed.init_process_group()的完整参数很多核心其实只有三个backend、init_method、world_size/rank。backend通信后端。常见的有nccl、gloo、mpi。GPU训练几乎都用nccl它是为英伟达GPU的高带宽通信设计的性能最好CPU训练或者GPU之间网络环境比较特殊时用gloo它是纯TCP实现稳定性高但慢不少。mpi需要额外编译支持我没怎么在实际项目里用过一般不建议碰。init_method进程之间如何互相发现。最常用的有两种env://和tcp://。env://表示从环境变量读取MASTER_ADDR、MASTER_PORT、RANK、WORLD_SIZEtcp://ip:port则是显式指定一个rank0节点的地址其他进程都会去连这个地址。world_size/rankworld_size是所有参与训练的进程总数不是GPU总数也不是节点数。一个进程对应一张卡这种约定只是常见做法技术上world_size可以大于GPU数比如同一张卡上跑多个进程但现实里没人这么干。rank是当前进程的编号从0到world_size-1。还有个相关的local_rank表示当前进程在所在机器内的编号。多机训练时world_size是所有机器上的进程总数local_rank只在单机内看。需要注意的是nccl后端初始化时本身就会建立连接所以init_process_group这行代码往往比想象的耗时。如果后端初始化阶段就卡住优先怀疑网络连通性而不是代码写错了。2.2 Ring-AllReduce如何把梯度同步时间降下来进程组建好以后DDP在做的是对梯度执行all-reduce操作。朴素的做法很好理解每个进程把自己的梯度发给某个中心节点中心节点把所有梯度加起来求平均再把结果发给所有人。问题在于进程多了以后中心节点会变成巨大的瓶颈而且每次通信量等于总数据量乘以进程数网络很快被打爆。Ring-AllReduce的思路反过来了。它把所有进程想象成一个环每个进程只和相邻的两个进程通信。通信分成两个阶段reduce-scatter阶段把梯度张量分成N份每个进程只把自己对应位置那一份作为接收目标。过程中每个进程把自己手上的数据块传给下一个进程同时接收上一个进程的数据块并累加。这一阶段结束时每个进程都得到了一份完整的“部分求和结果”但每个人手里的块不一样。all-gather阶段把剩余缺失的数据块通过同样的环形接力补齐。每个进程把自己已经算好的块继续传给下一个同时接收来自上一个进程的块。最终每个进程都拥有完整的平均梯度。整个过程里每个进程的总通信量大约是2 × (N-1) / N × 数据总量。当N很大时接近2倍数据量与进程数几乎无关所以扩展性比中心化方案好得多。用生活化的比喻一群人围成一圈每个人只把自己那份答案传给旁边的人转两圈之后所有人都拿到了完整答案而不是每个人跑去找一个汇总员对答案。理解了Ring-AllReduce也就理解了为什么DDP在8卡、16卡甚至更大规模下还能保持不错的扩展性。DDP的性能问题很少出在算法本身而是出在通信被计算打断、或者数据没喂够导致计算和通信没有重叠。2.3 反向传播钩子、梯度桶与更新一致性DDP在初始化时会遍历模型的所有参数给每个参数注册一个自动求导钩子。这个钩子的作用是当某个参数的反向传播梯度计算完成后立刻把这个梯度写入一个“缓冲区”而不是马上触发通信。PyTorch之所以这么做是因为每个参数的梯度是分批到达的如果每来一个梯度就通信一次通信次数等于参数量级那训练基本就废了。DDP的解决方案是梯度桶bucket。默认情况下参数按照模型定义的顺序分组成若干个桶每个桶大约25MBbucket_cap_mb参数控制。当一个桶里的所有梯度都到齐时DDP才对这个桶内的梯度做一次all-reduce。通过这种方式梯度同步次数从参数量级降到了桶数量级通信效率大幅提升。这里有个非常重要的细节DDP的梯度同步发生在反向传播结束的时候而不是优化器step()的时候。这是很多从单卡切换到DDP的人踩坑的地方。在单卡代码里我们习惯写loss.backward()然后optimizer.step()再optimizer.zero_grad()因为反正只有一个进程梯度就放在那里慢慢算。但DDP因为有钩子存在backward()返回之前集群范围内的梯度同步就已经发生了。所以如果你想做多步梯度累积每一小步的backward都会触发一次all-reduce通信这和单卡的行为完全不同。这个坑后面专门讲。各个进程用同步好的梯度各自执行optimizer.step()因为初始参数一样、同步后的梯度也一样、优化器逻辑也一样所以更新后的参数仍然保持一致。看起来好像每张卡都在“自己练自己的”实际上每一步都是整个集群协同后的结果。3. 手写一个最小可运行的DDP训练脚本3.1 初始化阶段代码从env变量到进程组理论讲完直接上一段能跑的骨架。这是最常见的基于环境变量的初始化方式import os import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def init_distributed(): dist.init_process_group(backendnccl, init_methodenv://) torch.cuda.set_device(int(os.environ[LOCAL_RANK]))这段代码里注意几个点init_methodenv://表示从环境变量读取配置。使用torchrun或后文说的启动脚本时环境变量会被自动设置好不需要手动读RANK、WORLD_SIZE再传进去。torch.cuda.set_device(local_rank)非常重要。多进程环境下每个进程对应一张卡如果不把当前进程的默认CUDA设备设置成自己的local_rank后续所有张量都可能被创建到cuda:0上轻则浪费显存重则直接OOM。初始化代码要在模型创建之前执行因为通信后端需要提前建好连接。如果你的机器是单机8卡那WORLD_SIZE就是8LOCAL_RANK从0到7。多机的话每台机器上的LOCAL_RANK都是0到N-1但全局RANK不同。3.2 启动器怎么选torchrun与launch的差异早期PyTorch官方教程里大量使用python -m torch.distributed.launch但新版本已经把推荐方式切换成了torchrun。两个工具做的事情差不多根据nproc_per_node在每个节点上拉起N个进程自动设置RANK、WORLD_SIZE、LOCAL_RANK、MASTER_ADDR、MASTER_PORT等环境变量。# 单机8卡 torchrun --nnodes1 --nproc_per_node8 --rdzv-endpointlocalhost:29500 train.py # 多机场景假设有两台机器每台8卡 torchrun --nnodes2 --nproc_per_node8 \ --rdzv-endpointmaster节点ip:29500 \ --rdzv_backendc10d \ train.pytorchrun相比老的launch有几个我很看重的改进自动处理了节点的重新调度和故障重启训练中途有进程崩了不至于整个作业彻底归零。环境变量规范统一脚本代码不用关心是被哪个launcher拉起来的。多机场景下地址发现更灵活用一个共享的redis或者c10d后端就能让不同节点上的进程互相找到对方。在脚本里区分进程身份只需要local_rank int(os.environ[LOCAL_RANK]) global_rank int(os.environ[RANK]) world_size int(os.environ[WORLD_SIZE])如果你看到老代码里用--local_rank传参新版本里已经不推荐了尽量换成环境变量读取。3.3 完整训练循环能被直接改造成自己项目的骨架下面这个骨架可以直接抄到自己的脚本里进行改造import os import torch import torch.nn as nn import torch.optim as optim import torch.distributed as dist from torch.utils.data import DataLoader, DistributedSampler from torch.nn.parallel import DistributedDataParallel as DDP def build_model(): model nn.Sequential( nn.Linear(1024, 2048), nn.ReLU(), nn.Linear(2048, 10), ) return model.cuda() def train(): dist.init_process_group(backendnccl, init_methodenv://) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) model build_model() ddp_model DDP(model, device_ids[local_rank], output_devicelocal_rank) dataset YourDataset() # 替换成自己的数据集类 sampler DistributedSampler(dataset, num_replicasdist.get_world_size(), rankdist.get_rank()) loader DataLoader(dataset, batch_size32, samplersampler, shuffleFalse, num_workers4, drop_lastTrue) optimizer optim.SGD(ddp_model.parameters(), lr0.01, momentum0.9) loss_fn nn.CrossEntropyLoss() for epoch in range(20): sampler.set_epoch(epoch) for x, y in loader: x, y x.cuda(), y.cuda() optimizer.zero_grad() out ddp_model(x) loss loss_fn(out, y) loss.backward() optimizer.step() if dist.get_rank() 0: print(fepoch {epoch} done) dist.destroy_process_group() if __name__ __main__: train()有几个点值得特别解释device_ids[local_rank]是明确告诉DDP这个进程对应的GPU编号。单卡时也可以不传但多卡时最好显式指定。DistributedSampler传入后DataLoader的shuffle必须设为False因为随机打乱的工作由sampler负责。如果不设False每个进程会先内部shuffle再把索引交给sampler不仅浪费还会导致数据划分不符合预期。drop_lastTrue在多机多卡下很有必要否则最后一批不满时各进程数据量可能不一致触发等待问题。日志打印只让rank0做否则8个进程各打一遍日志直接糊墙。这段代码跑起来说明DDP的环境、启动和训练循环已经通了。后面要做的就是在真实项目里处理各种细节问题。4. Sampler、batch与混合精度多进程下逃不掉的三个细节4.1 DistributedSampler的set_epoch与shuffle关系很多人在DDP脚本里写了sampler DistributedSampler(...)但忘记在epoch循环里调用sampler.set_epoch(epoch)。后果是什么每个epoch的数据划分完全一样。单卡训练里DataLoader的shuffle每次epoch都不同但DistributedSampler自己管理一份索引如果外部不告诉它“又过了一个epoch”它就认为你还在重复使用同一个随机划分模型相当于每个epoch都用同样的顺序看到同样的样本随机性被严重削弱。源码里的逻辑是shuffleTrue时set_epoch(epoch)会以epoch为随机种子重新打乱整体索引再按rank分片如果不调用set_epoch所有epoch的索引划分一模一样。所以正确的写法是for epoch in range(epochs): sampler.set_epoch(epoch) for batch in loader: ...这个细节虽然简单却直接影响收敛效果而且跑在小数据集上尤其明显过拟合现象会来得更早。4.2 总batch翻倍之后学习率到底要不要动DDP最常见的使用方式是把原本单卡的batch拆到每张卡上单独算最后梯度聚合相当于更大的batch。比如单卡batch size是328卡DDP下如果每张卡仍然是32那么全局有效batch就是256。这是DDP带来的第一个参数变化全局batch size变大了。线性缩放规则是一个比较实用的起点学习率大约按卡数翻倍。原先单卡用0.018卡时用0.08。但这不是绝对的。对于SGD这类优化器梯度本身是batch内样本梯度的平均batch变大时梯度噪声变小可以适当增大步长但增大太多又会不稳定所以实际调参一般是把学习率提高到原来的1倍到4倍之间看loss曲线的震荡程度。也可以用warmup来缓解高学习率初期的不稳定。如果不想调整学习率一个折中方案是每张卡上的本地batch减小到原来的1/N让全局batch仍然等于单卡时的batch。代价是每张卡的batch很可能太小显存利用率低、计算效率差。正常实践中还是建议维持较大的本地batch让每张卡吃饱再配合调学习率。4.3 梯度累积、drop_last与AMP在DDP下的交互梯度累积是个大坑。它原本是为了解决显存不够、batch又不想减小的问题连续几次backward()但只在累积一定步数后调用一次optimizer.step()等效于更大的batch。单卡上这个逻辑很自然因为梯度在backward()后就躺在那里不动。DDP下问题来了前面说过DDP的gradient hook在backward()返回前就会触发all-reduce所以哪怕你只是想累计三个mini-batch实际上每三个mini-batch也会同步三次梯度通信量是普通训练的3倍。正确的做法是用model.no_sync()在累积的前几步不触发同步for idx, (x, y) in enumerate(loader): with ddp_model.no_sync(): loss loss_fn(ddp_model(x), y) loss.backward() if (idx 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()注意no_sync()上下文管理器只影响这个block内的backward退出来以后再接管的backward又会正常同步。如果你的累积步数不是固定的需要在最后一步确保同步发生。这种做法能大幅削减通信次数但代价是梯度在最后一步all-reduce时个别进程可能因为前几步数据分布差异导致梯度方差略大实际影响一般可接受。drop_last的存在影响了最终有效样本数。世界大小是8、数据集大小是1000时每个sampler大概分配到125条样本如果每张卡读不满一个batch最后一批对齐不上。不设置drop_last时各进程的batch数可能不同出现一个进程还在算、其他进程等待的情况训练时间被拉长且行为不齐。所以多数生产代码都直接drop_lastTrue。混合精度和DDP的交互多见于显存优化场景。PyTorch的torch.cuda.amp.autocast配合GradScaler时梯度的scale和unscale在每个进程本地完成不影响all-reduce的数值一致性所以基本能直接组合使用。唯一要注意的依然是checkpointscaler.state_dict()和optimizer.state_dict()都要正确保存恢复否则后续的loss scale和优化器状态就对不上。5. checkpoint与训练恢复多进程状态下怎么保存才不丢5.1 模型参数保存只在rank0存还是每卡各存一份先明确一个事实DDP进程之间同步的是梯度每个进程用自己的优化器更新参数。因为初始参数一致、每次梯度一致、优化器逻辑一致所以训练过程中各进程的模型参数理论上是一致的。这意味着保存模型权重时不需要每张卡都存一份完整权重通常只在rank0上保存一份就够了if dist.get_rank() 0: torch.save({model: ddp_model.module.state_dict()}, ckpt.pt)不要直接把ddp_model.state_dict()拿去存。DDP包装后的state_dict()会把参数名加上module.前缀如果图省事存了这个加载时又忘了去前缀后面会遇到一堆key对不上的报错。稳妥做法是存ddp_model.module.state_dict()得到的就是原模型的权重字典。但如果要恢复训练问题就变得稍微复杂一点。因为每个进程的优化器状态不完全一样——数据顺序不同局部梯度不同Adam这种带一阶二阶动量累积的优化器在每张卡上累积的统计值会有细微差别。虽然权重一致但优化器状态不一致意味着从中间断点恢复时如果不处理这个差异后续训练表现可能会偏离预期。实操中常见两种策略简单策略只在rank0保存optimizer状态恢复时所有进程都加载rank0的那份。这样做优点是代码简单、文件只有一个缺点是后续各进程的优化器状态被人为统一了相当于发生了一次“强制同步”对训练收敛的影响通常不大但理论上不是最严谨的。严谨策略每个进程保存自己那份优化器状态到独立文件恢复时各进程加载各自对应的文件。比如rank0保存ckpt.rank0.pt、rank1保存ckpt.rank1.pt。这样每个优化器都恢复到自己原本的状态续训行为与中断前一致。我自己的实践是单机多卡且任务比较短时用简单策略就够了长训练任务、大规模分布式训练优先用严谨策略。5.2 恢复训练的完整流程与map_location处理分布式场景加载checkpoint时最容易踩的坑就是设备映射。单卡代码里你经常看到torch.load(ckpt.pt)直接把模型load到cuda:0DDP下如果你不看LOCAL_RANK8个进程可能全部去读cuda:0的显存不仅慢而且把0号卡的显存撑爆。正确的加载写法需要把checkpoint里的张量映射到当前进程对应的卡ckpt torch.load(ckpt.pt, map_location{cuda:0: fcuda:{local_rank}}) model.load_state_dict(ckpt[model]) optimizer.load_state_dict(ckpt[optimizer])map_location{cuda:0: fcuda:{local_rank}}的作用是如果checkpoint里的张量原本在cuda:0加载时自动搬到第local_rank张卡。如果文件里存的是CPU张量直接map_locationfcuda:{local_rank}即可。加载顺序也有讲究。建议先初始化进程组、再把模型装进DDP、然后加载权重而不是先load再包装。因为load_state_dict需要操作原模型的参数key包装前后key不一样容易出错。推荐顺序dist.init_process_group(...) model build_model() ddp_model DDP(model, device_ids[local_rank]) ckpt torch.load(...) ddp_model.module.load_state_dict(ckpt[model]) optimizer.load_state_dict(ckpt[optimizer])5.3 混合精度scaler与EMA状态如何一起恢复如果你的训练用了torch.cuda.amp.GradScaler那么checkpoint里至少应该包含以下内容ckpt { model: ddp_model.module.state_dict(), optimizer: optimizer.state_dict(), scaler: scaler.state_dict(), epoch: epoch, global_step: global_step, }恢复时记得把scaler.load_state_dict(ckpt[scaler])补上。很多人恢复完模型和优化器就把scaler忘了结果后面loss scale从默认值重新开始可能出现梯度下溢或溢出导致训练突然异常。EMA指数移动平均参数在DDP下也要小心。如果在模型外围单独维护了一组EMA参数EMA的更新发生在optimizer.step之后每张卡的本地模型初始一致、更新一致所以EMA参数在各进程间也是一致的。恢复时只要从rank0的checkpoint恢复即可。但要注意EMA参数不要放进DDP的模型参数里被梯度更新否则就失去了EMA的意义。通常做法是单独维护一个dict保存param_name: ema_value在需要评估或生成时替换模型参数或者在每轮结束时用model.module.load_state_dict(ema_dict)这种形式灌回模型。6. 实测中我踩过的坑与性能调优清单6.1 训练吞吐上不去的三个常见原因DDP不是加了就加速很多情况下跑起来比单卡还慢原因往往不在DDP本身而在数据加载和计算节奏上。第一个常见原因是DataLoader成了瓶颈。多卡并行时每张卡都在等数据如果数据增强和磁盘读取都在CPU上num_workers设得又小模型算得再快也没用。用nvidia-smi看GPU利用率如果卡上的利用率长期低于80%先查数据管道。建议把num_workers至少调到CPU线程数的一半同时把pin_memoryTrue打开。第二个原因是梯度同步打断了计算流水。DDP的all-reduce理论上应该和下一轮的forward重叠但如果你模型小、batch也小计算时间比通信时间还短通信就成了瓶颈。这时可以考虑把bucket_cap_mb调大一点减少通信次数也能减少同步粒度或者增大本地batch让计算时间覆盖通信时间。第三是多进程之间负载不均衡。比如某些进程被调度到更慢的GPU上、某张卡显存被打满导致swapping或者数据集本身划分后长度不均。DDP是同步训练最慢的进程决定整体速度。解决思路主要是检查GPU型号是否一致、给每张卡分配合适的数据量、必要时用find_unused_parameters和static_graph减少额外检查开销。6.2 从卡死到超时典型问题排查链路DDP最让人烦的问题就是卡死。遇到过几次之后我发现排查思路是有固定套路的直接按顺序检查看日志输出位置如果init_process_group之后第一轮forward都没跑大概率是通信初始化问题。先确认所有节点的MASTER_ADDR、MASTER_PORT、WORLD_SIZE是否一致防火墙有没有放开端口。看进程数torchrun拉起进程数要与WORLD_SIZE严格一致。如果手动启动时漏算一个进程就会出现某进程一直等着第N个rank加入表现就是毫无报错地卡住。看报错时间点如果卡在第一次backward大概率是梯度同步链路有问题。这时候打开环境变量NCCL_DEBUGINFO再看日志里all-reduce的进度能定位是哪两台机器之间通信超时。看超时时间init_process_group里的timeout参数默认是1800秒如果连接异常等待半小时才报超时相当痛苦。调试阶段可以设短一点比如timeouttimedelta(minutes5)让错误更快暴露。如果NCCL进程反复报各种连接错误先看不同机器上的网卡是否互通、是否处于同一网络段。跨机训练时最好有稳定的内网连接不要依赖公网IP。6.3 几个立刻能用的性能参数设置这里给一个比较稳的参数起点是我在单机4卡和单机8卡上都验证过的组合DDP( model, device_ids[local_rank], bucket_cap_mb50, static_graphTrue, )bucket_cap_mb50把梯度桶从默认25MB调大到50MB减少通信次数。模型特别大时收益明显模型很小或者batch很小时过大的桶反而会让第一轮通信等太久建议根据实测来回调整。static_graphTrue当模型结构在训练过程中完全不变大多数CV、NLP任务都满足时开启这个参数可以跳过一些动态图相关的重检查逻辑小幅提升性能。数据管道层面建议这样设置DataLoaderDataLoader( dataset, batch_sizebatch_size, samplersampler, num_workers8, pin_memoryTrue, persistent_workersTrue, prefetch_factor4, drop_lastTrue, )persistent_workersTrue让dataloader的worker子进程在epoch之间不退出减少反复创建进程的开销prefetch_factor4提前预取更多样本让GPU少等CPU。注意这些参数要匹配实际内存prefetch_factor太大内存会涨得厉害。另外如果你在多机场景下用NCCL建议把环境变量NCCL_DEBUGINFO只在排查问题时打开平时关掉否则日志量巨大可以用NCCL_DEBUGWARN保留关键告警。最后说一个我反复强调的经验DDP下的“本地batch变大”和“学习率调整”一定要分开验证。先把通信和训练跑通再一次性调优这些参数。实际操作中我更喜欢先固定一个小batch把DDP流程完全跑通确认loss能稳定下降、checkpoint能恢复然后再逐步增大batch、调学习率。否则一上来把所有因素混在一起出了问题很难定位是DDP配置的错还是超参变化导致的波动。从原理到代码再到踩坑DDP的运作逻辑其实没有想象中那么黑盒。掌握进程组、梯度同步方式、sampler规则和checkpoint策略这四件事DDP基本就能在你的项目里站住脚了。如果你正在从单卡往多卡迁移不妨先把这篇里的最小骨架跑通再往里面填自己的模型和数据遇到问题按最后那节的排查链路一步步走会比直接拿完整训练代码硬调试省力得多。
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
