1. 从一张卡跑不动说起DDP到底替你干了什么第一次把模型和数据规模同时往上抬的时候多数人都会经历同一个场景单卡显存直接爆掉改小 batch size 又掉精度于是开始找多卡方案。搜到的第一篇文章往往教你DataParallel写起来确实简单模型外面包一层就完事但真正跑起来你会发现主卡显存比副卡高一大截GPU 利用率曲线像心电图训练速度提升远低于卡数。这不是你代码写错了而是DataParallel的运行模型决定的它只开一个进程靠 Python 多线程把 batch 切到各张卡上前向之后梯度全部汇集到主卡做归约主卡既要算自己那份又要负责汇总Python 的全局解释器锁还会让线程调度卡顿。卡越多主卡越像瓶颈。DistributedDataParallel后面统一叫 DDP换了一套思路一张卡一个进程。每个进程持有完整的模型副本、完整的优化器状态各自读自己那份数据做前向和反向只在反向计算完梯度之后做一次跨进程的梯度归约让所有副本的梯度保持一致再由各自的优化器独立更新参数。因为初始参数相同、每步梯度相同、更新规则相同各副本在训练全程都保持同步逻辑上等价于用了一个超大 batch 的单卡训练。这个等价是理解 DDP 所有行为的基础后面遇到的 batch size 换算、学习率调整、BatchNorm 统计量偏差根子都在这里。它解决的不只是显存问题。多进程意味着每个进程可以独立设置 CUDA 上下文、独立开 DataLoader worker、独立做混合精度通信走的是针对 GPU 互联优化的集合通信库卡间走高速互联时归约开销可以压到很低。适合谁来上手只要你的训练脚本已经能在单卡跑通并且手上有两台以上 GPU单机多卡也算DDP 就是当前最值得投入时间掌握的并行方案。下面这些内容是我在实际项目里反复改脚本、反复排障之后沉淀下来的东西不是官方文档的复述。1.1 单卡脚本到多卡的三种路线为什么最后都收敛到 DDP把单卡训练改成多卡市面上常见三条路。第一条是手写torch.multiprocessing自己起进程、自己写通信逻辑灵活但极易写出死锁第二条是DataParallel改动最小但性能天花板很低第三条就是 DDP。三者放在一起对比差距非常直观对比项单进程多线程DataParallel多进程 DDP手写多进程进程模型1 进程 N 线程N 进程一卡一进程N 进程自行管理梯度同步方式汇集到主卡做归约环形/树形集合通信去中心化自己实现主卡显存偏斜明显主卡常多出 1~2GB基本均衡取决于实现Python GIL 影响有线程调度受限无无代码改动量最小中等需改数据加载与启动方式大多机扩展不支持原生支持需自行处理DataParallel的问题不在 API 层面而在架构层面它把分发和归约都压在主卡所在进程里卡数一多主卡既要承担通信中转又要承担计算扩展性天然受限。我实测过 4 卡场景同样的模型和 batchDataParallel的吞吐只有 DDP 的六成左右而且卡数越多差距越大。所以哪怕只是单机 2 卡我也建议直接上 DDP省得后面重构。1.2 一次 all-reduce 的账梯度到底是怎么对齐的DDP 在每个进程里注册了反向传播的钩子。当某一层参数的梯度算完钩子会把这个梯度标记为待同步等一个 bucket默认 25MB攒满或者反向结束就触发一次集合通信。通信的语义不是简单广播而是对所有进程的梯度做求和再除以进程数也就是求平均。之所以用平均而不是求和是为了让梯度的量级与单卡训练保持一致学习率不用因为卡数变化而重调当然全局 batch 变了学习率还是要按规则缩放这个后面单独讲。这里有个常被忽略的细节DDP 的归约是分 bucket 异步进行的反向计算和通信可以重叠。靠近输出层的梯度先算完先开始通信此时更浅层的反向还在算两者并行。这个设计是 DDP 吞吐能接近线性的关键之一。bucket 太小通信次数多、启动开销占比高bucket 太大通信要等更久才能发起重叠效果变差。默认 25MB 对大多数模型是合理值但如果你的模型单层特别大比如超大 embedding或者层数极多但每层很小就值得手动调bucket_cap_mb。还有一点DDP 只同步参与反向计算的参数的梯度。如果你在模型里写了条件分支某些参数这一步没参与计算它们的梯度是NoneDDP 默认会报错需要显式告诉它find_unused_parametersTrue。这个参数不是免费的它会让 DDP 在每次反向时额外遍历计算图去判断哪些参数没被用到开销随模型复杂度上升。所以正确做法是尽量把分支写清楚只在确实需要时开这个开关而不是图省事一直开着。1.3 进程组、rank、world_size三个词先掰扯明白DDP 的所有协调都建立在进程组这个概念上。world_size是参与训练的总进程数单机 4 卡就是 4双机各 4 卡就是 8。rank是全局进程编号从 0 到world_size - 1全局唯一用来做只有 rank 0 打印日志只有 rank 0 存模型这类判断。local_rank是进程在本机内的编号决定这个进程该用哪张卡torch.cuda.set_device(local_rank)就是靠它绑定的。这三个值的关系必须刻在脑子里因为绝大部分 DDP 报错都能追溯到它们没对齐world_size设错会让进程互相等待直到超时rank重复会让两个进程抢同一份数据、写同一个文件local_rank和实际使用的卡不一致会出现明明 4 卡nvidia-smi 上只有 2 张在跑的诡异现象。用torchrun启动时这三个值会通过环境变量自动注入所以你不需要手动传如果自己用mp.spawn起进程就得自己算好再传进去。这也是我推荐优先用torchrun的原因少写一层容易出错的封装。提示调试期在init_process_group之后立刻打印一次rank / local_rank / world_size和当前绑定的卡号这行日志能省掉后面大量猜测。2. 把单卡脚本改成DDP改动量比想象中小但顺序不能错一个写得规整的单卡训练脚本改成 DDP 通常只需要动五个地方进程组初始化、设备绑定、DataLoader 的 sampler、模型包一层 DDP、启动命令。听起来不多但每一步都有顺序依赖顺序错了就会卡住或者静默算错。我见过最常见的情况是脚本改完能跑loss 也在降但最后评估指标比单卡差一截回头查了半天才发现是验证集数据被重复采样了。动手之前先明确一件事DDP 下每个进程都是完整的训练程序它有自己的优化器、自己的学习率调度器、自己的日志。你要做的是让这些完整的程序看到不同的数据、在梯度上达成一致而不是把程序拆成几段分给不同进程。理解这一点改代码时就不会纠结这段该放在 rank 0 还是所有 rank——除了写文件、打日志、上传指标这类副作用操作其余全部放在所有 rank 上执行。2.1 进程组初始化backend、超时和设备绑定初始化的核心就三行但每行都有讲究import os import torch import torch.distributed as dist def setup_dist(): # torchrun 会注入 RANK / WORLD_SIZE / LOCAL_RANK local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) dist.init_process_group( backendnccl, timeouttorch.timedelta(minutes30), ) return local_rankbackend选nccl是 GPU 训练的默认答案它针对卡间通信做了深度优化单机内走 NVLink 或 PCIe跨机走网卡。CPU 训练才考虑gloo。set_device必须放在init_process_group之前或紧随其后原因是 NCCL 在初始化时会读取当前设备如果此时还没绑定所有进程可能都默认用 0 号卡结果就是显存全压在一张卡上其余卡空转。这个问题在日志里完全看不出来只有在nvidia-smi里才能发现。超时设置值得单独说。默认超时是 30 分钟不同版本略有差异对于大模型单步耗时长的场景可能不够一旦某个 rank 卡住其他 rank 会等满超时然后一起崩报出一堆NCCL timeout或者Watchdog caught collective operation timeout。把超时显式设大一点配合后面要讲的日志手段能让你在进程被杀之前有机会看清楚是谁卡住了。另外init_process_group是阻塞的它会等所有 rank 都到齐所以如果你看到程序停在初始化不动八成是某个 rank 根本没起来或者world_size和实际进程数对不上。2.2 DistributedSampler数据切分与那几个必须改的细节单卡时DataLoader的shuffleTrue就能工作DDP 下必须换成DistributedSampler否则每个进程都会读到全量数据等于把 batch 放大了world_size倍还做了重复计算from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler sampler DistributedSampler( dataset, num_replicasdist.get_world_size(), rankdist.get_rank(), shuffleTrue, drop_lastTrue, # 样本数不能被 world_size 整除时丢弃尾部保证各 rank 步数一致 ) loader DataLoader( dataset, batch_sizeper_gpu_bs, samplersampler, num_workers8, pin_memoryTrue, drop_lastTrue, )drop_lastTrue在这里几乎是必须的。如果不丢最后一个 batch 在每个 rank 上的大小可能不同反向时各进程的通信节奏就对不上轻则报形状不匹配重则直接 hang。代价是每个 epoch 丢掉少量样本对训练影响可以忽略但换来的稳定性很值。还有一个极易漏掉的调用每个 epoch 开始前要sampler.set_epoch(epoch)。DistributedSampler的 shuffle 依赖一个内部的随机种子这个种子由epoch决定。如果你不调用set_epoch每个 epoch 的切分方式完全一样模型看到的样本顺序恒定数据增强的随机性也被削弱表现为训练前期 loss 降得还行后期明显乏力。这个问题极其隐蔽因为代码不报错指标也不崩只是慢一点、差一点很多人会误以为是模型容量不够。提示把sampler.set_epoch(epoch)写在 epoch 循环的第一行和自己约定这行不能删。我在一个项目里因为它漏写白跑了两天实验才从对比曲线上看出来。另外num_workers的总量要算总账。单进程 8 个 worker8 卡就是 64 个进程如果机器只有 32 核worker 之间会互相抢 CPU数据加载反而变慢。经验值是num_workers * world_size CPU 核数 * 1.5超了就往下调。容器环境还要注意共享内存/dev/shm给小了会导致 DataLoader 卡死表现为进程在跑但 GPU 利用率长期为 0。2.3 模型包DDP之前这三件事必须先做完模型不是随便包一层就行的包之前有三件事必须完成顺序不能颠倒第一模型要先搬到对应的卡上。model.to(local_rank)要在构造 DDP 之前执行DDP 不接受跨设备的模型。第二优化器要在包 DDP 之后创建或者用model.parameters()时注意引用。严格说优化器在包之前创建也能工作因为 DDP 包装的是同一个 module 的参数不复制参数但为了可读性我习惯先包再用ddp_model.parameters()建优化器避免后面加载 checkpoint 时 key 名前缀不一致。第三确认没有未参与计算的参数。如果模型里有只在特定条件下才走的分支要么重构掉要么在构造时传find_unused_parametersTrue。判断方法很简单先不加这个参数跑一个 step报错了再加。from torch.nn.parallel import DistributedDataParallel as DDP model MyModel().to(local_rank) ddp_model DDP( model, device_ids[local_rank], output_devicelocal_rank, broadcast_buffersTrue, # 默认 True每次 forward 前广播 buffer保证 BN 统计量一致 find_unused_parametersFalse, # 确认没有未使用参数后再关掉省开销 gradient_as_bucket_viewTrue, # 梯度直接写入 bucket省一次内存拷贝 ) optimizer torch.optim.AdamW(ddp_model.parameters(), lrlr)broadcast_buffers这个开关值得多解释一句。模型里的 BatchNorm 有 running_mean、running_var 这类 buffer它们不参与梯度计算但会影响前向结果。DDP 默认在每次 forward 前把 rank 0 的 buffer 广播给所有进程保证各副本的前向行为一致。这在训练时是必要的但推理阶段会带来额外开销可以在 eval 前把ddp_model.broadcast_buffers False或构造时就设 False只在训练时注意代价是各卡的 BN 统计量可能不同步。我的习惯是训练保持默认纯推理场景改掉。2.4 启动方式torchrun 与 mp.spawn 的取舍启动命令看着简单但它决定了环境变量怎么注入进而决定了你的脚本能不能读到正确的 rank# 单机 8 卡 torchrun --nnodes1 --nproc_per_node8 --master_port29500 train.py --config cfg.yaml # 双机每机 8 卡两台机器上分别执行node_rank 不同 torchrun --nnodes2 --nproc_per_node8 \ --node_rank0 \ --master_addr10.0.0.11 \ --master_port29500 \ train.py --config cfg.yamltorchrun会自动设置RANK、WORLD_SIZE、LOCAL_RANK、MASTER_ADDR、MASTER_PORT还会在某个进程挂掉时把整组进程一起收掉避免留下僵尸进程占着显存。mp.spawn需要你自己算rank node_rank * gpus_per_node local_rank并且异常处理要自己写一旦某个子进程崩了其他进程可能还在傻等最后靠超时才退出。新项目没有历史包袱的话直接用torchrun。--master_port冲突是新手常踩的坑同一台机器上跑多个训练任务端口不换就会互相抢占症状是第二个任务起不来或者起一半就断。我的做法是每次起任务前随机挑一个 20000 到 30000 之间的端口或者在脚本里用--master_port$((20000 RANDOM % 10000))。3. 第一次跑通不等于跑对loss曲线背后的问题脚本能起来、loss 在下降只能说明流程通了不能说明算对了。DDP 的错误有一大类是静默的不报错、不崩溃但结果就是偏差。我见过最典型的一次是4 卡训练跑完验证集指标比单卡低了三个点团队排查了两天最后定位到验证集用了DistributedSampler但没有做跨进程收集每个 rank 各自算自己那部分样本的指标然后简单平均而各 rank 的样本数不一样尾数没对齐平均值被拉偏了。这类问题的共同特征是数值层面的错误不会触发任何异常只能靠对照实验发现。所以养成一个习惯——改完 DDP 之后先用小数据、小步数做一次单卡与多卡的对照。同样的种子、同样的数据子集跑 50 步看 loss 曲线是否一致允许多卡因为有 BN 和随机性带来的微小差异但量级应该吻合。这一步花十分钟能省掉后面几天的返工。3.1 全局batch、学习率与等价batch的换算DDP 最容易被误解的一点是DataLoader里写的batch_size是单卡 batch全局 batch 等于它乘以world_size。单卡 batch 8、8 卡训练全局就是 64。如果你从单卡 batch 64 的实验迁过来只改了启动方式没改 batch_size 参数那实际全局 batch 变成了 512学习率却还是原来那个收敛行为会明显不同。线性缩放规则是业界常用的起点全局 batch 扩大 k 倍学习率也扩大 k 倍。但这只是起点不是定律。batch 很大时线性缩放会让训练前期不稳定通常配合 warmup 使用前若干步把学习率从一个小值线性升到目标值。单卡 batch卡数全局 batch建议初始 lr以单卡 lr1e-4 为基准是否建议 warmup161161e-4否164644e-4建议500 步1681288e-4建议1000 步3282568e-4 ~ 1.6e-3必须1000~2000 步表里的数字是经验起点实际还要看优化器类型。AdamW 对学习率的敏感度比 SGD 高大 batch 下更保守一些SGD 配 momentum 时线性缩放的效果更稳定。我的做法是先用线性缩放跑 500 步看 loss 是否平稳如果前期抖动明显就降一半再试别一上来就上 cosine 调度那样出了问题分不清是调度策略的锅还是初始学习率的锅。还有一个隐性坑学习率调度器的步数要按全局步数算。单卡时一个 epoch 有 N 个 step8 卡时每个 rank 一个 epoch 只有 N/8 个 step因为数据被切了。如果你的 scheduler 是按 step 数配置的比如 cosine 的T_max迁移时忘了改训练会在计划外的位置结束衰减后段学习率过低模型欠拟合。3.2 日志、打印与副作用只留rank 0说话DDP 下所有 rank 都会执行同一份代码如果你在训练循环里直接print(loss)8 卡就会打出 8 行几乎一样的日志交错在一起根本没法看。正确做法是用 rank 判断包一层def log_info(msg, rank0): if dist.get_rank() rank: print(f[rank{rank}] {msg}, flushTrue) # 训练循环里 if step % 50 0: log_info(fstep{step} loss{loss.item():.4f} lr{scheduler.get_last_lr()[0]:.2e})flushTrue不要省。多进程环境下标准输出有缓冲不 flush 的话日志可能攒着一起吐或者进程被杀时直接丢失排障时你会怀疑自己是不是没打这行日志。写文件同理只让 rank 0 写其他 rank 跳过否则多个进程同时写同一个文件内容会交错甚至损坏。指标上报无论是写本地文件还是推到监控系统也要做单 rank 处理。但要注意如果你要在日志里记录全局平均 loss直接拿 rank 0 的 loss 是不准确的它只是那一份数据的 loss。严谨做法是用dist.all_reduce把各 rank 的 loss 求和再除以 world_size代价是一次很小的通信。日常调试用 rank 0 的值足够正式实验建议做全局归约。3.3 验证集上的重复样本与指标对不上前面提到的那个坑展开说一下正确写法。评估阶段有两种常见做法第一种每个 rank 用DistributedSampler(shuffleFalse, drop_lastFalse)分到不同子集各自前向把预测结果和标签all_gather到 rank 0 上拼起来再统一算指标。这种做法效率最高但要处理样本总数不能被 world_size 整除导致的 padding拼接时要把 padding 部分裁掉。第二种只让 rank 0 跑完整验证集其他 rank 用dist.barrier()等待。实现简单但浪费了其他卡的算力验证集大时会明显拖慢。# 方式一的核心片段跨进程收集预测 def gather_predictions(preds, labels): world_size dist.get_world_size() gathered_preds [torch.zeros_like(preds) for _ in range(world_size)] gathered_labels [torch.zeros_like(labels) for _ in range(world_size)] dist.all_gather(gathered_preds, preds) dist.all_gather(gathered_labels, labels) return torch.cat(gathered_preds, dim0), torch.cat(gathered_labels, dim0)注意all_gather要求各进程的张量形状一致所以验证集的 batch 要保证每个 rank 拿到的数量相同用 padding 补齐或者drop_last否则会在通信阶段报形状错误。这类错误信息通常不直观会提示某个维度不匹配得回头查数据切分。3.4 checkpoint存一份还是每卡都存保存模型只让 rank 0 做这个大家都知道。但有两个细节容易出问题。第一保存前要确保其他 rank 不会继续往前跑。如果 rank 0 在写几百 MB 的 checkpoint其他 rank 已经进入下一轮训练并开始集合通信rank 0 还没到它们就会阻塞等待。短时间没问题但如果写盘很慢网络存储可能触发通信超时。稳妥写法是在保存前后加 barrierif dist.get_rank() 0: torch.save({ model: ddp_model.module.state_dict(), # 注意 .module optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict(), epoch: epoch, }, ckpt_path) dist.barrier()第二保存的 state_dict 要去掉 DDP 包装。ddp_model.state_dict()和ddp_model.module.state_dict()在多数版本里 key 是一致的DDP 不额外加前缀但为了保险和可移植统一用.module.state_dict()这样存出来的文件能直接被单卡脚本加载不依赖 DDP 环境。断点续训时模型权重和优化器状态加载完后别忘了scheduler的步数也要恢复否则学习率会从初始值重新开始续训效果和预期完全不符。这个坑我在两个项目里都见过因为 scheduler 的状态往往被当成不重要而忽略。4. 卡住不动、报NCCL错误、速度不升反降排障链条DDP 的故障排查和单卡完全不是一个路数。单卡报错堆栈直接指到出问题的代码行DDP 出问题你看到的往往是某个进程报错其他进程在等或者所有进程都在等谁也没报错。这时候需要的不是读堆栈而是建立一套从外到内的排查顺序先确认进程是否都起来了再确认集合通信是否对齐最后才看具体是哪个 rank 的计算出了问题。我把这套顺序总结成一句话先看有几个进程活着再看它们在等谁最后看谁先掉队。下面按这个顺序展开。4.1 进程hang住的几种典型形态与定位顺序第一种形态程序启动后没有任何输出GPU 占用为 0。这通常是init_process_group没通过原因多是world_size和实际进程数不一致或者MASTER_ADDR/PORT配错导致进程连不上。排查方法是先看nvidia-smi上有几个进程再对照你期望的卡数。如果进程数就少问题在启动命令如果进程数够但都卡在初始化问题在网络地址或端口。第二种形态训练跑了几十步突然全体卡住GPU 利用率掉到 0。这多半是集合通信次数不对齐——某个 rank 多跑了一次all_reduce或者某个 rank 提前退出了循环。典型诱因是数据量不能被整除导致各 rank 步数不同或者代码里有if rank 0包住的集合通信调用集合通信必须所有 rank 都参与包在条件里必然死锁。第三种形态跑着跑着报NCCL timeout或Watchdog caught collective operation timeout然后整组退出。这种情况通常是某个 rank 在计算上花了异常长的时间比如遇到了超大样本、显存不足触发换页、磁盘 IO 卡顿其他 rank 等超时。定位的关键是看报错前最后一条日志来自哪个 rank、卡在第几步。# 排查时常用的两个命令 # 1. 看每张卡上的进程和显存 nvidia-smi --query-compute-appspid,used_memory --formatcsv # 2. 设置通信库日志级别拿到更详细的握手信息 export NCCL_DEBUGINFO export NCCL_DEBUG_SUBSYSINIT,COLLNCCL_DEBUGINFO会在启动时打印每张卡用的网卡、互联类型、拓扑结构很多跨机通信慢的问题从这段日志就能看出来——比如本该走高速互联却走了普通网卡。4.2 NCCL相关环境变量先确认互联拓扑NCCL 的默认行为是自动探测拓扑并选择最优路径多数情况下不需要手动干预。但在容器、多网卡、混合互联的环境里自动探测经常选错表现为跨机通信极慢或者初始化失败。这时候需要手动指定网卡# 指定用于跨机通信的网卡名多网卡环境必须显式指定 export NCCL_SOCKET_IFNAMEeth0 # 关闭 IB纯以太网环境可以试 export NCCL_IB_DISABLE1 # 打开 P2P单机内 NVLink/PCIe 直连场景有用 export NCCL_P2P_LEVELNVL这些变量不是开了就快而是环境不匹配时用来纠偏。我的建议是默认什么都不设先跑一次基准记录下单步耗时如果跨机明显慢于单机再逐个加变量对比。盲目照搬网上的环境变量组合很可能把本来正常的自动探测覆盖掉。单机多卡还有一个常见误区以为所有卡之间都是高速互联。实际上很多机器的卡间是 PCIe 分组的跨组的卡通信要经过 CPU 中转速度差好几倍。如果你的模型对通信敏感参数多、层数少、计算量小卡的分组方式会直接影响扩展效率。这一点在租用机器时要问清楚或者在NCCL_DEBUGINFO的拓扑输出里自己确认。4.3 慢节点拖累全局找出摸鱼的那个rank分布式训练的吞吐由最慢的那个 rank 决定。8 个进程里 7 个都算得快1 个因为磁盘读得慢或者数据预处理重整体速度就会被它拖住。这种问题在总吞吐上看是扩展效率不高但单看每张卡的利用率都是 100%因为快的在等慢的等待时也显示占用非常难定位。有效手段是给每个 rank 的关键阶段打时间戳训练若干步后汇总对比import time t0 time.time() batch next(iter(loader)) t1 time.time() loss ddp_model(**batch).loss loss.backward() optimizer.step() t2 time.time() stats torch.tensor([t1 - t0, t2 - t1], devicelocal_rank) gathered [torch.zeros_like(stats) for _ in range(dist.get_world_size())] dist.all_gather(gathered, stats) if dist.get_rank() 0: for r, s in enumerate(gathered): print(frank{r}: data{s[0]:.3f}s compute{s[1]:.3f}s)跑十几步取平均哪个 rank 的 data 时间明显偏高问题就在它的数据加载上。常见原因包括该 rank 分到的样本恰好都是大文件、num_workers不够、该进程所在 NUMA 节点内存带宽被抢占。定位到具体 rank 之后再去看它绑定的卡和 CPU 亲和性问题通常就清楚了。4.4 find_unused_parameters 的代价与显存不均衡find_unused_parametersTrue是很多人在遇到Expected to have finished reduction in the prior iteration报错后随手加上的开关。它确实能让代码跑起来但代价有两个一是每次反向多一次计算图遍历模型层数多时开销可达 10% 以上二是它会改变 DDP 对 bucket 的划分逻辑可能让通信与计算的重叠效果变差。更麻烦的是这个开关会掩盖真正的设计问题。模型里出现未参与计算的参数通常意味着有分支逻辑或者有参数压根没被用到。前者应该通过调整前向路径让所有参数每步都参与或者用no_sync配合后者说明模型定义有问题白占显存。我处理过的案例里有一半以上在清理掉无用参数之后find_unused_parameters就可以关掉单步耗时直接降了一成。显存不均衡是另一个隐蔽问题。理论上各 rank 的模型副本大小相同显存占用应该接近但实际经常看到 rank 0 多出几百 MB。原因通常是 rank 0 额外承担了日志张量、指标缓存、checkpoint 缓冲区。如果卡本来就快满了rank 0 会先 OOM而报错信息看起来像是模型太大误导排查方向。解决办法是把这些额外操作尽量放到 CPU 内存里做或者在保存前先释放缓存。5. 调优与工程化让DDP真正吃满机器跑通和跑好之间隔着一段距离。同样的硬件、同样的模型两个人写出来的 DDP 脚本吞吐可能差 30% 以上。差距主要来自几个地方通信与计算的重叠程度、数据加载是否成为瓶颈、混合精度与梯度累积的配合方式、以及一些容易被忽略的显存与带宽细节。这一节讲的是把 DDP 从能用推到用满的具体手段。5.1 bucket调优与梯度视图减少无谓的内存搬运前面提过 bucket 的作用这里给出调优方法。判断 bucket 是否合理的指标是通信时间占单步时间的比例如果通信占比很低比如 5% 以下说明重叠做得好不用动如果占比超过 20%就值得试几个不同的bucket_cap_mb值做对比。ddp_model DDP( model, device_ids[local_rank], bucket_cap_mb50, # 默认 25模型层大时可以调大 gradient_as_bucket_viewTrue, # 梯度直接写进 bucket省一次拷贝 static_graphTrue, # 计算图结构固定的模型可以开进一步优化 )gradient_as_bucket_viewTrue的效果是把归约后的梯度直接以 bucket 的视图形式暴露给优化器避免梯度先拷到 bucket、归约后再拷回来的来回搬运。这个开关在大模型上能省下可观的显存和时间但要求你的代码不去原地修改梯度比如手动做梯度裁剪时要小心。static_graphTrue适用于每步计算图完全一致的模型它让 DDP 可以提前规划好通信顺序进一步减少协调开销但如果模型有动态分支开了会报错或者算错用之前要确认。我用一组 8 卡、约 3 亿参数模型的实测做过对比只改这两个开关单步耗时从 1.42 秒降到 1.28 秒大约 10% 的提升改动成本几乎为零属于性价比最高的调优项。5.2 混合精度、梯度累积与no_sync的配合混合精度和 DDP 的组合有一个真实的坑GradScaler是每个进程独立的它会根据本进程是否出现 inf/nan 动态调整 scale 值。如果某个 rank 因为数值问题把 scale 降低了而其他 rank 没降那么各进程的梯度经过了不同的缩放再做平均就失去了数学意义训练会悄悄跑偏。规避方式有三种。第一种用 bf16 代替 fp16bf16 的动态范围足够大通常不需要 GradScaler从根上绕开这个问题。第二种如果必须用 fp16在optimizer.step()之前用集合通信把各进程的 scale 同步成一致的值取最小值。第三种关掉动态调整用固定 scale代价是需要提前试出合适的值。现在硬件支持 bf16 的话我优先选第一种省事且稳。梯度累积的目的是在小显存下模拟大 batch。DDP 下有一个专门的优化累积的中间步骤不需要通信只在最后一次反向时同步for i, batch in enumerate(loader): with ddp_model.no_sync() if (i 1) % accum_steps ! 0 else contextlib.nullcontext(): loss ddp_model(**batch).loss / accum_steps loss.backward() if (i 1) % accum_steps 0: optimizer.step() optimizer.zero_grad(set_to_noneTrue)no_sync()的作用是临时关闭梯度同步让累积的前几步只在本进程内计算。如果不加每一步都会触发一次全量归约通信量变成accum_steps倍累积就完全失去了意义。这里要注意 loss 要除以accum_steps否则梯度量级会放大等效学习率变了。提示no_sync()必须覆盖除了最后一次之外的所有累积步条件判断写反了不会报错只是通信量翻倍从监控上看是训练变慢了但不知道为什么。5.3 多机启动参数与容错的实际处理跨机训练比单机多卡复杂的地方在于网络。启动参数里--master_addr要填一个所有节点都能访问的地址--node_rank从 0 开始编号且不能重复。两台机器上都要执行torchrun参数除了node_rank之外完全一致。容错方面PyTorch 提供torchrun的弹性训练能力配合--rdzv_backend、--rdzv_endpoint、--max_restarts某个节点挂掉后可以自动拉起重新加入。但弹性训练要求你的脚本能在启动时从 checkpoint 恢复否则重启等于从头开始。我的建议是先把 checkpoint 的保存与恢复做扎实再考虑弹性。很多团队上了弹性训练却发现恢复逻辑有 bug重启后模型状态不对反而比直接失败更糟。跨机还要注意时钟同步。日志里各节点的时间戳如果差几分钟排障时根本对不上事件顺序。容器环境下时钟漂移更常见部署前确认一下各节点时间是同步的。5.4 项目里沉淀下来的一些习惯最后分享几个我在多个项目里固定下来的做法都是踩过坑之后形成的条件反射。第一所有 DDP 相关代码集中在一个dist_utils.py里包括初始化、清理、日志封装、指标归约、checkpoint 保存。训练脚本只调用封装好的函数不在业务代码里散落dist.get_rank()判断。这样换项目时可以整包拷过去也不会出现某处忘了判断 rank 导致 8 个进程同时写文件的问题。第二启动脚本里固定打印一次环境摘要world_size、每台机器的卡数、NCCL 版本、是否启用了 IB、各 rank 的 hostname。这份摘要存在日志开头出问题时第一时间能确认环境是否符合预期省掉大量来回询问。第三用小规模数据做冒烟测试。改完 DDP 相关代码先用 100 条样本、跑 20 步、4 卡验证一遍确认 loss 有下降、checkpoint 能存能读、验证集指标计算正常再上全量数据。全量跑一次几小时冒烟测试两分钟这个投入产出比没什么好犹豫的。第四显存预留留出余量。DDP 的通信 bucket、NCCL 的内部缓冲、日志张量都会占显存按单卡能跑的最大 batch 去配置很容易在多卡场景下 OOM。我的经验是留 10% 左右的显存余量并且用torch.cuda.memory_summary()定期看一眼峰值。第五也是我认为最重要的一条不要迷信扩展效率这个数字。8 卡跑到 7 倍吞吐是很好的结果但如果你为了追这个数字把单卡 batch 调得极小、通信占比极高实际训练效果可能还不如 4 卡配大 batch。分布式训练的目标是缩短达到目标精度的时间不是把 GPU 利用率数字做漂亮。我见过为了追求线性扩展把 batch 压到 1 的实验最后收敛慢得离谱白白浪费了机器。真要把 DDP 用顺手关键不在于记住多少个 API而在于脑子里始终有一张图几个进程各自持有一份完整模型吃不同的数据在梯度上对齐独立更新。所有诡异现象——hang、慢、指标不对、显存不均——都能从这张图里推出可能的原因然后再用日志和对照实验去验证。这套思路比任何一份配置模板都管用。
