ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

从算法岗转Infra:五种并行策略(TP/DP/PP/CP/EP)详解与选型实战

从算法岗转Infra:五种并行策略(TP/DP/PP/CP/EP)详解与选型实战 1. 从算法岗转Infra为什么绕不开这五种并行策略做算法的同学第一次打开一个千亿参数模型的训练脚本大概率会被满屏的tp_size、pp_rank、dp_group、ep_degree搞懵。明明在单卡上跑得好好的模型一上多机就报 OOM或者跑起来了但 GPU 利用率只有百分之十几吞吐量惨不忍睹。这时候你会发现光懂 Transformer 结构、懂反向传播、懂优化器已经不够了——模型能不能训起来、训得快不快取决于你怎么把这些计算切分到几十上百张卡上。这就是 LLM 分布式计算要解决的问题。TP、DP、PP、CP、EP 这五个缩写基本覆盖了当前大模型训练和推理中最主流的并行切分维度。它们不是互斥的选项而是可以叠加组合的正交维度。一个典型的千卡训练任务往往是 TP8、PP4、DP16 同时开启再加上 CP 处理长序列、EP 处理 MoE 专家层。你不需要成为 Infra 专家但必须理解每种策略在切什么、通信开销在哪、什么场景下该用哪个。这篇文章面向的是有算法背景、但对分布式训练基础设施还不够熟悉的同学。我会从每种并行策略的核心逻辑讲起说清楚它们各自解决什么问题、代价是什么、实际配置时怎么选再补充一些从实际调试中积累的经验。读完你应该能看懂主流框架的并行配置并且在自己遇到 OOM 或吞吐瓶颈时知道该往哪个方向调。2. 数据并行DP最直观的切分方式也是最容易踩坑的起点2.1 DP到底在切什么数据并行的逻辑最简单每张卡上都放一份完整的模型副本把训练数据切成 N 份每张卡用自己那份数据算梯度然后所有卡把梯度做一次全局归约AllReduce保证每张卡上的模型参数始终一致。用生活化的类比好比五个学生做同一套题库每人分到不同的题目各自做完后把错题本汇总统一更新学习方法。每个人手里的教材模型参数始终是同一本。DP 的优势在于实现简单、对模型结构无侵入。你几乎不需要改模型代码框架层面加个DistributedDataParallel就能跑。但它的约束也很明显每张卡都要存一份完整的模型参数、梯度、优化器状态。以一个 7B 模型为例FP16 参数占 14GB梯度 14GBAdam 优化器状态一阶矩二阶矩56GB加起来接近 84GB——单张 80GB 的卡根本放不下。这就是为什么纯 DP 只能训小模型大模型必须叠加其他并行策略。2.2 梯度累积与DP的配合实际训练中我们经常用梯度累积来模拟更大的 batch size。假设你只有 8 张卡但想要等效 batch size 512可以设置每张卡的 micro batch 为 8梯度累积步数为 8那么等效全局 batch 8卡 × 8 micro batch × 8 累积步 512。这里有个容易搞混的点梯度累积是在单卡内串行执行多次前向反向累积梯度后再做一次 AllReduce。不是每步都通信这能显著降低通信频率。但要注意BatchNorm 这类依赖 batch 统计量的层在梯度累积下行为会变LLM 里通常用 LayerNorm 或 RMSNorm不受影响但如果你在模型里加了自定义的 BatchNorm 层就要小心。2.3 DP的通信瓶颈与ZeRO系列优化标准 DP 的通信量是 2倍模型参数量AllReduce 的 reduce-scatter all-gather。当模型大到几十 GB每次迭代都要传输这么多数据通信就会成为瓶颈。ZeROZero Redundancy Optimizer系列正是为此而生ZeRO-1只切分优化器状态每张卡存 1/N 的优化器状态通信量和标准 DP 相同。ZeRO-2再切分梯度每张卡只存自己负责的那部分参数的梯度。ZeRO-3连模型参数也切分每张卡只存 1/N 的参数前向和反向时需要临时从其他卡 gather 参数。ZeRO-3 本质上已经模糊了 DP 和模型并行的边界它用通信换显存。实测中ZeRO-3 能把 175B 模型的显存需求从数 TB 降到单卡可承受的范围但代价是通信量大幅增加需要高速互联NVLink 或 InfiniBand才能跑出可接受的效率。注意ZeRO-3 在反向传播时需要重新 gather 参数如果网络带宽不够GPU 会有大量时间在等通信利用率可能掉到 30% 以下。上 ZeRO-3 之前先确认你的互联带宽。3. 张量并行TP把矩阵乘法拆开通信换显存3.1 TP的切分逻辑按列切还是按行切张量并行是把单个矩阵运算拆到多张卡上。以 Transformer 里的线性层 Y XW 为例有两种切法按列切分Column Parallel把权重矩阵 W 按列切成 [W1, W2]每张卡算 Y1 XW1、Y2 XW2最后把结果拼接起来。这种切法每张卡都需要完整的输入 X输出需要 AllGather。按行切分Row Parallel把 W 按行切成 [W1; W2]同时把输入 X 按列切成 [X1, X2]每张卡算 Y_partial X1W1 或 X2W2最后把部分和相加得到完整输出。这种切法输出需要 AllReduce。Megatron-LM 的经典设计是Attention 层的 QKV 投影用列并行输出投影用行并行FFN 的第一个线性层用列并行第二个用行并行。这样一列一行交替每两层只需要一次 AllReduce通信量被压到最小。3.2 TP的通信代价为什么这么大TP 的通信发生在每一层的前向和反向中频率极高。以列并行行并行为例每个 Transformer block 至少需要两次 AllReduce前向一次、反向一次。一个 80 层的模型每步迭代就有 160 次 AllReduce。每次 AllReduce 的数据量是 batch × seq_len × hidden_dim在长序列场景下非常可观。这就是为什么 TP 通常只在单机内使用NVLink 互联带宽 600GB/s 以上跨机 TP 在以太网或普通 InfiniBand 上几乎不可行。实践中 TP size 一般不超过 8正好对应一台 8 卡服务器。超过 8 路 TP通信开销会吃掉所有并行收益。3.3 TP与DP的组合3D并行的基础最常见的组合是 TP DP。假设你有 64 张卡可以设 TP8、DP8。每 8 张卡组成一个 TP 组组内做张量并行8 个 TP 组之间做数据并行。这样模型被切了 8 份每份放在一组卡上组间用 DP 扩大吞吐。配置时要注意TP 组内的卡最好在同一台机器内保证 NVLink 互联DP 组间的通信频率低每步一次 AllReduce可以跨机。框架通常用tensor_model_parallel_size和data_parallel_size两个参数控制两者的乘积等于总卡数不考虑 PP 和 CP 时。4. 流水线并行PP按层切分用调度填满气泡4.1 PP的基本思路与气泡问题流水线并行把模型按层切成多个 stage每个 stage 放在不同的卡或卡组上。数据像流水线一样从第一个 stage 流到最后一个 stage。比如一个 32 层的模型PP4那么每 8 层为一个 stage。PP 的最大问题是气泡Bubble在流水线的开始和结束阶段有些卡在等数据有些卡算完了在等下一批。假设 PP4不做任何优化理想情况下 GPU 利用率只有 1/4 左右实际上因为前向反向交错会好一些。解决方案是微批次Micro-batch把一个 batch 切成多个 micro-batch让它们像工厂流水线一样错开执行。当 micro-batch 1 在 stage 2 计算时micro-batch 2 可以在 stage 1 计算以此类推。micro-batch 越多气泡占比越小。理论上气泡占比约为 (PP-1)/(micro_batch_num PP - 1)。4.2 GPipe与1F1B调度策略GPipe是最朴素的调度所有 micro-batch 的前向做完再做所有反向。这需要缓存所有 micro-batch 的中间激活值显存占用高。1F1BOne Forward One Backward是改进方案每个 stage 在做完一个 micro-batch 的前向后紧接着做上一个 micro-batch 的反向。这样中间激活值的缓存量从 O(micro_batch_num) 降到 O(PP)显存大幅节省。Megatron-LM 默认用 1F1B。还有Interleaved 1F1B虚拟流水线把每个 stage 再切成多个虚拟 stage让每个卡负责不连续的层。比如 PP4、虚拟 stage2 时每张卡负责第 1、5 层和第 3、7 层这样的组合。这能把气泡进一步缩小但通信模式更复杂。4.3 PP的适用场景与配置建议PP 适合模型层数多、单层计算量相对均匀的场景。它的通信量比 TP 小得多只在 stage 边界传激活值可以跨机使用。但 PP 对负载均衡敏感如果某个 stage 的计算量明显大于其他比如 embedding 层和 lm_head 层通常计算量小就会出现木桶效应。实践中PP size 一般选 2 到 16配合较大的 micro-batch 数如 16 到 64来压气泡。配置时要注意pipeline_model_parallel_size和num_micro_batches的配合micro-batch 数最好是 PP size 的整数倍调度更规整。提示如果发现 PP 的某个 stage 明显比其他慢检查层切分是否均匀。Embedding 和输出层通常单独放一个 stage不要和 Transformer 层混在一起。5. 上下文并行CP长序列训练的救命稻草5.1 为什么长序列需要CP当序列长度从 4K 涨到 128K 甚至 1M激活值的显存占用线性增长。Attention 的计算复杂度是 O(n²)显存也是 O(n²)如果不做 FlashAttention 之类的优化。单卡放不下这么长的序列就需要把序列维度切到多张卡上这就是上下文并行。CP 的核心思路把输入序列按 token 维度切成 N 份每张卡处理一段。但 Attention 需要每个 token 看到所有其他 token所以切分后需要跨卡通信来交换 KV。5.2 Ring Attention的工作机制Ring Attention 是目前最主流的 CP 实现。它把 KV 块在卡之间像传环一样传递每张卡先用自己的 Q 和本地 KV 算 attention然后把 KV 块传给下一张卡同时从上一张卡接收新的 KV 块继续计算。经过 N 轮传递后每张卡都见过了所有 KVattention 结果完整。这种方式的通信量和计算量可以重叠在算当前 KV 块的 attention 时下一块 KV 已经在传输了。如果计算时间大于通信时间通信就能被完全隐藏。实测中在 100Gbps 以上的互联带宽下Ring Attention 的通信开销可以控制在 10% 以内。5.3 CP与TP、DP的组合注意事项CP 通常和 TP、DP 一起用。配置时要注意CP 组内的卡需要频繁交换 KV最好放在同一台机器或高带宽互联的节点内。CP size 一般不超过 8再大通信开销就压不住了。另一个坑是位置编码。序列切分后每张卡上的 token 位置索引必须正确否则 RoPE 或 ALiBi 会算错。框架通常会自动处理但如果你自己实现 CP一定要检查位置编码的偏移量。6. 专家并行EPMoE模型的专属切分方式6.1 EP解决的是MoE的显存和计算问题MoEMixture of Experts模型里FFN 层被替换成多个专家每个 token 只激活其中几个。比如 Mixtral 8x7B 有 8 个专家每个 token 激活 2 个。如果所有专家都放在一张卡上显存占用是稠密模型的 8 倍显然不现实。EP 就是把不同的专家放到不同的卡上。每张卡只存一部分专家token 经过路由后被发送到对应专家所在的卡上计算算完再发回来。这样每张卡的显存只增加自己负责的那几个专家。6.2 All-to-All通信EP的核心代价EP 的通信模式是 All-to-All每张卡都要把 token 发给所有其他卡因为 token 可能路由到任何专家同时从所有其他卡接收 token。这是所有并行策略里通信模式最复杂的。All-to-All 的通信量取决于 token 路由的分布。如果路由均匀每张卡发送和接收的数据量差不多如果路由倾斜某些专家被大量激活就会出现热点某些卡通信量暴增。实践中需要加负载均衡损失auxiliary loss来鼓励均匀路由。6.3 EP与TP、DP的配合EP 通常和 TP 组合使用专家内部的矩阵乘法可以用 TP 切分专家之间用 EP 分布。比如 8 个专家、TP2、EP4那么每两张卡组成一个 TP 组负责 2 个专家4 个 TP 组覆盖 8 个专家。配置 EP 时要特别注意expert_model_parallel_size和tensor_model_parallel_size的关系。通常要求 EP size 能整除专家总数且 EP 组内的通信要尽量走高速互联。7. 五种并行策略的选型逻辑与组合实战7.1 一张表看清五种策略的差异策略切分维度通信模式通信频率适用场景典型sizeDP数据AllReduce每步一次扩大吞吐任意TP模型宽度AllReduce/AllGather每层两次单机内切大模型2-8PP模型深度P2Pstage边界跨机切大模型2-16CP序列长度Ring AllGather每层多次长序列训练2-8EP专家维度All-to-All每层一次MoE模型2-647.2 实际配置的计算逻辑假设你有 512 张 GPU要训一个 70B 的稠密模型序列长度 8K。怎么配首先算显存70B 参数 FP16 占 140GB加上梯度、优化器状态、激活值单卡肯定放不下。用 TP8 把模型切 8 份每份 17.5GB 参数加上其他开销大约 70GB单卡 80GB 能放下。TP 组内 8 卡用 NVLink。然后看吞吐512 卡TP8 用掉 8 卡一组还剩 64 组。可以用 PP4 把这 64 组再分成 16 个 PP 组每个 PP 组内 4 个 stage。最后 DP1616 个 PP 组之间做数据并行。验证TP×PP×DP 8×4×16 512正好。序列长度 8K 不算特别长CP 可以不开。模型是稠密模型EP 不需要。7.3 调优时的排查顺序训练跑起来后如果吞吐不达标按这个顺序排查先看 GPU 利用率如果所有卡利用率都低可能是数据加载瓶颈或通信瓶颈。看通信占比用 profiling 工具如 PyTorch Profiler、Nsight Systems看 AllReduce、All-to-All 占了多少时间。如果超过 30%说明通信是瓶颈。检查 TP 是否跨机TP 跨机会导致通信急剧变慢确认 TP 组都在同一台机器内。检查 PP 气泡如果 PP size 大但 micro-batch 少气泡占比会很高。增加 micro-batch 数或减小 PP size。检查 EP 负载均衡如果用了 MoE看专家激活分布是否均匀不均匀就调大 auxiliary loss 的系数。注意不要一上来就调并行配置。先确认单卡能跑通、数据管道没问题、模型代码没有性能陷阱比如没开 FlashAttention、用了低效的 LayerNorm 实现再动并行策略。8. 几个容易搞混的概念和实际踩坑记录8.1 TP和ZeRO-3到底有什么区别很多人觉得 ZeRO-3 和 TP 都是切模型参数应该差不多。实际上两者的通信模式完全不同TP 是在计算过程中切矩阵每层都要通信但通信量小且可预测ZeRO-3 是在参数级别切分前向反向时需要 gather 完整参数通信量大但频率低每层一次 gather。选择上如果单机内 NVLink 充足TP 效率更高如果跨机且带宽有限ZeRO-3 可能更合适因为它对带宽的要求相对低一些可以用通信重叠来隐藏。8.2 PP的stage切分不是越均匀越好理论上每个 stage 的计算量应该相等但实际上要考虑通信。Embedding 层和 lm_head 层参数量大但计算量小如果和 Transformer 层放一起会导致那个 stage 显存吃紧但计算空闲。通常的做法是把 Embedding 和 lm_head 单独放一个 stage或者和少量 Transformer 层搭配。另外第一个和最后一个 stage 通常有额外的通信开销接收输入数据、输出 loss切分时可以适当减少这两个 stage 的层数。8.3 CP在推理场景下的特殊考量训练时 CP 切的是输入序列每张卡处理一段。但推理时如果用户请求的序列长度不一CP 的负载均衡会成问题。有些框架用 continuous batching 来动态调整但 CP 组内的序列长度必须一致否则 Ring Attention 的轮次对不上。实践中推理场景下 CP 用得比训练少更多是用 PagedAttention 之类的显存管理技术来支持长序列。如果非要用 CP 做推理建议把相近长度的请求分到同一个 CP 组。8.4 我踩过的一个EP配置坑有一次配 MoE 模型EP8、TP2总卡数 16。我以为 EP 组和 TP 组是正交的结果发现框架要求 EP 组内的卡必须属于同一个 TP 组。也就是说TP2 时每两张卡组成一个 TP 组EP 只能在 TP 组之间分布。最后改成 EP4、TP4 才跑通。这个坑的根源是专家内部的矩阵乘法用 TP 切分所以一个专家的参数分布在 TP 组内的多张卡上。EP 切的是专家之间的分布必须在 TP 组之上进行。配置时一定要看清框架的约束文档不同框架的实现可能不一样。8.5 通信重叠是性能优化的关键无论用哪种并行策略通信重叠都是提升效率的核心手段。所谓通信重叠就是在计算的同时进行通信。比如 TP 的 AllReduce 可以和下一层的计算重叠PP 的 P2P 传输可以和当前 stage 的计算重叠。实现通信重叠通常需要框架层面的支持如 Megatron-LM 的 async communication但作为算法同学你能做的是确保 batch size 足够大让计算时间大于通信时间选择合适的并行粒度不要让通信过于频繁。9. 从算法视角看分布式你不需要成为Infra专家但必须懂这些回到最初的问题算法同学为什么要学 Infra因为模型越来越大单卡训不动是常态。你不一定要自己写分布式框架但必须能看懂配置、能定位瓶颈、能和 Infra 同学有效沟通。我自己的经验是掌握这五种并行策略后遇到 OOM 时至少知道该加 TP 还是 PP遇到吞吐低时知道该查通信还是查气泡。这比盲目调参高效得多。最后分享一个实用建议如果你刚开始接触分布式训练先从 TPDP 的组合入手这是最成熟、坑最少的方案。跑通之后再逐步加入 PP 和 CP。EP 留到最后因为 MoE 模型本身的调试复杂度就高再加上 EP 的 All-to-All 通信出问题很难定位。每次只改一个维度观察吞吐和显存的变化逐步建立对每种策略的直觉。
返回列表