资讯动态

PyTorch DDP单机多卡训练实战:从踩坑到加速比优化

发布时间:2026/10/28 3:01:52 来源:尧图企业网站定制
单卡跑一个 epoch 要四个小时八张卡跑完还是四个小时——这是我第一次动手把PyTorch 分布式训练的DDP接进自己项目时遇到的真实状况。程序没报错日志照打loss 照降但nvidia-smi一看八张卡里只有一张在干活另外七张挂在原地等。后来的两天里我陆续踩完了init_process_group端口冲突、sampler 长度不一致卡死、BatchNorm 统计量不对、学习率没跟着缩放导致模型直接发散这一整套坑。所以这篇东西不打算写成 API 手册而是按“我当初想知道但没人告诉我”的顺序来组织DDP 到底做了什么、脚本怎么改、数字怎么重算、哪里最容易翻车、怎么把线性加速比从 0.5 拉回到 0.9 附近。如果你手上有个能单卡跑通的训练脚本急着用多卡把训练时间压下来或者你已经挂上了 DDP 但加速比难看、结果复现不了、偶尔莫名其妙卡死这篇基本对得上你的场景。全文假定你用的是一台机器上的多张同型号加速卡最后也会给多机多卡的接入方式但重心放在单机多卡因为九成的人第一步卡在这里。1. 动手前先想清楚DDP 解决的是哪一层问题1.1 单卡慢的真正瓶颈分布很多人以为多卡是为了“跑更大的模型”实际上绝大多数人上多卡的第一动机是“同样的模型跑得快一点”。这两件事在实现难度上差着一个量级前者要考虑模型并行、张量切分、流水线调度后者只需要数据并行把不同的数据分给不同的卡各自算梯度再把梯度同步一下。DDP 做的就是后者如果你要切模型那是另一个话题FSDP、张量并行那一套别指望 DDP 帮你解决。先把时间账算清楚。一个训练 step 的耗时大致等于数据加载时间加前向时间加反向时间加优化器更新时间。单卡时这四项是串行的你换多卡能加速的其实只有前向和反向这两块数据加载如果num_workers不够、IO 打满多卡只会让你更快地撞上同一个 IO 瓶颈优化器更新在参数量大的时候也不算小头Adam 一步要读写好几份和参数同尺寸的状态张量但 DDP 是按卡复制模型的每个 rank 都自己算优化器更新这部分不会加速反而因为通信插在中间被拉长了。所以一个粗略的判断标准是单卡时用 profiler 看一下前向加反向占比超过七成上 DDP 收益可观如果数据加载占了一半以上先把 DataLoader 调好再谈分布式否则你加了卡只会看见等待时间被放大。这个结论我吃过亏——早期给一个读大量小文件的任务上八卡加速比只有 1.6后来把num_workers从 4 提到 12、配上persistent_workersTrue和prefetch_factor八卡才跑到 6.8 倍。1.2 DDP 与 DataParallel 的根本差别老代码里常见torch.nn.DataParallel一行model nn.DataParallel(model)就能跑非常诱人。但它的工作方式是单进程多线程主线程把输入切成小块 scatter 到各卡各卡算完 forward 后 gather 回主卡算 loss再 scatter 回去做 backward最后梯度汇总在主卡上做更新。问题在于整个 Python 解释器只有一把 GIL模型参数和优化器也只存在于主卡的显存里主卡显存占用天然比别的卡高出一截通信全挤在 PCIe 上而且是单点汇聚卡越多主卡越堵。八卡规模下 DP 能给你 2 到 3 倍加速就算不错了还经常主卡先 OOM。DDP 是另一种结构每个进程独占一张卡模型在每个进程里各有一份完整副本前向反向全在本地算只有在反向传播算完梯度时才做一次全局的梯度同步。没有主卡没有单点通信走的是环形或者树形可以跑在卡间的高速互联上。代价是你要管进程、要管数据切分、要管日志只打一份代码不再是加一行那么简单。但真正的区别在扩展性DDP 的通信量和卡数基本无关后面讲 Ring-AllReduce 会说清楚而 DP 的通信量随卡数线性上涨所以 DP 在八卡以上基本没救DDP 可以一路扩到几百卡。对比项DataParallelDDP进程模型单进程多线程每卡一个进程显存分布主卡额外承担汇聚开销各卡基本均衡通信模式单点汇聚走 PCIe环形/树形可用高速互联通信量随卡数近似线性增长基本持平八卡典型加速比2 到 3 倍6 到 7 倍代码改造量一行中等需要改启动方式和数据切分1.3 什么项目值得上 DDP什么项目纯属折腾我的经验判断线是这样单卡单个 epoch 超过两小时且你已经确认数据加载不是瓶颈那就值得上单卡一小时以内能跑完的实验改代码加调坑的时间大概率超过你省下的时间先把手上的实验排完更划算。另外一个容易被忽略的点是显存DDP 不会帮你省显存它是复制模型的每张卡都要放完整的一份参数、梯度、优化器状态和激活值。想把“单卡放不下的模型”跑起来是 FSDP 或者梯度检查点的事DDP 爱莫能助。还有一种情况是纯粹为了“跑大 batch”。有些任务确实对 batch size 敏感比如对比学习、大 batch 的检测训练。这时候 DDP 的价值不在速度而在能凑出更大的全局 batch只要你能把学习率和 warmup 调对收益是实打实的。反过来如果你只是想跑得快而模型本身很小、单步耗时只有几十毫秒那通信开销占比会很高八卡的加速比可能只有五倍出头这种任务更适合先把torch.compile、更大 batch、混合精度这些单卡优化做完再考虑分布式。2. 揭开盖子DDP 每一步到底在干什么2.1 一进程一卡进程组、rank 与通信后端DDP 的所有魔法都建立在“进程组”这个概念上。启动时每个进程调用init_process_group通过一个约定的地址和端口互相找到对方握手成功后形成一个组。组里每个进程有个全局编号叫rank从 0 到 world_size-1还有一个本地编号local_rank指的是这台机器上的第几张卡。单机八卡时 rank 和 local_rank 数值相同多机时就不是了——第二台机器的 rank 0 对应 local_rank 0这个区别在多机脚本里写错set_device会导致所有进程抢同一张卡是很经典的翻车点。通信后端的选择很直接有加速卡就用nccl纯 CPU 环境用gloo。早期有人为了图省事在 GPU 任务里用 gloo结果发现慢了一个数量级因为 gloo 根本不走卡间互联数据要绕回主机内存再发出去。有一点要注意init_process_group的默认超时是 30 分钟新版本里首次连接的超时更短如果模型一大、初始化慢或者某张卡正在做很长的预处理就可能直接抛超时中断。我习惯在初始化时显式给一个更宽的超时尤其是多机或者节点间网络一般的情况import datetime import torch.distributed as dist dist.init_process_group( backendnccl, timeoutdatetime.timedelta(minutes60), )另外一个常被忽略的细节是设备绑定。init_process_group之后必须紧接着torch.cuda.set_device(local_rank)或者用device torch.device(cuda, local_rank)显式指定。如果忘了所有进程默认都用 0 号卡你会看到 0 号卡显存爆掉其他卡一根毛都没动。这个错误的表象是 OOM很容易被误导去找 batch size 的问题。2.2 梯度同步的核心AllReduce 与分桶流水线反向传播时每个 rank 算出的梯度只反映了自己那部分数据的贡献。要保证所有卡上的模型保持一致就得让每张卡拿到全局的梯度平均。这个“所有进程把各自的数据汇总、算出结果后再发回给所有人”的操作就是 AllReduce。最朴素的实现是每张卡把自己的梯度发给 rank 0rank 0 求平均再广播回去通信量是 2×(N-1)/N 倍的数据量而且 rank 0 是瓶颈。实际用的是 Ring-AllReduce所有卡排成一个环把梯度切成 N 份先做一轮 scatter-reduce每张卡只负责累加其中一份转 N-1 步之后每张卡手里都有一份完整的部分和再做一轮 all-gather把这 N 份结果转着圈传一遍最后每张卡都拿到完整结果。总通信量是 2(N-1)/N 倍数据量当 N 很大时趋近于 2 倍数据量跟卡数几乎无关。这就是 DDP 能扩展到几百卡的理论基础。更关键的是分桶。如果等整个模型的梯度都算完再一次性 AllReduce那通信和计算是完全串行的卡在通信上的时间会非常可观。DDP 的做法是反向传播一开始就按参数逆序把梯度往桶里塞桶满了默认 25MB就立刻对这个桶发起异步 AllReduce同时反向传播继续算下一层。这样通信和计算重叠起来实际暴露出来的等待时间就只剩最后一两个桶。这也解释了一个调参点bucket_cap_mb设太小会导致启动很多次通信、延迟占主导设太大会让重叠变差、等待变长25MB 这个默认值在多数模型上是合理的参数量特别大的模型可以试着调到 50 到 100 观察一下。2.3 一个反直觉的点DDP 平均的是梯度这里必须单独拎出来讲因为它直接决定你的 loss 写法对不对。DDP 在 AllReduce 的时候对梯度做的是求平均除以 world_size不是求和。也就是说如果你的每个 rank 上 loss 用的是mean归约那你什么都不用改直接loss.backward()就得到正确的全局平均梯度。对比一下 DPDP 是把各卡的输出 gather 起来算一个总 loss 再反传梯度天然是“总 batch”的量级。而 DDP 是各自算各自的 loss 再平均梯度数学上等价于“每张卡的平均 loss 再平均”也就是全局平均。这一点等价性在 batch 长度一致时成立如果你的各卡 batch 大小不一样比如最后一个 batch 没 drop那“平均的平均”就不等于真正的全局平均梯度会有轻微偏差。这就是为什么分布式训练里强烈建议drop_lastTrue。还有一个相关的坑在梯度累积。假设你设了累积 4 步你写了loss criterion(...) / 4再 backward看起来对但 DDP 每次 backward 都会做一次 AllReduce 把梯度平均掉所以累积 4 步之后你得到的其实是 4 个 micro-batch 梯度的平均而不是总和。正确做法是仍然除以累积步数但要理解这里的语义loss/accum配合 DDP 的平均等价于“对 4×world_size 个 micro-batch 求平均”这正好是你想要的。如果你在累积里忘了除梯度会放大 4 倍配合缩放后的学习率很容易炸。3. 从单卡脚本到 DDP改造清单与完整骨架3.1 改造清单需要动的地方一共七处我把改造拆成一张清单每次新项目照着走一遍二十分钟能搞定比对着文档东找西找快得多。序号改造点单卡写法DDP 写法1进程初始化无init_process_groupset_device2模型包装model.to(device)再包一层DDP(model, device_ids[local_rank])3数据采样shuffleTrueDistributedSampler 每 epochset_epoch4数据加载DataLoader(dataset, shuffleTrue)DataLoader(dataset, samplersampler)5日志打印直接 print只在 rank 0 打印6保存检查点直接 save只在 rank 0 存model.module.state_dict()7启动方式python train.pytorchrun --nproc_per_nodeN train.py有一个顺序问题要强调DistributedSampler必须配shuffleFalse传给 DataLoader因为你把 sampler 显式传进去了DataLoader 的 shuffle 参数必须关掉否则会报冲突。另外drop_last在两个层面都要考虑sampler 的drop_lastTrue保证各 rank 分到的样本数一致DataLoader 的drop_lastTrue保证每个 rank 上最后那个不完整的 batch 被丢掉。两个都开最省心。3.2 完整可跑的骨架代码下面这份骨架我删掉了业务逻辑只留 DDP 相关的结构你可以直接往里填模型和数据。import os import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, DistributedSampler def setup_dist(): dist.init_process_group(backendnccl) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) return local_rank def cleanup_dist(): dist.barrier() dist.destroy_process_group() def build_loaders(dataset_train, dataset_val, batch_size, num_workers): train_sampler DistributedSampler(dataset_train, shuffleTrue, drop_lastTrue) val_sampler DistributedSampler(dataset_val, shuffleFalse, drop_lastFalse) train_loader DataLoader( dataset_train, batch_sizebatch_size, samplertrain_sampler, num_workersnum_workers, pin_memoryTrue, drop_lastTrue, persistent_workersnum_workers 0, ) val_loader DataLoader( dataset_val, batch_sizebatch_size, samplerval_sampler, num_workersnum_workers, pin_memoryTrue, drop_lastFalse, persistent_workersnum_workers 0, ) return train_loader, val_loader, train_sampler def main(): local_rank setup_dist() rank dist.get_rank() world_size dist.get_world_size() torch.manual_seed(1234) torch.cuda.manual_seed_all(1234) model MyModel().to(local_rank) model DDP(model, device_ids[local_rank], find_unused_parametersFalse) optimizer torch.optim.AdamW(model.parameters(), lr0.0) scaler torch.cuda.amp.GradScaler() train_loader, val_loader, train_sampler build_loaders( train_set, val_set, batch_size32, num_workers8 ) for epoch in range(total_epochs): train_sampler.set_epoch(epoch) model.train() for step, (x, y) in enumerate(train_loader): x x.to(local_rank, non_blockingTrue) y y.to(local_rank, non_blockingTrue) optimizer.zero_grad(set_to_noneTrue) with torch.cuda.amp.autocast(): out model(x) loss criterion(out, y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() if rank 0 and step % 50 0: print(fepoch {epoch} step {step} loss {loss.item():.4f}) evaluate(model, val_loader, local_rank, rank) if rank 0: torch.save(model.module.state_dict(), fckpt_e{epoch}.pt) dist.barrier() cleanup_dist()几个容易写错的点DDP包装要在建 optimizer 之前虽然包装后model.parameters()仍然能取到同一批参数但顺序错了有些版本的梯度钩子行为不符合预期torch.manual_seed要在init_process_group之后、建模型之前设置各 rank 用同一个种子保证初始权重一致DDP 也会从 rank 0 广播一次参数所以就算不一致也能被纠正但保持同种子能避免其他随机性带来的不一致保存时一定要用model.module.state_dict()用model.state_dict()会把module.前缀一起存进去加载时对不上。3.3 启动命令与多机参数单机八卡就一行torchrun --nproc_per_node8 --standalone train.py多机两节点、每节点八卡节点 0 上执行torchrun --nnodes2 --nproc_per_node8 --node_rank0 \ --master_addr10.0.0.11 --master_port29500 train.py节点 1 上把--node_rank改成 1其他一致。--master_addr填节点 0 的地址端口挑一个没被占用的29500 是默认值多人共用一台机器时最容易撞的就是它。这里不需要任何环境变量的手动设置torchrun会自动把RANK、WORLD_SIZE、LOCAL_RANK、MASTER_ADDR、MASTER_PORT注入到子进程里你的代码只管读LOCAL_RANK就好。老代码里你可能见过python -m torch.distributed.launch它已经被标记为弃用行为上有几个坑比如默认给每个进程传--local_rank参数需要你在 argparse 里接住新项目一律用torchrun。另外较大的集群上torchrun会有一个 elastic 的容错模式节点挂了可以自动重启但那套东西涉及 rendezvous 后端本地实验用不上先别碰。3.4 DataLoader 与 Sampler 的正确接法数据这块是新手翻车最密集的区域值得多说几句。DistributedSampler做的事是把整个数据集按 rank 切成 world_size 份rank i 取下标i, iW, i2W, ...这些样本。所以每个 epoch 的迭代步数是len(dataset) // W // batch_size。如果你没开drop_lastDistributedSampler在样本数不能被 W 整除时会补齐——重复取前面几个样本让每个 rank 长度一致。这在训练时会导致少量样本被重复训练影响不大但在评估时会重复计入指标算出来的准确率会有一点偏差。所以评估集要么让 sampler 也开drop_lastTrue丢掉尾巴要么在汇总指标时按真实样本数加权。set_epoch这一步千万别漏。DistributedSampler用 epoch 号作为随机种子来打乱顺序如果你不调set_epoch每个 epoch 的样本顺序完全一样等于没打乱。这个错误不会报错也不会让 loss 明显异常但会让模型的泛化轻微变差属于那种“找不到原因但就是差一点”的问题。num_workers要按进程数重新算。单卡时你设 8八个进程加起来就是 64 个 worker 进程CPU 核心不够就会疯狂上下文切换反而更慢。合理值大概是CPU 核心数 // world_size再减去主进程本身的开销。我第一次改的时候忘了这茬八卡跑起来机器直接卡死top一看 load average 到了 200 多。4. 超参与显存必须重新算的那几个数4.1 全局 batch 与学习率的换算DDP 下最容易忘记的一件事你传给 DataLoader 的batch_size是每张卡的不是总的。八卡各 32全局 batch 就是 256。如果你原来是单卡 256改了 DDP 之后 batch_size 还写 32那你的有效 batch 变成原来的八分之一学习率相对就太大了训练会变得不稳定。常用的换算规则是线性缩放新学习率 基准学习率 × (新全局 batch / 基准全局 batch)。原来是单卡 32、lr 1e-3现在八卡全局 256那 lr 应该放到 8e-3。这个规则在 batch 不大、优化器是 SGD 的时候比较可靠换到 Adam 系列、或者 batch 放大超过十倍时线性缩放往往过于激进社区里也有用平方根缩放的放大 8 倍 lr 放 2.83 倍效果介于两者之间。我自己的做法是batch 放大不超过 4 倍时用线性超过就用平方根然后重点看前 500 步的 loss 曲线如果前几百步就抖得像心电图说明放大过头了。如果你不想放大学习率另一个选择是配梯度累积想保持全局 batch 不变就把每卡 batch 减小、累积步数加上去。比如原来单卡 256、lr 1e-3现在八卡各 16、累积 2 步全局还是 256lr 一点都不用改。这种“等价替换”在迁移老实验时最省事缺点是训练变慢了一些累积步的通信仍然存在除非用no_sync。4.2 warmup 在多卡下为什么更必要大 batch 训练的直接后果是梯度估计的方差变小、方向更准这听起来是好事但它意味着你可以也应该用更大的步长而大学习率在训练初期非常危险因为那时候参数离最优点很远梯度方向虽然准但幅度大一步跨过头就发散。warmup 的作用就是在最初几千步里把学习率从一个很小的值线性拉到目标值给优化器一个缓冲期。多卡场景下 warmup 的重要性比单卡高得多因为你的有效学习率往往放大了好几倍。经验值是 warmup 占总步数的 1% 到 5%全局 batch 越大、总步数越少这个比例要往高的方向取。有个具体的现象可以帮你判断如果训练开始后 100 步内 loss 直接从 2.3 跳到 NaN 或者 8 以上八成是学习率放大之后没有配 warmup或者 warmup 的初始值设得太高。4.3 显存账怎么算一张表看清DDP 是数据并行每张卡上都有一份完整的模型副本所以显存不会因为多卡而分摊。这是新手最大的误解之一。显存的主要构成是这样的占用项大致大小说明模型参数P × 4 字节fp32P 为参数个数梯度P × 4 字节反向传播时产生优化器状态AdamP × 8 字节一阶动量 二阶动量混合精度额外的主权重副本P × 4 字节AMP 下的 fp32 主权重激活值取决于 batch、序列长度、层结构通常是最大头拿一个 1 亿参数的模型举例参数 400MB梯度 400MBAdam 状态 800MB加起来 1.6GB这是固定的底座。真正的变量在激活值——它和每卡 batch、输入分辨率、序列长度线性相关。所以八卡上如果你想让每卡 batch 保持和单卡一样显存占用是完全相同的如果 OOM 了问题不在 DDP在 batch 或者别的什么地方。有一类 OOM 很有迷惑性所有卡里只有 0 号卡 OOM。原因通常是 rank 0 上多做了事情——打日志时把 loss 张量留在显存里没释放、在 rank 0 上额外跑了一遍验证、或者保存 checkpoint 时没先挪到 CPU。如果看到 rank 0 显存比别的卡高出好几个 G先查这些地方。4.4 梯度累积与 no_sync省显存但不能省通信梯度累积的写法前面提过逻辑上就是累积 K 步再optimizer.step()。但有一个性能细节值得单拿出来说DDP 在每次backward()时都会触发 AllReduce也就是说累积 4 步就要通信 4 次通信量完全没有减少。前 K-1 步的通信纯属浪费因为它们马上会被后面的梯度叠加上去。DDP 提供了no_sync上下文来跳过这些无用的同步for i, (x, y) in enumerate(loader): is_last (i 1) % accum_steps 0 ctx contextlib.nullcontext() if is_last else model.no_sync() with ctx: with torch.cuda.amp.autocast(): loss criterion(model(x), y) / accum_steps scaler.scale(loss).backward() if is_last: scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_noneTrue)no_sync的作用是让当前的 backward 只在本卡累积梯度不发起通信直到最后一个 micro-batch 才真正同步一次。实测在累积 4 步的场景下这个改动能把 step 时间压掉 15% 到 25%取决于模型大小和互联带宽。要注意no_sync期间的梯度是被累加在.grad上的所以最后一个 micro-batch 同步时AllReduce 平均的是“已经累加了 K 次”的梯度除以accum_steps的写法依然是正确的。5. 踩坑实录DDP 的高频故障与排查路径5.1 卡死类问题从握手到第一个 iteration分布式训练最折磨人的不是报错是卡死。程序不退出日志不再更新八张卡的状态是“有占用但 GPU 利用率 0%”。我把遇到过的卡死分成三个阶段来定位。卡在 init_process_group所有 rank 都不往下走。九成是连接层面的问题——端口被占多人共用机器时 29500 最容易撞、master_addr写错或者用了从别的节点访问不到的地址、多机时防火墙拦了端口。诊断办法是先手动测一下端口通不通或者在启动前换个随机端口。多机上还有一个隐蔽的原因机器有多张网卡进程选错了网卡出去。这种情况下可以显式指定用于通信的网卡名让进程走对的那张。卡在第一个 iteration 之前模型都建完了DataLoader 也建完了但训练循环不开始。常见原因是数据集初始化里有阻塞操作比如某个 rank 卡在读文件上或者某个 rank 上抛了异常但被try/except吞了——DDP 里只要有一个 rank 退出其余 rank 会在下一次集合通信时永远等下去表现出来就是全体卡死。排查时把日志级别调高确保每个 rank 的异常都能打出来别只让 rank 0 打日志。卡在训练中途通常是某个 rank 提前结束了循环。比如你用DistributedSampler但不同 rank 的 loader 长度不一致有的 rank 跑完了 while 别的还在跑或者代码里有if rank 0: ... break这类分支。集合通信是“所有人必须到场”的机制缺一个人全体僵住。这里有个实用技巧给init_process_group设个较短超时比如 10 分钟让卡死变成超时异常至少能看到栈信息告诉你是哪个集合操作在等。另外torch.distributed有一个调试开关可以打开集合通信的日志能看到每一步卡在哪个 AllReduce 上。5.2 精度类问题为什么八卡的结果和单卡对不上这类问题最烦因为程序不报错指标就是差那么一两个点你怀疑是玄学。第一个怀疑对象是 BatchNorm。DDP 是按卡独立做 BN 的每张卡用自己那份数据的均值和方差来归一化。八卡各 32每个 BN 层看到的是 32 个样本的统计量而单卡 batch 256 看到的是 256 个样本的统计量两者的统计量稳定性完全不同。更严重的是验证阶段如果验证时batch_size变得很小BN 的 running stats 会明显偏移。解决办法是换成SyncBatchNorm它在训练时跨卡同步统计量代价是多两次 AllReducemodel torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(local_rank) model DDP(model, device_ids[local_rank])convert_sync_batchnorm必须在 DDP 包装之前调用包完再转有些层替换不到。另外 SyncBN 在小 batch 下每卡小于 8收益不大反而拖慢速度那种情况一般直接用 GroupNorm 或者 LN。第二个怀疑对象是随机性。即使你设了种子DistributedSampler的分片方式决定了每个 rank 看到的数据顺序和单卡完全不同dropout 的随机数也不同。所以多卡和单卡的 loss 曲线不可能逐步对齐只在统计意义上接近。如果你想让结果可复现需要固定种子的同时固定cudnn.deterministic会损失性能并且保证每次运行的卡数一致。第三个怀疑对象是学习率和 batch 的配合前面讲过这里不重复。还有一个冷门的drop_lastFalse导致最后一个 batch 各卡长度不同DDP 平均梯度时权重不对训练末期会有轻微抖动。改成True就好。5.3 find_unused_parameters 的代价与替代方案模型里有参数在某些前向路径上不参与计算多任务模型、带条件分支的模型、检测里的辅助头DDP 默认会在反向时检查所有参数是否收到梯度发现有的参数没收到梯度就报错。打开find_unused_parametersTrue能绕过这个报错但代价是每个 step 都要遍历整张计算图去标记哪些参数被用过实测会带来 10% 到 20% 的额外开销模型越大越明显。我的建议是先把这些参数找出来。思路很简单把 DDP 包去掉单卡跑一个 step然后遍历model.named_parameters()看哪些.grad is None。找到之后有两条路如果这个参数确实不该参与比如只是挂在模型里但逻辑上没用到改成requires_gradFalse如果它是条件分支的产物考虑把分支重构掉或者接受开销。还有一个中间方案是static_graphTrue告诉 DDP“计算图在每个 step 里是一样的”这样它能缓存未使用参数的集合把额外开销降下来。但它有前提图必须真的不变如果某个 step 因为数据不同走了另一个分支行为会出错。5.4 日志、保存与评估的分布式写法日志只在 rank 0 打印。其他 rank 也 print 的话你会看到八份一模一样的输出交错在一起根本没法读。如果要监控每张卡的负载用nvidia-smi或者专门的监控别用 print。保存只在 rank 0 存且一定要存model.module.state_dict()。如果还想保存 optimizer 状态也要只存一份。有一个常见错误是所有 rank 都去写同一个文件路径结果是文件被写坏或者只有最后一个 rank 的内容生效。如果做了分片保存多个文件每个 rank 写自己的分片那就要注意所有 rank 都要有对应的写操作否则 rank 0 在保存时其他 rank 已经跑到下一轮反而可能影响一致性。稳妥起见保存前后加一个dist.barrier()。评估两种写法。一是在 rank 0 上单独跑完整验证集其他 rank 用 barrier 等着简单但浪费了其他卡。二是各 rank 用DistributedSampler(shuffleFalse)分担验证集各自算指标后汇总。第二种效率高但要注意在汇总时用all_reduce求和正确的分子分母而不是简单平均各 rank 的准确率——各 rank 样本数不一样时简单平均会算错。def reduce_metrics(correct, total, device): t torch.tensor([correct, total], dtypetorch.float64, devicedevice) dist.all_reduce(t, opdist.ReduceOp.SUM) return (t[0] / t[1]).item()数据类型要用float64用float32在样本数上千万时会有精度损失导致最后的准确率差个小数点后几位虽然不影响结论但看着膈应。5.5 常见故障速查表症状最可能的原因处理方式全体卡在 init_process_group端口被占、地址不可达、网卡选错换端口、检查 master_addr、显式指定网卡只有 0 号卡显存涨、其他卡不动忘了set_device初始化后立刻绑定 local_rank训练中途全体僵住某个 rank 提前退出循环检查各 rank loader 长度、去掉带 rank 条件的 break报 unused parameter 错误有条件分支或未参与前向的参数找出并冻结或开find_unused_parameters单卡正常、多卡 loss 不降学习率没按 batch 缩放、没加 warmup重算 lr加 1% 到 5% 的 warmup多卡指标比单卡差一截BN 统计量不一致换 SyncBatchNorm或改用 LN/GN加速比只有 1.5 倍DataLoader 瓶颈、num_workers 过多用 profiler 看耗时分布减少 worker 并开 pin_memoryrank 0 OOM日志或保存占着显存打印前.item()保存前挪到 CPU保存的权重加载后 key 对不上存了model.state_dict()改存model.module.state_dict()每个 epoch 数据顺序一样忘了set_epoch在 epoch 循环开头调用6. 性能优化把加速比从 0.6 拉到 0.96.1 先判断你卡在通信还是计算优化之前必须先定位。用最朴素的办法单卡跑 100 步记下平均 step 时间八卡跑 100 步再记一次算出实际加速比再除以 8 得到扩展效率。0.9 以上说明很好0.7 到 0.9 说明有优化空间低于 0.6 说明有明显瓶颈。要进一步定位看两个指标就够了。一是nvidia-smi里各卡的 GPU 利用率如果长期在 30% 到 50% 之间跳动、呈现锯齿状典型的是“算一会儿等一会儿”通信或数据加载在拖后腿。二是用 profiler 抓一段看ncclAllReduce的累计时间和数据加载的时间占比。如果ncclAllReduce占比超过 30%是通信瓶颈如果数据加载占比高是 IO 瓶颈跟 DDP 没关系先解决 DataLoader。我遇到过一个案例八卡加速比只有 2.1nvidia-smi 显示利用率锯齿严重。加上 profiler 之后发现真正的问题在数据侧——num_workers设了 16八卡就是 128 个进程在抢同一块盘的 IO每个 batch 都要等几百毫秒。把 worker 降到 6、把数据集预处理成打包好的二进制格式加速比直接到了 6.9。所以别一上来就怀疑通信库。6.2 通信侧的三个有效手段调 bucket 大小。前面讲过25MB 是默认值。对参数量小、层数多的模型比如小卷积堆叠桶会在很短时间内被填满导致频繁发起通信延迟成为主导这种情况下把bucket_cap_mb调大能减少发起次数。反过来单个参数就超过 25MB 的大模型桶永远是“一个参数一个桶”调大反而会让通信和计算没法重叠得调小或者干脆接受。用gradient_as_bucket_viewTrue。这个选项让梯度直接写到通信桶里省掉一次显存拷贝能省几个百分点到十几个百分点的显存和时间。前提是你的代码没有在 backward 之后、step 之前去手动改.grad如果有自定义的梯度裁剪逻辑要注意它作用在视图上。减少通信总量。混合精度训练时如果用的是 fp16 通信通信量能减半。DDP 在 AMP 下默认通信的是 fp16 梯度前提是梯度本身是 fp16这一块通常是自动的。如果你的模型必须用 fp32 通信可以考虑用bf16训练通信和计算都受益而且不需要 GradScaler稳定性比 fp16 好。6.3 计算侧那些被忽略的单卡优化分布式调了半天没效果有时候问题根本不在分布式。torch.compile在多数视觉和 NLP 模型上能带来 20% 到 40% 的单步提速这个提速在多卡上同样有效而且通常比通信优化的收益大。打开方式很简单包一下模型model torch.compile(model, modedefault)但要注意和 DDP 的组合顺序先torch.compile再 DDP还是先 DDP 再 compile不同版本行为不同推荐先编译再包 DDP并在开启find_unused_parameters时谨慎使用两者的交互在某些版本里会带来额外的图遍历开销。第一次编译会比较慢几十秒到几分钟训练步数少的小实验不一定划算。另一个常被忽略的是set_to_noneTrue。optimizer.zero_grad()默认是把梯度张量清零这会触发一次显存写入设成True之后是把.grad置为None下一次反向时重新分配省掉清零的写入开销在小模型上能省几个百分点的显存和时间。还有一个是关闭 cudnn 的 benchmark 或者反过来打开它——输入尺寸固定时打开能选到更快的卷积算法尺寸变化大时打开会让它反复搜索、反而变慢。6.4 一个真实的调优过程记录最后分享一次完整的调优过程数字我记得比较清楚。任务是某图像分类训练模型规模中等几千万参数八卡同型号单卡 batch 32。阶段操作单步耗时加速比基线单卡410 ms1.0初版 DDP直接改num_workers16490 ms6.7修 workernum_workers6430 ms7.6开 AMPautocast GradScaler260 ms12.6对单卡 fp32 基线开 no_sync累积 2 步275 ms/步等效 550ms/2步相对提升约 6%开 compilemodedefault215 ms15.2要注意这里的加速比已经不能直接和 8 相比了因为中间加了混合精度单卡的耗时基线也变了。正确的对比是“每个配置下分别测单卡和八卡”。但趋势很清楚数据侧的问题修掉带来了最大的一档提升混合精度是第二档通信侧的两个小优化加起来不到 10%。这也印证了我一直的看法——DDP 本身的通信开销在多数中等模型上不是主要矛盾把数据管线和单卡效率做好加速比自然就上来了。多机的时候结论会变。跨节点通信的带宽比机内低一个数量级通信占比会明显上升这时候 bucket 大小、梯度压缩、甚至把 AllReduce 换成分层通信节点内先 reduce、节点间再 reduce就变成主要优化方向了。不过那是另一篇的篇幅单机多卡跑到 0.85 以上的扩展效率之后再考虑跨机不然你只是把一个问题带到了更贵的环境里。我个人在实际操作中的体会是DDP 的代码改造本身不难难的是改完之后那一堆“不报错但结果不对”的状态。我的习惯是每改一步就停下来验一次改完 sampler 先看各 rank 的 batch 数量对不对改完 BN 先看单卡和多卡的验证指标能不能对上改完学习率先跑 200 步看 loss 曲线形状。别一口气改完再跑出了问题你连是哪一步引入的都说不清——这个教训我交了不止一次学费。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑