ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

算法同学必读:LLM分布式并行策略TP/DP/PP/CP/EP详解

算法同学必读:LLM分布式并行策略TP/DP/PP/CP/EP详解 1. 为什么算法同学必须搞懂分布式并行策略如果你是一个算法同学平时的工作重心在模型结构设计、损失函数调优、数据清洗和实验对比上突然有一天发现单卡显存装不下模型了或者训练速度慢到无法忍受这时候你就必须面对一个现实问题分布式计算。LLM时代的模型参数量从几亿飙升到几千亿甚至万亿单张GPU的显存和算力早就扛不住了。你不可能永远靠“换更大显存的卡”来解决问题因为硬件的发展速度远远跟不上模型膨胀的速度。分布式计算的核心目标就两个把模型拆开放到多张卡上以及让多张卡协同干活。但怎么拆、怎么协同这里面有非常多的门道。TP、DP、PP、CP、EP这五个缩写就是目前LLM训练和推理中最主流的五种并行策略。你不需要成为Infra专家但你必须理解它们各自解决什么问题、代价是什么、什么时候该用哪个。否则你跟Infra同学沟通时对方说“我们这次用TP8、PP4、DP2”你只能点头假装听懂了实际上心里一片茫然。这篇文章就是写给算法同学的Infra入门指南。我会用最直白的方式把TP、DP、PP、CP、EP这五个概念讲清楚包括它们的工作原理、通信开销、适用场景以及在实际项目中怎么组合使用。读完你至少能做到看到并行配置不懵能判断某个方案是否合理能跟Infra同学进行有效沟通。提示本文假设你已经了解Transformer的基本结构知道什么是Attention、FFN、LayerNorm。如果你对这些还不熟悉建议先补一下基础。2. 五种并行策略的核心原理拆解2.1 DP数据并行最直观的加速方式DPData Parallelism是最好理解的一种并行方式。它的逻辑非常简单每张卡上都放一份完整的模型副本然后把训练数据切成N份每张卡用自己的那份数据独立计算梯度最后把所有卡的梯度汇总求平均再统一更新模型参数。打个比方就像五个学生做同一套试卷的不同部分做完之后互相核对答案取平均分作为最终成绩然后每个人根据这个平均分来调整自己的学习方法。每张卡上的模型是完全一样的只是看到的数据不同。DP的优势在于实现简单、通信模式清晰。PyTorch的DistributedDataParallelDDP就是干这个的基本上你只要把模型包一层DDP启动脚本改一下就能跑起来。对于中小规模模型比如参数量在10B以下DP通常是最先被考虑的方案。但DP有几个硬伤。第一每张卡都要存一份完整模型包括参数、梯度、优化器状态。以一个7B模型为例FP16精度下参数占14GB梯度占14GBAdam优化器的两个动量状态各占28GBFP32加起来就是84GB。一张80GB的A100根本装不下更别说还要留空间给激活值。第二通信量大。每次迭代结束都要做一次全局梯度同步AllReduce参数量越大通信时间越长。当模型大到一定程度通信开销会吃掉大部分加速收益。所以DP适合的场景是模型能单卡装下或者用ZeRO等优化手段后能装下主要瓶颈在数据量太大、训练太慢。这时候加卡做DP线性加速比还是比较理想的。2.2 TP张量并行把矩阵乘法切开TPTensor Parallelism的思路是把模型中的大矩阵运算拆到多张卡上。Transformer里最耗显存和算力的就是那些大矩阵乘法比如Attention里的QKV投影、FFN里的两层线性变换。TP就是把这些矩阵按行或按列切开每张卡只算一部分最后再把结果拼起来。以FFN为例假设第一层线性变换是 Y XA其中A是[h, 4h]的权重矩阵。如果TP2我们可以把A按列切成两半A1和A2各是[h, 2h]分别放到两张卡上。每张卡计算Y1 XA1和Y2 XA2得到两个[*, 2h]的结果。然后第二层线性变换B是[4h, h]按行切成B1和B2各是[2h, h]。每张卡计算Z1 Y1B1和Z2 Y2B2最后把Z1和Z2相加得到最终结果。这样每张卡只需要存一半的权重算一半的矩阵乘法。Attention部分的TP稍微复杂一点。多头注意力本身就是天然可并行的每个头可以独立计算。所以TP在Attention上的做法通常是按头切分每张卡负责几个头的QKV计算和Attention输出最后再拼接。TP的通信模式是每层都要通信。在FFN中前向传播需要一次AllReduce把Z1和Z2相加反向传播需要两次AllReduce。Attention部分也类似。这意味着TP的通信频率非常高对卡间带宽要求极高。所以TP通常只在同一台机器内的NVLink卡之间使用跨机器做TP会因为网络带宽不足而严重拖慢速度。注意TP的通信量跟模型参数量成正比跟TP度数也有关。TP度数越大每张卡上的计算量越小但通信次数不变通信占比就越高。实践中TP一般不超过8再大就得不偿失了。2.3 PP流水线并行按层切分模型PPPipeline Parallelism的思路更符合直觉把模型按层切成若干段每段放到不同的卡上。比如一个32层的TransformerPP4那就每8层放一张卡。数据从第一张卡进去算完8层后把中间结果传给第二张卡以此类推。PP最大的问题是流水线气泡。假设你有4张卡每张卡负责8层。如果一次只处理一个batch的数据那么当第一张卡在算的时候后面三张卡都在闲着等第一张卡算完传给第二张卡第一张卡又闲下来了。这样GPU利用率极低大部分时间都在等。为了解决这个问题PP通常采用微批次micro-batch技术。把一个大的batch切成多个小micro-batch让它们像流水线一样依次进入。当第一个micro-batch进入第二张卡时第二个micro-batch进入第一张卡这样所有卡都能同时干活。但即使这样流水线的开始和结束阶段仍然会有气泡气泡的大小跟PP度数和micro-batch数量有关。PP的通信量相对较小因为只在层与层之间传递激活值而且通信频率是每个micro-batch一次不像TP那样每层都要通信。所以PP可以跨机器使用对带宽要求没那么苛刻。PP的另一个难点是负载均衡。如果模型各层的计算量不均匀比如有些层有MoE有些没有简单的按层数切分会导致某些卡成为瓶颈。实践中需要根据实际计算量来调整切分点。2.4 CP上下文并行专治长序列CPContext Parallelism是最近几年随着长上下文模型火起来的一种并行方式。它的核心思路是把输入序列切分成多段每段放到不同的卡上但Attention计算需要跨段通信。为什么需要CP因为当序列长度达到32K、128K甚至更长时Attention的计算量和激活值显存占用会爆炸。以标准Attention为例计算复杂度是O(n²)序列长度翻倍计算量翻四倍。而且中间产生的Attention矩阵n×n会占用大量显存。单卡根本扛不住。CP的做法是把序列切成N段每张卡负责一段的Q、K、V计算。但Attention需要每个token看到所有token所以每张卡在计算Attention时需要拿到其他卡上的K和V。这就需要一个环形通信模式每张卡把自己的K、V发给下一张卡同时接收上一张卡的K、V经过N-1次传递后每张卡就拥有了所有K、V的副本可以完成完整的Attention计算。CP的通信量跟序列长度和模型维度有关但相比TPCP的通信频率较低每个Attention层一次环形通信所以可以跨机器使用。CP特别适合长文本理解、长文档摘要、代码生成等场景。2.5 EP专家并行MoE模型的专属方案EPExpert Parallelism是专门为混合专家模型MoE设计的并行策略。MoE的核心思想是把FFN层拆成多个“专家”每个token只激活其中少数几个专家。比如Mixtral 8x7B有8个专家每个token只走其中2个。EP就是把不同的专家放到不同的卡上。每张卡负责几个专家的计算。当token进来时通过一个路由网络决定它该去哪些专家然后把token发送到对应的卡上算完后再把结果传回来。EP的通信模式是All-to-All因为每个token可能被路由到任意专家所以需要跟所有卡交换数据。All-to-All的通信量跟token数量和专家数量有关通常比较大。而且EP面临负载不均衡的问题如果路由网络倾向于把大部分token发给少数几个专家那这几张卡就会过载其他卡闲着。实践中需要加负载均衡损失来缓解。EP的优势在于大幅减少计算量。因为每个token只激活部分专家总FLOPs比稠密模型低很多。所以MoE模型可以用更少的计算量达到接近稠密大模型的效果。但EP的通信开销和负载均衡问题需要仔细调优。3. 五种策略的对比与组合使用3.1 一张表看清五种并行的差异并行策略切分维度通信模式通信频率适用场景主要缺点DP数据AllReduce每步一次模型能单卡装下显存冗余通信量大TP矩阵运算AllReduce每层多次单机内大矩阵通信频繁跨机差PP模型层P2P每micro-batch跨机深层模型流水线气泡CP序列环形通信每Attention层长序列实现复杂EP专家All-to-All每MoE层MoE模型负载不均衡这张表建议你保存下来以后看到并行配置时对照着看基本能判断出这个方案的侧重点和潜在瓶颈。3.2 实际项目中怎么组合真实的大模型训练从来不是只用一种并行策略而是多种策略组合。最常见的组合是TPPPDP也就是常说的“3D并行”。以GPT-3 175B为例OpenAI用的配置大概是TP8、PP8、DP8总共512张A100。TP8在同一台机器内做利用NVLink的高带宽PP8跨机器做通信量小DP8在剩余的维度上做数据并行。这样每张卡上的显存占用和计算量都被控制在合理范围内。具体怎么选组合取决于你的硬件拓扑和模型结构。一般来说单机内优先用TP因为NVLink带宽高TP的频繁通信可以接受。跨机用PP因为PP通信量小对带宽不敏感。DP放在最外层用来扩展吞吐量。长序列场景加CP比如训练128K上下文的模型。MoE模型加EP把专家分散到不同卡上。实操心得组合并行时TP和PP的度数乘积不能超过总卡数DP度数 总卡数 / (TP × PP)。比如你有64张卡TP8、PP4那DP2。如果还要加CP那DP 总卡数 / (TP × PP × CP)。配置时一定要算清楚否则启动会报错。3.3 通信开销的量化分析理解通信开销对算法同学来说很重要因为你需要判断一个并行方案是否合理。通信开销主要取决于两个因素通信量和带宽。通信量方面DP的AllReduce通信量是2×参数量×N-1/NN是DP度数。TP的AllReduce通信量跟激活值大小有关每层大概是2×batch×seq×hidden×TP-1/TP。PP的P2P通信量是batch×seq×hidden每个micro-batch一次。CP的环形通信量是batch×seq×hidden×CP-1/CP每层一次。EP的All-to-All通信量跟token数量和专家数量有关。带宽方面同一台机器内NVLink的带宽大概是600GB/sA100到900GB/sH100跨机器InfiniBand大概是200GB/s到400GB/s。所以TP放在单机内、PP跨机就是因为TP通信量大需要高带宽PP通信量小可以走网络。你可以粗略估算一下假设模型有100层hidden8192batch4seq4096TP8。每层TP的AllReduce通信量大概是2×4×4096×8192×7/8 ≈ 470MB。100层就是47GB。如果NVLink带宽是600GB/s那通信时间大概是78ms。如果计算时间也是几十毫秒量级那通信占比就很高了。这就是为什么TP度数不能太大。4. 实操中怎么配置和调试并行策略4.1 从单卡到多卡迁移步骤如果你已经有一个单卡训练脚本想迁移到多卡并行可以按以下步骤操作第一步确定并行方案。先算一下模型参数量和显存占用。如果单卡能装下直接用DDP做DP就行。如果装不下看是层数多还是单层大。层数多用PP单层大用TP。如果序列特别长加CP。如果是MoE加EP。第二步修改模型代码。TP需要把线性层替换成列并行和行并行版本比如Megatron-LM里的ColumnParallelLinear和RowParallelLinear。PP需要把模型按层切分每段包成一个Module。CP需要修改Attention实现加入环形通信逻辑。EP需要实现路由网络和All-to-All通信。第三步配置通信组。用PyTorch的dist.new_group()创建不同的通信组。TP组、PP组、DP组、CP组、EP组各是一个独立的通信组。初始化进程组时要注意顺序通常先初始化全局组再按维度创建子组。第四步调整数据加载。DP维度上要确保每张卡拿到不同的数据。通常用DistributedSampler来实现。PP维度上要注意micro-batch的切分和传递。第五步启动训练。用torchrun或mpirun启动指定nproc_per_node和nnodes。启动脚本里要设置好RANK、WORLD_SIZE、MASTER_ADDR、MASTER_PORT等环境变量。# 示例用torchrun启动8卡TP2卡PP4卡DP的训练 torchrun \ --nproc_per_node8 \ --nnodes8 \ --node_rank$NODE_RANK \ --master_addr$MASTER_ADDR \ --master_port29500 \ train.py \ --tensor_parallel_size8 \ --pipeline_parallel_size2 \ --data_parallel_size4注意并行配置的乘积必须等于总进程数。上面这个例子中8×2×464所以需要64个进程也就是8台机器每台8卡。4.2 显存估算与参数选择选择并行策略时显存估算是关键。以FP16训练为例每张卡的显存占用包括模型参数参数量×2字节 / TP度数梯度参数量×2字节 / TP度数优化器状态参数量×8字节 / TP度数Adam的FP32动量和方差激活值跟batch、seq、hidden、层数有关PP可以分摊通信缓冲区跟并行策略有关TP需要额外的AllReduce缓冲区假设一个13B模型TP4PP2DP4总共32卡。每张卡的参数显存是13B×2/4 ≈ 6.5GB梯度6.5GB优化器状态26GB加起来39GB。激活值大概每层几GBPP2的话每张卡承担一半层数激活值也减半。总体下来每张卡大概50-60GB80GB的A100可以装下。如果TP2那参数显存翻倍到13GB梯度13GB优化器52GB加起来78GB再加上激活值就爆了。所以TP度数不能太小否则单卡显存压力太大。4.3 常见报错与排查思路并行训练最容易出的问题就是通信超时和显存溢出。通信超时通常是因为某个进程卡住了导致其他进程在等它。排查方法是看日志里哪个rank最后输出然后检查那个rank的代码逻辑。常见原因包括数据加载不均匀导致某些rank提前结束、条件判断导致某些rank跳过了通信操作、网络抖动导致NCCL超时。显存溢出的话先看是哪张卡爆了。如果是所有卡都爆说明并行配置不够需要增加TP或PP度数。如果是某张卡单独爆可能是负载不均衡比如PP切分点没选好某张卡层数太多。或者是EP路由不均某些专家过载。NCCL相关的报错也很多比如“NCCL error: unhandled system error”通常是网络问题“NCCL error: invalid argument”通常是通信组配置错了。建议开启NCCL_DEBUGINFO看详细的通信日志。实操心得调试并行训练时先用小模型比如2层、hidden256跑通流程确认通信组配置正确再换成大模型。这样能快速定位是配置问题还是规模问题。另外NCCL_DEBUGWARN可以过滤掉大量无用日志只看警告和错误。5. 算法同学需要关注的关键问题5.1 并行策略对训练效果的影响并行策略本身不改变模型数学等价性但实现细节可能影响数值精度。比如TP的AllReduce是浮点加法不同卡上的部分和相加顺序不同可能导致微小的数值差异。这种差异在训练初期可能不明显但长期累积可能影响收敛。实践中通常用FP32做AllReduce来减少精度损失。PP的micro-batch数量会影响梯度累积的效果。micro-batch越多流水线气泡越小但每个micro-batch的batch size越小BatchNorm的统计量可能不准。不过LLM通常用LayerNorm不受batch size影响所以这个问题不大。CP的环形通信会改变Attention的计算顺序但数学上是等价的。不过如果实现有bug比如K、V传递不完整会导致Attention结果错误表现为loss异常或生成乱码。EP的负载均衡损失是个超参数需要调。如果负载均衡损失太大路由网络会倾向于均匀分配token可能牺牲模型效果。如果太小又会导致专家过载。通常从0.01开始调。5.2 跟Infra同学沟通的正确姿势算法同学跟Infra同学沟通时最容易出现的问题是只提需求不提约束。比如你说“我要训练一个100B模型序列长度32K”Infra同学会问你有多少卡什么型号卡间带宽多少这些信息决定了并行方案的选择。正确的沟通方式是给出模型规模、序列长度、batch size、目标吞吐量、可用硬件然后让Infra同学推荐并行方案。你也可以自己先估算一下提出一个初步方案让Infra同学评估可行性。另外要理解Infra同学的难处。并行策略的调优是个多目标优化问题显存、吞吐、通信、负载均衡往往顾此失彼。有时候为了跑通不得不牺牲一些吞吐。所以沟通时要明确优先级是先跑通还是先跑快是显存优先还是速度优先5.3 什么时候该自己动手什么时候该用现成框架现在有很多现成的分布式训练框架比如Megatron-LM、DeepSpeed、FairScale、Colossal-AI。这些框架已经实现了TP、PP、DP、CP、EP的各种组合你只需要改配置文件就行。对于大多数算法同学来说优先用现成框架不要自己从头实现。Megatron-LM是NVIDIA出的TP和PP实现最成熟适合训练GPT类模型。DeepSpeed是微软出的ZeRO系列优化很强适合显存受限的场景。Colossal-AI是潞晨科技出的支持多种并行策略组合文档比较友好。但如果你有特殊需求比如自定义的模型结构、特殊的通信模式那可能需要在框架基础上改。这时候你需要理解框架的并行实现原理知道怎么扩展。建议先读框架的源码从简单的例子入手逐步修改。提示Megatron-LM的代码结构比较清晰建议从megatron/model/目录下的Transformer实现开始读理解ColumnParallelLinear和RowParallelLinear的用法。这是理解TP的最佳入口。6. 从理论到落地一个完整的配置案例假设你要训练一个70B参数的稠密模型序列长度8K可用硬件是64张A100 80GB单机8卡NVLink互联机器间InfiniBand 200GB/s。目标是跑通训练吞吐量尽量高。第一步显存估算。70B模型FP16参数是140GB梯度140GBAdam优化器状态560GB总共840GB。64张卡平均每张13GB看起来不多。但激活值是大头8K序列、hidden8192、80层每层激活值大概几GB总共可能上百GB。所以需要PP来分摊激活值。第二步选择并行方案。TP8单机内PP4跨机DP264/8/42。这样每张卡承担70B/8/4 ≈ 2.2B参数显存压力不大。激活值也被PP4分摊每张卡承担20层。CP不需要因为8K序列不算特别长。EP不需要因为不是MoE。第三步配置通信组。TP组是同一台机器内的8张卡PP组是跨机器的4组DP组是剩余的2组。用dist.new_group()创建这些组注意group的创建顺序和rank的映射关系。第四步调整micro-batch。PP4的情况下micro-batch数量建议至少是PP度数的4倍也就是16个micro-batch这样流水线气泡比较小。全局batch size micro-batch size × DP度数 × 梯度累积步数。假设micro-batch size1DP2梯度累积8那全局batch size16。第五步启动训练。用torchrun启动8台机器每台8卡总共64进程。配置TP8、PP4、DP2。监控GPU利用率、通信时间占比、loss曲线。这个配置下TP的AllReduce通信量每层大概是2×1×8192×8192×7/8 ≈ 117MB80层就是9.4GB。NVLink带宽600GB/s通信时间约16ms。计算时间取决于FLOPs70B模型每token大概140GFLOPs8K序列就是1.1TFLOPsA100算力312TFLOPS理论计算时间约3.6ms。通信占比偏高但可以接受。如果TP降到4通信量减半但单卡显存翻倍需要权衡。PP的P2P通信量是1×8192×8192×2 ≈ 134MB每个micro-batch一次16个micro-batch就是2.1GB。InfiniBand 200GB/s通信时间约10ms。相比计算时间可以忽略。DP的AllReduce通信量是2×70B×1/2 ≈ 70GB每步一次。这个通信量很大但DP的AllReduce可以跟反向传播重叠实际影响没那么大。整体来看这个配置是可行的。实际跑起来后根据监控数据再微调比如增加micro-batch数量、调整TP/PP比例、优化数据加载等。实操心得第一次跑大规模并行训练时建议先用小规模比如8卡验证配置正确再扩展到全量。扩展时注意检查每个rank的日志确保没有rank掉队。另外保存checkpoint时要考虑并行策略的兼容性TP和PP的切分方式会影响checkpoint的加载。建议用框架提供的分布式checkpoint功能不要自己手动存。7. 一些容易踩的坑和独家建议7.1 TP和PP的度数选择有讲究TP度数不是越大越好。TP越大单卡计算量越小但通信占比越高。经验法则是TP度数不要超过单机卡数因为跨机TP的通信开销太大。而且TP度数最好是2的幂次这样矩阵切分比较均匀。PP度数也不是越大越好。PP越大流水线气泡越多而且跨机通信次数增加。通常PP度数不超过16再大就得不偿失了。PP的切分点要尽量选在计算量相近的层之间避免某张卡过载。TP和PP的乘积决定了模型并行的总度数。这个乘积不能超过总卡数否则DP就没法做了。如果总卡数不够可以考虑用ZeRO来减少DP的显存冗余这样可以用更少的卡跑更大的模型。7.2 CP的实现细节很容易出错CP的环形通信需要仔细处理边界条件。比如序列长度不能被CP度数整除时需要做padding。padding的部分在Attention计算时要mask掉否则会影响结果。另外环形通信的顺序很重要每张卡要确保在正确的时间发送和接收否则会死锁。CP和TP组合时通信组要分开创建。CP组和TP组是不同的维度不能混在一起。通常先做CP的环形通信再做TP的AllReduce顺序不能反。7.3 EP的负载均衡是个持续调优的过程EP的负载均衡损失需要根据实际路由情况调整。如果发现某些专家过载可以增大负载均衡损失的权重。但权重太大会影响模型效果因为路由网络会为了均衡而牺牲专业性。实践中通常用辅助损失auxiliary loss来鼓励均衡权重从0.01开始调。EP的All-to-All通信可以用NCCL的all_to_all_single来实现但要注意token的排列顺序。通常需要先按专家分组再发送接收后再还原顺序。这个过程涉及大量的permute和unpermute操作实现起来比较繁琐。7.4 监控和调优的实用工具推荐几个监控工具NVIDIA的nsys可以看GPU利用率和通信时间线PyTorch的profiler可以看每个算子的耗时NCCL的调试日志可以看通信是否正常。另外wandb或tensorboard可以记录loss、吞吐量、显存占用等指标方便对比不同配置的效果。调优时优先看GPU利用率。如果利用率低于50%说明有瓶颈可能是通信、数据加载或计算。然后看通信时间占比如果超过30%说明通信是瓶颈需要调整并行策略。最后看显存占用如果接近上限说明需要增加并行度数或优化激活值。提示PyTorch的torch.cuda.memory_summary()可以打印详细的显存分配情况包括参数、梯度、激活值、通信缓冲区各占多少。这个对定位显存瓶颈非常有用。7.5 从算法角度理解并行的代价最后想跟算法同学说一点并行不是免费的。每增加一种并行策略就增加一层通信开销和实现复杂度。TP的通信频率最高PP有流水线气泡CP实现最复杂EP负载均衡最难调。所以能用简单方案解决就不要用复杂方案。如果你的模型能在单卡上跑就别搞分布式。如果单机8卡能跑就别跨机。如果DP能解决就别上TP。只有当模型大到单卡装不下、单机跑不动时才考虑更复杂的并行组合。理解这些并行策略的代价也能帮你更好地设计模型。比如你知道TP的通信量跟hidden size成正比那在设计模型时就可以考虑用GQA分组查询注意力来减少KV头的数量从而减少TP的通信量。你知道PP的气泡跟micro-batch数量有关那就可以在训练时增大batch size增加micro-batch数量来减少气泡。这些权衡和取舍才是算法同学学Infra的真正价值所在。你不需要会写CUDA kernel但你需要知道什么样的模型结构对Infra友好什么样的并行方案对训练效果影响最小。这样你跟Infra同学合作时才能做出对双方都最优的决策。
返回列表