分布式训练并行策略全解析：DDP、ZeRO、张量并行与流水线并行选型指南

1. 分布式训练到底在解决什么问题先把场景摆出来。你手里有一块24GB显存的卡想训一个7B参数的模型。光是把权重、梯度、优化器状态这三样东西用fp16存下来粗算一下权重14GB梯度14GBAdam的优化器状态一阶动量加二阶动量又是28GB加起来56GB打底还没算激活值。单卡直接爆。这时候你有两条路要么换更大的卡要么把这份“家当”拆开分给多块卡一起扛。分布式训练干的就是后面这件事。但“拆”这个字说起来轻巧拆法不同代价天差地别。最朴素的拆法是数据并行每块卡都存一份完整模型只是喂不同的数据算完梯度再同步。这个方案简单可它有个硬伤模型本身必须塞得进单卡。当模型大到单卡放不下时数据并行就失效了于是才有了后面ZeRO、张量并行、流水线并行、上下文并行这些更细粒度的切分手段。我见过不少刚入门的朋友一上来就纠结“该用哪种并行”其实这个问题问反了。正确的顺序是先问“我的瓶颈在哪”——是显存不够还是通信太慢还是单卡算力吃不满。瓶颈不同选型完全不同。这篇文章我就按这个思路把DDP、ZeRO、张量并行、流水线并行、上下文并行这几条路线挨个拆开讲讲清楚它们各自解决什么、代价是什么、什么场景该上哪个最后再聊聊这几种方案怎么组合起来用。适合谁看如果你已经能跑通单卡训练正准备往多卡扩或者你已经在用DDP但发现显存还是不够、想搞清楚下一步该往哪走那这篇就是写给你的。纯小白也能看我会尽量用生活化的类比把原理讲透但涉及具体配置的地方你得有点PyTorch基础才跟得上。2. 数据并行DDP为什么是分布式训练的起点2.1 DDP的核心机制与它解决的边界DDPDistributedDataParallel是绝大多数人接触分布式训练的第一站也是理解后面所有方案的基石。它的逻辑非常直白假设你有N块卡就把同一个模型完整复制N份每块卡拿一份然后整个训练集切成N份每块卡只吃自己那份数据。前向传播各算各的互不干扰反向传播时每块卡算出自己那份数据产生的梯度然后所有卡把梯度做一次全局平均保证大家更新后的权重完全一致。这里的关键动作是梯度同步。DDP不是等所有梯度都算完才开始通信而是把梯度按参数分桶bucket某个桶的梯度算完就立刻发起AllReduce让通信和计算重叠起来。这个设计很聪明因为反向传播本身就是从后往前逐层算梯度的算完一层就能同步一层不用干等。实测下来只要桶大小调得合适通信开销能被计算掩盖掉大半。那DDP的边界在哪就一句话模型必须完整地放进单卡。因为每块卡都存了一份完整模型包括权重、梯度、优化器状态。前面算过7B模型用Adam训练光这三样就56GB24GB的卡根本放不下。所以DDP适合的是中小模型或者你用了梯度检查点、混合精度这些省显存的手段之后单卡勉强能塞下的场景。一旦模型大到单卡放不下DDP就无能为力了必须换思路。2.2 DDP实操配置与几个容易踩的坑启动DDP最省事的方式是用torchrun不用自己写多进程启动逻辑。一个典型的命令长这样torchrun --nproc_per_node4 --nnodes1 --node_rank0 \ --master_addr127.0.0.1 --master_port29500 \ train.py代码里初始化进程组的写法import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group(backendnccl) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) model MyModel().to(local_rank) model DDP(model, device_ids[local_rank])几个我踩过的坑值得单独拎出来说。第一DataLoader必须用DistributedSampler否则每块卡都会读到全量数据等于白切。第二BatchNorm在DDP下要换成SyncBatchNorm因为普通BN只统计单卡的均值和方差各卡统计量不一致会导致训练不稳定。第三随机种子要按rank设置不然所有卡的数据增强、dropout掩码都一样等于变相减小了batch size。第四保存checkpoint只在rank 0上做其他卡写了会互相覆盖。提示DDP的通信后端在GPU上一定用nccl别用gloo。gloo是给CPU用的在GPU上跑会慢到让你怀疑人生。还有一个隐蔽的坑如果你的模型里有不参与梯度计算的参数比如冻结的embedding记得在DDP里设置find_unused_parametersTrue否则会报错说某些参数没收到梯度。但这个开关会拖慢速度因为它每次都要遍历计算图找未使用参数。更好的做法是提前把这些参数从优化器里排除而不是靠这个开关兜底。3. ZeRO把冗余彻底干掉的三级切分3.1 ZeRO的三级切分逻辑DDP最大的浪费在于冗余。N块卡每块都存一份完整的优化器状态、梯度和权重可这些东西在同步之后是完全一样的。ZeROZero Redundancy Optimizer的思路就是既然大家存的东西一样那就别各存各的了切开分给不同的卡需要的时候再通信凑齐。ZeRO分三个阶段切得越来越狠ZeRO-1只切优化器状态。这是最划算的一刀因为优化器状态占大头Adam是权重的两倍。切完之后每块卡只存1/N的优化器状态显存直接省下一大截通信量增加很少。ZeRO-2在ZeRO-1基础上再切梯度。每块卡只保留自己负责那部分参数对应的梯度其余梯度算完就扔掉。显存进一步下降通信量略有增加。ZeRO-3连权重也切。每块卡只存1/N的权重前向和反向传播时需要哪层参数就临时从其他卡把参数收集过来用完再释放。这个阶段显存省得最狠但通信量也最大因为每次计算都要AllGather参数。用生活化的类比DDP是N个人每人背一整套工具箱ZeRO-1是大家把不常用的工具分开放谁用谁去拿ZeRO-2是连常用工具也分开放ZeRO-3是连手上正在用的工具都只拿零件用的时候现组装。3.2 ZeRO-3的通信代价与适用判断ZeRO-3省显存是真的猛但代价也真实。因为每次前向传播都要把参数从其他卡收集过来反向传播又要再收集一次或者用预取来重叠通信量比DDP高出好几倍。在带宽有限的集群上ZeRO-3可能慢得让你想砸机器。那什么时候该上ZeRO-3我的经验判断是当你用ZeRO-2还是放不下模型时才考虑ZeRO-3。因为ZeRO-2已经切了优化器状态和梯度通常能省下60%以上的显存很多场景到这一步就够了。只有当权重本身都成了瓶颈比如百B级模型才需要动ZeRO-3。配置上DeepSpeed和PyTorch FSDP都实现了ZeRO。DeepSpeed的配置文件里zero_optimization的stage字段就是选级别{ zero_optimization: { stage: 2, offload_optimizer: { device: cpu, pin_memory: true }, allgather_partitions: true, overlap_comm: true, reduce_scatter: true } }这里有个很实用的选项offload_optimizer。它把优化器状态卸载到CPU内存进一步省显存代价是CPU和GPU之间的数据传输。如果你的机器CPU内存够大、PCIe带宽还行这个选项能让你在有限的显存下训更大的模型。我实测过offload之后显存能再降30%左右但每步耗时增加15%到25%值不值取决于你的显存有多紧张。注意overlap_comm和reduce_scatter这两个开关建议都打开它们能让通信和计算重叠、用更高效的通信原语实测能提速10%以上。3.3 ZeRO与DDP的取舍对照维度DDPZeRO-1ZeRO-2ZeRO-3优化器状态每卡全量切分切分切分梯度每卡全量每卡全量切分切分权重每卡全量每卡全量每卡全量切分显存节省基准约4倍约8倍约N倍通信开销基准略增增加显著增加适用场景单卡放得下显存略紧显存紧张单卡放不下这张表我建议存下来选型时对着看。核心判断逻辑是能用DDP就用DDPDDP放不下就上ZeRO-1还放不下上ZeRO-2再放不下才上ZeRO-3。别一上来就ZeRO-3那是杀鸡用牛刀还慢。4. 张量并行把单层矩阵乘法切开4.1 张量并行的切分原理ZeRO解决的是“存不下”的问题但它没有解决“单层算得太慢”或者“单层参数太大”的问题。张量并行Tensor Parallelism切的是模型内部的矩阵运算把一层里的权重矩阵按行或按列切开分到不同卡上算算完再拼起来。拿一个最简单的全连接层举例输出Y XAA是权重矩阵。列并行就是把A按列切成[A1, A2]每块卡算X和Ai的乘积得到Y的一部分最后拼起来。行并行是把A按行切每块卡算一部分输入和对应权重的乘积最后做AllReduce求和。Transformer里最典型的是把多头注意力按头切分每个卡负责几个注意力头FFN层则按中间维度切分。Megatron-LM是这方面的标杆实现它的切分策略至今被广泛借鉴。张量并行的优点是通信量相对可控因为每层算完只需要一次AllReduce或AllGather而且通信量正比于激活值大小不像ZeRO-3那样要传参数。缺点是它要求把模型结构拆得很细实现复杂度高而且对卡间带宽要求极高因为每层都要通信。所以张量并行通常只在单机内NVLink互联用跨机用会被网络带宽拖死。4.2 张量并行的实操要点与通信瓶颈用Megatron-LM或者它的衍生框架时张量并行度tensor parallel size一般设成2、4、8这种且必须整除注意力头数和FFN中间维度。比如模型有32个注意力头你设tp8那每块卡分4个头没问题设tp5就除不尽直接报错。通信瓶颈是张量并行的命门。我做过一个对比测试同样8卡tp8单机NVLink跑和tp2跨机跑前者吞吐是后者的2.5倍以上。原因就是张量并行每层都要AllReduce跨机走网络延迟太高。所以我的建议很明确张量并行尽量限制在单机内跨机交给流水线并行。还有一个细节张量并行下LayerNorm和Dropout这些操作要在切分后的局部做但它们的统计量需要跨卡一致。Megatron的做法是在这些地方插入额外的通信或者干脆把LayerNorm复制到每块卡上算。这个取舍会影响精度和速度调的时候要留意。提示张量并行度不是越大越好。tp增大单卡计算量减小但通信占比上升。存在一个最优值通常在4到8之间。超过8之后通信开销往往吃掉计算收益吞吐不升反降。5. 流水线并行按层切分的时间游戏5.1 流水线并行的基本模型与气泡问题流水线并行Pipeline Parallelism切的是模型的深度把不同的层分到不同的卡上。比如一个32层的模型4块卡每块卡放8层。数据从第一块卡流入逐层往后传最后一块卡输出loss梯度再反向传回来。这个方案听起来很美但有个致命问题气泡bubble。假设4块卡同一时刻只有一块卡在干活其他三块在等——第一块卡算完第一批数据传给第二块第二块才能开始这时候第一块卡闲着。为了让所有卡都忙起来就得把一个大batch切成多个micro-batch让它们像流水线一样错峰流动。但即使这样流水线的头尾仍然会有填充和排空阶段这段时间卡是闲置的这就是气泡。气泡的比例大致是(P-1)/(MP-1)P是流水线级数M是micro-batch数量。想让气泡小要么减少P要么增大M。但M增大会增加显存占用因为要同时保存多个micro-batch的激活值所以这是个权衡。5.2 1F1B调度与micro-batch调优为了压气泡业界提出了各种调度策略最经典的是1F1BOne Forward One Backward。它的做法是预热阶段先做P-1次前向填满流水线然后进入稳态每做一次前向就做一次反向保持流水线满载最后排空阶段把剩下的反向做完。相比朴素的“全前向再全反向”1F1B能把激活值显存占用降到最低因为同一时刻只需要保存P个micro-batch的激活。调micro-batch数量是个经验活。太少气泡大太多显存爆。我的做法是先按显存上限估一个M然后逐步减小看吞吐变化找到拐点。通常M在8到32之间比较常见。流水线并行的优点是通信量小因为只在层与层之间传激活值而且可以跨机用对带宽要求比张量并行低得多。缺点是气泡带来的效率损失以及实现上的调度复杂度。PyTorch的Pipe、DeepSpeed的Pipeline、Megatron的流水线都是成熟实现选一个用就行别自己造轮子。并行方式切分对象通信频率通信量跨机友好度数据并行数据每步一次梯度大小好ZeRO状态/梯度/权重每步多次较大好张量并行层内矩阵每层多次激活值差流水线并行层间每micro-batch激活值好这张表能帮你快速判断跨机部署优先考虑数据并行、ZeRO、流水线并行单机内追求极致可以上张量并行。6. 上下文并行长序列训练的新解法6.1 上下文并行要解决的长序列显存问题前面几种并行方式切的是模型参数或数据但有一类场景它们都搞不定超长序列。当序列长度拉到32K、128K甚至更长时注意力机制的显存占用是随序列长度平方增长的。这时候哪怕模型不大光是激活值就能把显存撑爆。上下文并行Context Parallelism切的就是序列维度。它把一条长序列切成若干段每块卡负责一段注意力计算时通过环形通信ring attention在卡间传递KV让每块卡都能看到完整的上下文。这样显存占用随卡数线性下降序列就能拉得更长。这个思路和前面几种并行是正交的可以叠加使用。比如你可以同时用张量并行切模型、上下文并行切序列、数据并行切batch三者互不冲突。这也是为什么现在训长上下文模型基本都是多种并行组合着上。6.2 环形注意力的通信模式与实现选择环形注意力的核心是每块卡持有自己那段序列的Q、K、V然后KV沿着环形在卡间流转。第i步时每块卡用当前的KV块算一部分注意力算完把KV传给下一块卡同时接收上一块卡传来的KV。转一圈下来每块卡都见过了所有KV注意力就完整了。这个通信模式的好处是通信量正比于序列长度和卡数而不是平方。但实现起来对通信和计算的overlap要求很高因为KV传递和注意力计算要尽量重叠否则通信就会成为瓶颈。目前支持上下文并行的框架不算多Megatron-LM、DeepSpeed、以及一些专门的长上下文训练库有实现。选型时要注意框架对注意力实现的兼容性比如FlashAttention和环形注意力能不能配合这个得看具体版本。注意上下文并行对序列切分的位置有讲究。如果序列里有跨段的依赖比如某些位置编码切分时要保证这些依赖不被破坏否则精度会掉。实践中通常按固定长度切并在边界处做特殊处理。7. 组合并行3D并行的实战配置思路7.1 3D并行的维度分配逻辑真实的大模型训练几乎不会只用一种并行。常见的是3D并行数据并行乘以张量并行乘以流水线并行。比如一个千亿模型可能配置成tp8、pp8、dp16总共1024块卡。维度怎么分配我的经验是按这个优先级来先看单卡显存能放下多少用张量并行和流水线并行把模型切开直到单卡放得下剩下的卡全部拿去做数据并行。因为数据并行的通信效率最高、实现最简单能用数据并行解决的就别用更复杂的。具体到数字张量并行度受限于单机卡数和NVLink拓扑一般不超过8流水线并行度受限于层数和气泡容忍度一般不超过16数据并行度就是剩下的卡数。三者相乘等于总卡数。7.2 通信组划分与配置落地组合并行最容易出错的地方是通信组的划分。不同的并行维度需要不同的通信组张量并行组内做AllReduce流水线并行组间传激活数据并行组内同步梯度。这些组必须正确初始化否则要么通信错乱要么死锁。以Megatron-LM为例它的初始化逻辑是先按张量并行分组再按流水线并行分组最后按数据并行分组。每块卡根据自己的rank计算出自己在各个维度上的坐标然后加入对应的通信组。这个逻辑框架帮你封装好了你只需要传对参数torchrun --nproc_per_node8 --nnodes16 \ pretrain_gpt.py \ --tensor-model-parallel-size 8 \ --pipeline-model-parallel-size 8 \ --num-layers 64 \ --hidden-size 8192 \ --num-attention-heads 64这里tp8、pp8单机8卡正好一个张量并行组16台机器组成8个流水线级数据并行度就是16/82。配置的时候要保证num-layers能被pp整除num-attention-heads和hidden-size能被tp整除否则直接报错。提示组合并行下全局batch size等于micro-batch乘以数据并行度乘以梯度累积步数。调参时这三个量要一起看别只盯着一个改。8. 常见问题与排查技巧实录8.1 显存与通信类问题速查分布式训练报错八成逃不出显存和通信这两类。我整理了一张速查表遇到问题先对着查现象可能原因排查方向OOM但单卡能跑并行配置不当检查ZeRO级别、offload是否开启训练卡住不动通信死锁检查各卡是否都进入了同一个集合通信吞吐远低于预期通信瓶颈用profiler看通信占比调整并行度loss不收敛梯度同步错误检查DDP的find_unused_parameters各卡loss不一致随机种子未按rank设置检查seed和sampler配置NCCL超时网络或拓扑问题检查NCCL环境变量和网卡绑定8.2 几个我踩过的深坑第一个坑NCCL超时。默认NCCL的超时是30分钟但有些集合通信在特定拓扑下会卡很久导致超时中断。解决办法是设NCCL_TIMEOUT环境变量调大或者用TORCH_NCCL_ASYNC_ERROR_HANDLING让错误及时抛出而不是干等。第二个坑显存碎片。长时间训练后显存碎片会导致明明还有空闲显存却OOM。设PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True能缓解这个问题它让分配器用可扩展的段来管理显存减少碎片。第三个坑checkpoint不兼容。用不同并行配置保存和加载checkpoint时参数切分方式不一样直接加载会错位。这时候需要用框架提供的转换工具把checkpoint从一种并行布局转成另一种。Megatron和DeepSpeed都有这类工具别自己手写转换脚本容易出错。第四个坑数据加载成为瓶颈。当计算很快时DataLoader的读取速度可能跟不上GPU在等数据。解决办法是增大num_workers、用pin_memory、把数据预处理做成流式。我见过一个案例光是把num_workers从4调到16吞吐就涨了20%。8.3 性能调优的实操顺序调优别瞎调按这个顺序来先确认单卡能跑满再上多卡看扩展效率然后逐步加并行维度。每一步都用profiler记录找到真正的瓶颈再动手。我常用的profiler是PyTorch自带的torch.profiler能看到每个算子的耗时和通信占比。扩展效率的计算很简单N卡吞吐除以单卡吞吐再除以N。如果这个值低于0.7说明扩展性有问题得查通信。低于0.5基本就是通信瓶颈要么调并行度要么换更快的互联。最后分享一个我个人的判断准则任何并行方案先在小规模上验证正确性和扩展效率再往大规模推。我见过太多人一上来就上几百卡结果配置错了跑了一天发现loss是nan白白烧钱。小规模跑通、扩展效率达标再放大这是最省时间的路径。