★ FEATURED ARTICLE 做深度学习的人基本都会经历一个阶段单卡训练跑得动但越往后越不对劲。模型从几亿参数涨到几十亿数据量翻了十倍单卡训练一次要等一周实验迭代慢到让人怀疑人生。这时候你听到最多的一句话就是“上分布式训练吧”——Pytorch里的分布式训练听起来高大上实际上却是每个做大规模训练的工程师迟早要跨过的一道坎。这篇内容不打算给你背一遍官方文档而是把分布式训练从“为什么要做”到“怎么落地”再到“踩了哪些坑怎么填”完整串一遍。包括DDP怎么改、torchrun怎么用、环境变量怎么理解、BatchNorm为什么在多卡下会“失灵”、断点续训为什么总报错这些实际问题。适合正好被训练时长折磨、准备上手多卡训练的工程师和研究者也适合那些已经能跑通但想搞明白“为什么多卡反而更慢”的人。1. 分布式训练的前置认知先想明白要解决的问题1.1 单卡训练到底慢在哪先别急着写代码我们先拆一下单卡训练的时间都花在哪了。一次完整的训练迭代大致可以分成这么几块前向传播算损失、反向传播算梯度、优化器更新参数、数据加载。前两项是纯计算吃的是GPU算力数据加载吃的是CPU和磁盘IO参数更新虽然也有开销但相比前两项通常小到可以忽略。单卡训练慢本质上不是某一个环节慢而是这几个环节串在一起谁慢谁就卡住整个流程。很多人第一步想到的优化是把数据加载改异步、把CPU换快一点这些都是局部优化。当你把模型的batch size开到显存极限数据加载也提前用DataLoader的num_workers铺满了仍然训练一周才能收敛那剩下的瓶颈就是GPU算力本身。想突破这张卡的天花板就得把活分给多张卡干这就是分布式训练上场的原因。1.2 同步训练与异步训练两种“开会”方式“把活分给多张卡”听起来简单但卡和卡之间怎么配合是个绕不开的问题。Pytorch里的多卡训练大体上分同步和异步两种思路。同步训练好理解每张卡各自拿一批数据算梯度算完之后大家把梯度汇总到一起所有人拿到一个统一的全局梯度再同时更新模型参数。整个过程就像开会所有人必须等最后一个人发言完毕才能散会。优势是收敛行为稳定等价于用一个更大的batch size做训练缺点是同步等待会带来通信开销。异步训练则是各卡算完自己的梯度之后不等人直接更新参数。节奏快但每一张卡看到的参数可能是别人已经更新过的收敛过程容易抖动调起来对经验和策略要求都高。工业界里绝大多数训练任务用的都是同步模式Pytorch的DistributedDataParallel默认也是同步更新。新手不要一上来就想着异步先把同步模式吃透大部分团队的训练任务其实同步就够了。1.3 数据并行与模型并行两种切分思路另一个绕不开的概念是数据并行和模型并行。很多人容易搞混其实一句话就能说清楚数据并行是每张卡放一份完整的模型各自处理不同的数据然后合并梯度模型并行是模型大到一张卡放不下把模型的不同层拆开放在不同的卡上数据依次流过。实际工作中90%以上的分布式训练需求都属于数据并行。原因很简单大部分模型单卡放得下用数据并行就能解决算力不够的问题代码改动还小。模型并行几乎只在模型大到单卡显存装不下时才需要考虑而且实现复杂度高得多属于“卡脖子才用”的方案。所以这篇内容主线就放在数据并行上模型并行只在概念层面讲清楚避免你在技术选型时走弯路。1.4 一张卡算力翻倍不等于速度翻倍通信代价的直觉为什么有人上了8卡训练速度不但没变成8倍反而比单卡还慢这里的核心原因是多卡并行之后“通信”成了一种新的开销。我拿一个直观的算式说明。假设你的模型有1亿参数反向传播后每张卡会产生一份梯度用单精度存储大约是400MB。8张卡做数据并行每张卡算完自己的梯度之后要和其他7张卡把梯度同步到一个平均值这个操作在分布式系统里叫all-reduce一次all-reduce过程中每张卡实际发送的数据量差不多是自身梯度大小的两倍也就是800MB。所以每迭代一次8张卡之间就要流动6GB以上的网络数据。如果数据加载和计算本身只要1秒网络同步又花掉0.5秒那8卡带来的计算加速就会被通信吃掉一大截。实际表现可能就是8卡比4卡只快了1.5倍极端情况下甚至倒挂。记住这个结论分布式训练优化的大方向永远是“减少通信量”和“让通信和计算重叠起来”后面所有实操技巧都是围绕这两点展开的。2. 环境规划与通信后端选型决定分布式上限的隐性地基2.1 单机多卡与多机多卡的硬件规划动手写代码前先看看你的硬件不同硬件拓扑下的分布式训练策略和预期收益完全不一样。单机多卡指的是同一台机器上插了多张GPU卡与卡之间的数据交换走的是主机内部的GPU互连总线比如PCIe或者更高速的NVLink特点是带宽高、延迟低通信开销相对小这是最常见的省钱又省力的方案。多机多卡则是多台机器通过网络连接常见的是万兆以太网或InfiniBand。跨机器的通信延迟和带宽都比单机内部差一个量级所以多机多卡对通信优化的要求高得多不然训练根本跑不起来。规划阶段记住几个原则优先把卡集中到尽量少的机器上同一台机器上的卡越多单机内部的通信越高效如果不得不多机训练机器之间的网络设计不能省网卡带宽和交换机容量都是硬瓶颈。我在实际项目里见过不少团队机器之间用千兆网跑分布式训练结果GP利用率常年不到30%代码怎么调都没用换了万兆网络立刻跑满。真就是“地基不牢上面白折腾”。2.2 Pytorch后端选型NCCL、Gloo与MPIPytorch的分布式训练有一个参数叫backend中文一般叫“通信后端”它决定了各进程之间用什么方式交换数据和梯度。官方常用的有三个NCCL、Gloo、MPI。NCCL是GPU场景下的首选。它专为多卡通信设计支持高带宽的GPU直连比如通过NVLink在GPU之间免去CPU中转对all-reduce这类集合通信做了深度优化。只要你的训练跑在GPU上就别犹豫选NCCL。Gloo更通用既能跑CPU也能跑GPU但性能不如NCCL通常拿来兜底比如某些调试场景下NCCL连不上先用Gloo验证逻辑。MPI是高性能计算领域的传统方案灵活但配置成本高Pytorch场景下一般用不上。一句话总结用GPU训练无脑选NCCL实在初始化失败再临时切Gloo排查问题。这个选择判断比绝大多数人想的要简单。2.3 关键环境变量RANK、LOCAL_RANK、WORLD_SIZE一次性讲透刚接触Pytorch分布式训练的人有一半时间是在跟RANK、LOCAL_RANK、WORLD_SIZE这三个环境变量搏斗。它们的含义不复杂但搞混的话代码一定跑不通。WORLD_SIZE整个训练任务一共有多少个进程参与。RANK当前进程在整个任务里的全局编号从0到WORLD_SIZE-1。LOCAL_RANK当前进程在“本台机器”上的编号。多机场景下LOCAL_RANK的“局部性”尤其容易让人晕。举个具体例子4台机器、每台8张卡一共32个进程。第一台机器上的8个进程RANK是0到7LOCAL_RANK也是0到7第二台机器上的8个进程RANK是8到15但LOCAL_RANK又变回0到7。模型加载的Device ID必须用LOCAL_RANK而日志打印、全局编号相关的逻辑用RANK。把这两个概念分清楚多机环境就不会手足无措。好在Pytorch新版本推荐使用的torchrun启动器会自动设置这些环境变量不需要手工export。但就算如此你也必须知道它们的含义因为后续的日志输出、数据切分、checkpoint保存逻辑都会用到。3. Pytorch分布式训练核心API拆解与代码改造3.1 从单卡训练到DDP的四步改造DistributedDataParallel简称DDP是Pytorch官方主推的数据并行方案它的改造套路非常固定我基本已经写成了一套模板。第一步初始化进程组。所有参与训练的进程必须先“握手”建立连接靠init_process_group完成。关键参数是backend和init_methodbackend用ncclinit_method最常见的是env://也就是让Pytorch从环境变量里读取上面说的RANK、WORLD_SIZE等信息。import torch import torch.distributed as dist dist.init_process_group(backendnccl, init_methodenv://)第二步把模型包上DDP。注意先设置好当前进程使用的GPU设备再把模型搬到对应设备上最后才交给DDP包装。local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) model model.to(local_rank) model torch.nn.parallel.DistributedDataParallel( model, device_ids[local_rank] )第三步改造DataLoader。这里一定要引入DistributedSampler否则每张卡取到的是同一批数据训练就白并行。from torch.utils.data import DistributedSampler sampler DistributedSampler(dataset) dataloader torch.utils.data.DataLoader( dataset, batch_sizebatch_size, samplersampler, num_workers4, )第四步改造训练循环头部每个epoch开始时给sampler打一个“换一轮随机顺序”的标记。for epoch in range(epochs): sampler.set_epoch(epoch) for data, target in dataloader: ...这套模板基本覆盖90%的DDP需求。改完就跑跑起来再考虑优化。3.2 启动方式的正确姿势torchrun与spawn代码改完之后怎么把多进程启动起来是另一个高发踩坑区。老教程会让你用python -m torch.distributed.launch这东西能用但Pytorch官方后来推了torchrun好处是它会自动处理异常重启、环境变量注入这些脏活。比如你想在一台8卡机器上起8个进程一条命令就够了torchrun --nproc_per_node8 --nnodes1 train.py想跑多机加一个节点列表torchrun --nnodes2 --nproc_per_node8 \ --rdzv_endpoint主节点IP:端口 \ train.pytorchrun启动后每个进程都会拿到RANK和LOCAL_RANK环境变量你只需要在代码里读取它们。我自己用过一段时间torch.distributed.launch后来全面切到torchrun理由很简单单机调试时它能帮你在进程崩溃后自动重启省得反复手动拉起多机时它的动态发现机制也更好用不用手动配置一堆互信和IP列表。除了命令行Pytorch还提供torch.multiprocessing.spawn在代码内部起进程。这种方式的好处是写进脚本就能跑不用额外命令适合做小规模验证坏处是嵌套太多时日志容易乱而且不如torchrun灵活。我的建议是正式训练一律用torchrunspawn只拿来写Demo。3.3 DistributedSampler的隐藏细节shuffle、epoch和断点续训DistributedSampler看起来就是往DataLoader里塞了个参数实际上坑最多。第一个坑是shuffle行为。如果你传了shuffleTrueDistributedSampler会在内部用torch.randperm生成一份随机索引。问题是这份随机索引是基于当前epoch生成的而DistributedSampler只有在被调用set_epoch的时候才会更新这份索引。如果你忘了在每个epoch开头调用set_epoch那每一轮训练拿到的样本顺序完全一样模型泛化能力会受影响。第二个坑跟断点续训有关。训练中断后恢复通常要恢复三样东西模型参数、优化器状态、当前epoch。但在分布式场景下还要把DistributedSampler的epoch状态一起恢复。因为Pytorch的随机数生成器在每个进程里都有独立状态如果不恢复epoch并重新set_epoch恢复后的数据顺序和原来对不上可能造成某些样本被重复训练或漏训练。实际做法是在保存checkpoint时把epoch和随机种子一起写进去恢复时先加载这些状态再set_epoch。这一步不能偷懒否则你在分布式下的断点续训就是薛定谔的续训看起来跑起来了实验结果却对不上。4. 一次真实的训练提速实践从单机四卡到多机扩展4.1 改造前的基线一个模型的训练困境光讲API容易飘我拿一个实际接触过的模拟项目来走一遍完整流程。前阵子帮一个做视觉模型训练的团队排查训练效率他们的网络模型结构不算大参数量大概8000万单卡训练一个epoch要接近40分钟一个完整实验要跑若干个epoch迭代效率很低。他们的原始代码是标准的单卡写法一个for循环遍历DataLoader模型简单包了个torch.nn.Module没做任何分布式处理。机器上其实有4张卡但一直躺着没用。我的第一个建议就是先上单机的4卡DDP别急着考虑多机。因为单机内部通信带宽高改动成本又低往往能解决大部分算力不足问题。4.2 分阶段改造与踩坑记录改造第一步我先做了最小化验证只改进程初始化和DDP包装先确保多进程能跑通、每张卡的GPU利用率能起来。结果第一次试跑就出了问题4张卡的数据Loader完全一致每张卡在计算一模一样的梯度。原因就是忘了加DistributedSampler加完之后数据开始分散loss曲线也正常了。第二步踩的是BN的坑。这个模型结构里有多个BatchNorm层4卡训练后在前几个epoch里验证集上出现了“loss下降但准确率反而比单卡低”的情况。排查下来是BatchNorm在DDP下统计的是单卡batch内的均值和方差而不是全局batch的。解决办法是等模型训练稳定后引入同步BatchNorm后面会专门讲这个问题。第三步是把学习率和batch size联动调整。原来单卡batch size是644卡并行后每张卡维持64全局batch变成256按线性缩放规则把学习率从0.001提到0.0025附近并加了一段warmup。这一步做完训练收敛速度明显加快loss曲线也不抖了。4.3 提速结果怎么看加速比、吞吐与GPU利用率改造结束后不能只看训练总时长变短就交差。我习惯用三个指标衡量上分布式到底值不值单卡吞吐量、整体加速比、GPU利用率。单卡吞吐量就是每秒每张卡能处理多少样本反映的是GPU有没有吃饱。如果一张卡上的吞吐量在训练过程中只有单卡基线的一半说明数据加载、通信和计算没有做好重叠卡在空等。整体加速比是总训练时间相对单卡训练的倍数。以这个模拟项目为例4卡DDP改造后加速比大约3.6倍左右没有到理论上的4倍原因就是all-reduce和负载不均的额外开销。GPU利用率则用常规监控工具看实用主义的标准就是尽量让利用率维持在90%以上的水平如果长期低于70%说明通信等待占了太大比重需要做性能分析。上线多机训练之前我先用单机4卡把代码、超参、日志体系都稳定下来。这个策略我一直很推荐先小规模把分布式训练的全部逻辑跑通再放大到更多卡问题定位成本低得多。5. Pytorch分布式训练常见问题与排查技巧实录5.1 初始化卡死与通信超时新手遇到最多的分布式训练问题是进程启动后卡在init_process_group不动或者跑着跑着报NCCL超时。这类问题绝大多数跟网络和端口有关代码本身的概率反而小。排查顺序我一般这样来先确认所有参与训练的进程都能互相访问多机场景先ping一下主机IP再确认通信端口没被防火墙拦掉然后看所有节点的GPU数量是不是和启动参数匹配。如果单机多卡优先检查是不是有人手动设置了CUDA_VISIBLE_DEVICES导致各进程看到的设备编号不一致多机场景检查一下各节点的torch版本和GPU驱动是否一致。版本不一致引发的奇怪报错大厂的工单系统里能拉出一整排。还有一个容易掉坑的点是NCCL的通信超时时间太短。Pytorch新版本里init_process_group支持一个timeout参数单位是秒默认是30分钟。碰上模型很大、梯度同步耗时很长的训练30分钟可能不够可以在初始化时显式调大例如设成3600秒。这不是逃避问题而是避免任务跑到一半被误杀。5.2 BatchNorm在DDP下的统计偏移BatchNorm大概是DDP里最阴险的陷阱它在单卡上表现良好一上多卡性能反倒变差而且特别隐蔽。BatchNorm的均值、方差原本是在一个batch内统计的。单卡训练时batch size如果是128那统计的就是128个样本的分布。到了DDP数据并行每张卡只看到自己那部分数据假设4卡、每卡batch为32那么每一张卡上BatchNorm统计的其实只有32个样本的分布而不是你以为的128。这样统计出来的均值和方差偏差很大尤其在目标检测、语义分割这类自带大量小目标样本的任务里表现会明显退化。解决办法是使用同步BatchNorm也就是SyncBatchNorm。它会在训练时把各卡的均值和方差做一次跨卡同步让统计量基于全局batch。Pytorch里改起来很省事把模型里的BatchNorm层替换成SyncBatchNorm即可。model torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)但这也不是免费的每一次BN层的前向传播都会引入一次跨进程通信模型里BN层多的话通信量不可忽视。所以我的建议是小batch训练时先看是不是真的出现了性能退化再决定要不要上SyncBatchNorm如果单卡batch本身就比较大统计量不太失真就尽量别开省下的带宽留给数据交换。5.3 梯度累积与梯度裁剪的坑分布式场景下做梯度累积很多人的代码是“看起来对实际不对”。DDP默认每个backward都会触发一次all-reduce梯度同步而梯度累积是想攒够多个batch的梯度再更新参数。如果只是简单地把多个backward累加再step就相当于每次反向都做了一次无效的all-reduce通信浪费巨大。正确做法是用no_sync上下文管理器把非更新步的backward包起来只在真正更新参数的那一步才做梯度同步。for data, target in dataloader: if acc_step % accumulation_steps ! 0: with model.no_sync(): loss model(data, target) loss.backward() else: loss model(data, target) loss.backward() optimizer.step() optimizer.zero_grad()梯度裁剪也有版本讲究。单卡下clip_grad_norm_直接对grad_norm做裁剪分布式下要特别注意DDP在all-reduce之后各卡上收到的是全局梯度理论上每张卡的梯度都是一样的此时裁剪没有任何问题。但如果你对每张卡的梯度只做了局部裁剪再同步同步后的梯度就和“裁剪后”的梯度不一致了模型更新方向就会偏。所以梯度裁剪一定要放在backward之后、optimizer.step之前并且是全局同步完成之后执行。5.4 模型保存与断点续训的正确姿势分布式训练下保存模型最常见的错误是每张卡都往同一个路径写checkpoint导致文件互相覆盖或者被写坏。正确做法是只让编号为0的主进程负责保存。if dist.get_rank() 0: torch.save({ model: model.state_dict(), optimizer: optimizer.state_dict(), epoch: epoch, sampler_epoch: sampler_epoch, }, checkpoint.pth)加载的时候同样只让主进程load完之后再把参数广播给所有进程避免每个进程单独读文件带来的内存峰值和文件句柄冲突。另一个很隐晦的坑是模型里有些buffer或state是“非确定性”的比如BatchNorm的running_meanDDP会把主进程上的buffer自动广播一次但如果你手动保存再加载最好统一在一开始就用torch.distributed.broadcast对所有参数做一次对齐省得后续莫名其妙的不一致。还有一点加载模型权重时先把checkpoint里的state_dict迁到CPU再往每张卡的GPU上搬能省下大量显存临时开销。别问为什么踩过的人都知道load大模型时显存是怎么被撑爆的。5.5 Batch Size放大后学习率怎么调把batch size从256放大到2048学习率还维持原来的0.1大概率loss直接飞掉。分布式训练里最常用的调法叫线性缩放规则把batch size放大K倍学习率也大致放大K倍同时配上足够长的warmup。举个例子原方案batch256、lr0.1、训练100个epoch改成batch2048后lr先设成0.8然后从0.08开始warmup在前几个epoch线性涨到0.8再按原来的衰减计划走。为什么需要warmup因为大batch下模型一开始的参数离收敛方向太远直接用大学习率容易冲过头先小步走一段等优化方向稳定了再放开学率。这只是一个经验起点不代表一定最优但比“batch翻倍学习率不动”靠谱得多。如果不想动这个脑筋也可以直接保持每卡batch不变让总batch随卡数放大再通过控制epoch总数保证训练总量一致很多大团队的视觉任务就是这么跑的。6. 从“能跑”到“好用”性能调优与调试经验6.1 把通信藏进计算里分布式训练的终极优化方向是“通信与计算重叠”。通俗讲梯度同步不需要等所有卡都算完才开始理想状态是每张卡算完一部分梯度就先同步一部分同时接着算后面的层让网络在计算的空档里悄悄把数据传完。对大部分使用者来说Pytorch的DDP已经内置了相当不错的通信重叠逻辑默认情况下它会尽量把梯度的all-reduce和后续的backward计算重叠。你要做的是别“手贱”打断它。实测中很多性能退化都是因为代码里强行加了一些不必要的同步点比如频繁打印loss时调用all_gather或者每个iteration都做一次torch.cuda.synchronize()这些都会把并行流程硬生生掰回串行。另一种见效明显的技巧是梯度压缩。通信量正比于梯度大小如果能用低精度格式比如fp16传梯度通信量直接减半。具体做法是在构造进程组时打开NCCL的梯度压缩开关或者通过混合精度训练AMP让反向传播和同步都跑在低精度上。Pytorch的自动混合精度开启后DDP会使用fp16梯度通信多卡场景下提速非常可观。6.2 定位性能瓶颈的简单方法出现“多卡还不如单卡快”的诡异情况先别急着改代码。我习惯用最简单的方式定位在训练循环里手动记录每个关键阶段的时间戳。大致思路是每个iteration里分别计时数据加载耗时、前向计算耗时、反向计算耗时、DDP等待耗时。如果DDP等待耗时占了大头说明各卡计算负载不均衡或网络通信太慢如果数据加载耗时过高说明DataLoader还没喂饱GPU如果前向反向本身就很慢那问题不在分布式。更精细的分析可以配合profiler工具但那是后话。日常我基本只用两种信息一看GPU利用率像不像锯齿状起伏二看单iteration耗时里的communication占比。前者可以分辨数据加载瓶颈后者可以直接分辨通信瓶颈。这套方法不需要读太多工具报告几分钟就能看出问题方向。6.3 一些我反复用到的经验贴士最后分享几个在多个项目里反复验证过的实操习惯都是文档里很少写的“土办法”但很管用。第一个正式跑多卡之前先用单机双卡跑通全流程。双卡能暴露绝大多数分布式逻辑bug忘加sampler、忘记set_epoch、checkpoint保存冲突这些问题在双卡下就能看出来排查成本比8卡低得多。第二个日志输出一定要带RANK前缀。分布式下几十个进程同时打印不带RANK的日志根本没法看。主进程负责打印训练指标其他进程只在出现错误时才被允许输出这是降低排障痛苦的有效手段。第三个多机训练之前先确认共享存储。分布式训练最好有一个所有节点都能读写的共享数据目录用来放checkpoint和数据集。没有共享存储的话每台机器都要维护一份数据副本一旦数据版本没对齐训练结果就是灾难。第四个也是我自己在实践里最深的一条体会分布式训练不是“把单卡代码跑成多卡”就结束了它真正的难度在于你愿不愿意把通信、调度、负载均衡这些底层因素纳入思考。代码能跑只是起点每一层优化都是在跟机器和网络“讨价还价”。踩过几次坑之后我现在接手一个训练任务第一件事永远是先问数据多大、模型多大、卡在哪、网络多快这四个答案基本决定了分布式方案的天花板。剩下的就是照着这篇内容的套路一步步落地的过程。
阅读完成 · 觉得有帮助？
