大模型训练并行策略详解：DP、TP、PP、CP、EP原理与实战调优

1. 算法同学为什么要补上分布式计算这一课做算法的同学尤其是从传统机器学习、深度学习一路走到大模型方向的大概率都经历过这样一个阶段单卡能跑通的模型参数量一上去就爆显存好不容易把 batch size 调到 1训练速度又慢得让人怀疑人生。这时候你去看那些大模型训练框架的文档满屏都是 TP、DP、PP、CP、EP 这些缩写每个字母都认识连在一起就不知道在说什么了。我刚开始接触大模型训练的时候也是这样。当时手里有一批 8 卡的机器想训一个 7B 左右的模型结果发现光是搞明白怎么把这些卡用起来就花了好几天。后来慢慢踩坑踩多了才意识到这些并行策略不是孤立存在的它们各自解决的是不同维度的问题而且在实际训练中往往是组合使用的。这篇文章就是把我自己学习和实践过程中对这些概念的理解整理出来尽量用算法同学能听懂的语言把 TP、DP、PP、CP、EP 这几个核心概念讲清楚。先说结论这五种并行策略分别对应着不同的切分维度。DP 切的是数据TP 切的是模型层内的参数矩阵PP 切的是模型的层CP 切的是序列长度维度EP 切的是专家网络。理解它们的关键不在于记住定义而在于搞清楚每种策略在通信开销、显存占用、计算效率这三个维度上的取舍。你选择哪种策略本质上是在这三个维度之间做权衡。这篇文章适合谁看如果你是大模型方向的算法工程师正在或者即将参与多卡训练任务对分布式并行只知道个大概但说不清楚细节那这篇内容应该能帮到你。如果你只是想了解一下大模型训练的基本原理不涉及具体实操那也可以看看前面的概念部分后面涉及具体配置和调优的部分可以跳过。2. 五种并行策略的核心原理与适用场景拆解2.1 DP最直观的数据并行但远没有你想的那么简单DP 全称 Data Parallel数据并行。这是最容易理解的一种并行方式每张卡上都放一份完整的模型副本然后把训练数据切成 N 份每张卡处理一份各自计算梯度最后把所有卡上的梯度做一次全局归约更新模型参数。听起来很简单对吧但这里有个关键问题每张卡上都要放一份完整的模型参数、梯度、优化器状态。以一个 7B 参数的模型为例FP16 精度下参数本身占 14GB梯度占 14GB如果用 Adam 优化器优化器状态一阶矩和二阶矩还要占 28GBFP32加起来就是 56GB。这还没算激活值。一张 80GB 的 A100 勉强能放下但如果你想把 batch size 调大一点或者模型再大一点单卡就放不下了。所以 DP 的局限性很明显它解决的是计算吞吐的问题不解决单卡显存放不下模型的问题。当模型大到单卡装不下的时候你就需要模型并行来帮忙了。不过 DP 本身也在进化。现在常说的 DP 其实有两种形态一种是传统的 DP每张卡独立计算梯度然后 AllReduce另一种是 ZeROZero Redundancy Optimizer它把优化器状态、梯度、参数分片到不同的卡上需要的时候再通过通信收集回来。ZeRO 的 Stage 1 只分片优化器状态Stage 2 分片优化器状态和梯度Stage 3 连参数也分片。实际训练中ZeRO Stage 1 或 2 配合 DP 是最常见的组合因为 Stage 3 虽然省显存但通信开销会显著增加。实操心得如果你的模型单卡能放下优先用 DP ZeRO Stage 1。Stage 2 在梯度分片后需要额外的 AllGather 通信对带宽压力更大。Stage 3 一般只在模型实在放不下的时候才用而且通常要配合 TP 或 PP 一起。2.2 TP把矩阵切开算通信是最大的代价TP 全称 Tensor Parallel张量并行。它的核心思想是把模型中的权重矩阵切分成多块每张卡只负责计算其中一块然后通过通信把结果拼起来。以 Transformer 中的多头注意力为例。假设有 H 个注意力头TP 度为 T那么每张卡负责 H/T 个头。Q、K、V 的投影矩阵按列切分每张卡计算自己那部分头的 Q、K、V做注意力计算然后输出投影矩阵按行切分最后把各卡的结果加起来。这就是 Megatron-LM 提出的经典 TP 方案。TP 的关键在于切分方式要和计算图的通信模式匹配。列切分Column Parallel之后通常接行切分Row Parallel这样只需要在行切分的输出做一次 AllReduce通信量最小。如果切分方式不对中间会产生大量不必要的通信性能会急剧下降。TP 的通信开销是它最大的问题。每一层 Transformer 都需要做两次 AllReduce前向一次反向一次而且这个通信是同步的所有卡必须等最慢的那张卡完成。所以 TP 通常只在单机内使用因为单机内 NVLink 的带宽远高于跨机的 InfiniBand。跨机 TP 的通信开销会大到让训练效率急剧下降。TP 度一般选 2、4、8很少超过 8。因为一台机器通常就是 8 卡TP 度超过 8 就意味着要跨机通信得不偿失。实际配置中TP 度通常等于单机 GPU 数或者单机 GPU 数的一半。注意事项TP 度不是越大越好。TP 度增加每张卡上的计算量减少但通信量不变AllReduce 的数据量是固定的所以计算通信比会下降。当 TP 度超过一定值后通信会成为瓶颈训练速度反而变慢。2.3 PP按层切分流水线是门艺术PP 全称 Pipeline Parallel流水线并行。它的思路很直接把模型的层分成若干段每张卡或每组卡负责一段。比如一个 32 层的模型PP 度为 4那么每张卡负责 8 层。但 PP 有个致命问题如果只是简单地按层切分那么同一时刻只有一张卡在计算其他卡都在等待。这就好比工厂流水线上只有一个工位在干活其他工位都在闲着。为了解决这个问题PP 引入了微批次Micro-batch的概念把一个大批次切成多个小批次让不同的小批次在不同的流水线阶段同时计算从而提高利用率。这就是所谓的流水线调度。常见的调度策略有 GPipe 和 1F1BOne Forward One Backward。GPipe 的做法是先做完所有微批次的前向再做所有微批次的反向优点是实现简单缺点是显存占用高因为要保存所有微批次的激活值。1F1B 则是前向和反向交替进行显存占用更低但实现更复杂。PP 的另一个问题是流水线气泡Bubble。在流水线的开始和结束阶段总有一些卡处于空闲状态。气泡的大小和流水线深度、微批次数量有关。微批次越多气泡占比越小但显存占用也越高。实际配置中微批次数量通常设为 PP 度的 2 到 4 倍。PP 的通信量相对较小因为只需要在层与层之间传递激活值和梯度而且这个通信可以和计算重叠。所以 PP 适合跨机使用通常和 TP 配合单机内用 TP跨机用 PP。2.4 CP序列维度的切分长上下文训练的利器CP 全称 Context Parallel上下文并行。这是相对较新的一种并行策略主要解决的是长序列训练的问题。在标准 Transformer 中注意力计算的时间复杂度是 O(n²)n 是序列长度。当序列长度达到 32K、128K 甚至更长时单卡的显存和计算能力都扛不住。CP 的思路是把序列维度切分开每张卡负责一部分序列然后在注意力计算时通过通信交换 K、V。具体来说假设序列长度为 SCP 度为 C那么每张卡负责 S/C 个 token。在计算注意力时每张卡需要其他卡上的 K、V 来计算自己的注意力输出。这就需要一个 AllGather 或者 Ring Attention 的通信模式。Ring Attention 是现在比较流行的方案它把 K、V 在卡之间循环传递每张卡计算一部分注意力最后汇总。CP 的通信模式和 TP 不同它的通信量随着序列长度增加而增加但和模型参数量无关。所以 CP 特别适合长上下文场景比如训练支持 128K 上下文的模型。CP 度一般选 2、4、8和 TP 类似通常限制在单机内。实操心得CP 和 TP 可以组合使用但要注意通信模式的冲突。TP 的 AllReduce 和 CP 的 AllGather 如果同时进行可能会争抢带宽。实际配置中通常会把 TP 和 CP 放在不同的通信组里或者调整它们的相对大小来平衡通信压力。2.5 EP专家并行的核心MoE 模型的关键EP 全称 Expert Parallel专家并行。它是专门为 MoEMixture of Experts模型设计的并行策略。MoE 模型的核心思想是把 FFN 层替换成多个专家网络每个 token 通过门控网络选择 top-k 个专家进行计算。EP 就是把这些专家分布到不同的卡上每张卡负责一部分专家。EP 的关键在于 token 的路由Routing。每个 token 需要被发送到它选择的专家所在的卡上计算完后再发送回来。这个过程中涉及大量的 All-to-All 通信通信量和 token 数量、专家数量、top-k 值都有关系。EP 的挑战在于负载均衡。如果某些专家被选择的频率远高于其他专家那么负责这些专家的卡就会成为瓶颈其他卡则在等待。为了解决这个问题MoE 模型通常会引入负载均衡损失Load Balancing Loss鼓励门控网络均匀地选择专家。EP 度通常等于专家总数或者专家总数除以每张卡上的专家数。比如 64 个专家每张卡放 8 个那么 EP 度就是 8。EP 通常和 DP、TP 组合使用形成复杂的并行策略。3. 组合策略实际训练中怎么搭配使用3.1 三维并行DP TP PP 的经典组合在实际的大模型训练中很少只用一种并行策略。最常见的是 DP TP PP 的三维组合也就是所谓的 3D 并行。假设你有 64 张 GPU想训练一个 70B 的模型。一种可能的配置是TP8PP4DP2。这样总卡数就是 8×4×264。TP8 在单机内完成PP4 跨 4 台机器DP2 复制两份。这种配置的逻辑是TP 负责单机内的模型切分利用 NVLink 的高带宽PP 负责跨机的模型切分通信量相对较小DP 负责数据并行提高吞吐量。三者各司其职共同支撑起大模型的训练。配置的时候有个经验法则TP 度不要超过单机 GPU 数PP 度不要超过机器数DP 度尽量大。因为 DP 的通信效率最高只需要 AllReduce 梯度TP 的通信效率最低每层都要 AllReducePP 居中。3.2 长上下文场景CP 的引入当序列长度超过 8K 时激活值的显存占用会急剧增加。这时候就需要引入 CP 来切分序列维度。一个典型的配置是TP4CP2PP4DP2。总卡数 4×2×4×264。CP2 把序列切成两半每张卡负责一半的序列注意力计算时通过 Ring Attention 交换 K、V。CP 的引入会改变通信模式。原本 TP 的 AllReduce 是在 TP 组内进行的现在 CP 组内也需要通信。实际配置中通常会把 TP 和 CP 放在同一台机器内因为它们的通信都比较频繁需要高带宽支持。3.3 MoE 模型EP 的加入MoE 模型的并行策略更复杂因为多了 EP 这个维度。一个典型的 MoE 训练配置可能是TP4EP8PP2DP2。总卡数 4×8×2×2128。EP8 意味着有 8 组专家每组负责一部分专家。token 路由的时候需要根据门控网络的输出把 token 发送到对应的专家组。这个 All-to-All 通信是 MoE 训练的主要开销之一。MoE 的并行配置有个特殊之处EP 和 DP 可以共享通信组。因为 MoE 的专家是分布式的DP 的数据并行和 EP 的专家并行可以合并成一个维度。这样可以在不增加通信开销的情况下提高并行度。4. 通信开销与性能调优的实操要点4.1 通信原语AllReduce、AllGather、All-to-All 的区别理解并行策略的通信开销首先要搞清楚几种基本的通信原语。AllReduce 是把所有卡上的数据加起来然后广播回所有卡。DP 的梯度同步用的就是 AllReduce。AllReduce 的数据量等于模型参数量通信量和卡数成对数关系Ring AllReduce或线性关系Tree AllReduce。AllGather 是把所有卡上的数据收集起来每张卡都得到完整的数据。ZeRO Stage 3 的参数收集、CP 的 K、V 交换用的都是 AllGather。AllGather 的数据量等于所有卡上数据的总和。All-to-All 是每张卡把自己的数据分发给其他所有卡同时接收其他卡发来的数据。EP 的 token 路由用的就是 All-to-All。All-to-All 的通信量最大因为每张卡都要和所有其他卡通信。ReduceScatter 是 AllReduce 的变体只做归约不做广播结果分散在各卡上。ZeRO 的梯度分片用的就是 ReduceScatter。这几种通信原语的效率排序是AllReduce ReduceScatter AllGather All-to-All。所以配置并行策略的时候要尽量减少 All-to-All 的使用或者把它限制在高速互联的范围内。4.2 通信组划分哪些卡应该放在一起通信组的划分直接影响训练效率。核心原则是通信越频繁的并行维度越应该放在高速互联的范围内。TP 的通信最频繁每层都要 AllReduce所以 TP 组应该放在同一台机器内利用 NVLink。CP 的通信也比较频繁通常和 TP 放在同一台机器内。PP 的通信相对较少可以跨机。DP 的通信最少每个 step 才同步一次梯度可以跨更多机器。EP 的 All-to-All 通信量很大最好也放在同一台机器内或者至少在同一组高速互联的机器内。实际配置中通常的优先级是TP CP EP PP DP。也就是说TP 和 CP 优先占用单机内的卡EP 次之PP 再次之DP 最后。4.3 显存优化激活值重计算与 ZeRO 的配合并行策略解决的是模型参数放不下的问题但激活值的显存占用同样不可忽视。尤其是在长序列场景下激活值可能比参数还大。激活值重计算Activation Checkpointing是常用的优化手段。它的思路是不保存中间层的激活值在反向传播时重新计算。这样可以把激活值的显存占用从 O(n) 降到 O(√n)代价是增加约 30% 的计算量。ZeRO 和激活值重计算可以配合使用。ZeRO 负责分片参数、梯度、优化器状态激活值重计算负责降低激活值占用。两者结合可以在有限的显存下训练更大的模型。注意事项激活值重计算不是万能的。如果序列特别长即使重计算也可能放不下。这时候就需要 CP 来切分序列维度。另外重计算会增加计算时间如果计算资源本来就紧张需要权衡是否值得。5. 常见问题与排查技巧实录5.1 训练速度慢如何定位瓶颈训练速度慢是最常见的问题。排查的时候先看是计算瓶颈还是通信瓶颈。如果是计算瓶颈GPU 利用率会很高接近 100%但吞吐量上不去。这时候要考虑是不是模型太大、batch size 太小、或者计算图有优化空间。如果是通信瓶颈GPU 利用率会波动很大有时候高有时候低。这时候要看是哪种通信导致的。可以用 NCCL 的调试日志或者 PyTorch Profiler 来定位。常见的通信瓶颈包括TP 度过大导致 AllReduce 频繁、PP 气泡过大导致流水线效率低、EP 的 All-to-All 通信量过大。针对不同的瓶颈调整策略也不同TP 过大就减小 TP 度增加 PP 度PP 气泡大就增加微批次数量EP 通信大就减少 EP 度或者优化路由算法。5.2 显存溢出从哪些维度入手排查显存溢出OOM是另一个高频问题。排查的时候先看是参数、梯度、优化器状态还是激活值导致的。参数、梯度、优化器状态的显存占用是固定的可以用公式算出来。如果这部分就超了那必须用 ZeRO 或者 TP/PP 来分片。激活值的显存占用和 batch size、序列长度、模型层数都有关系如果这部分超了可以用激活值重计算或者 CP 来优化。有个实用的技巧先用小 batch size 和小序列长度跑通然后逐步增加观察显存占用的变化。这样可以定位到具体是哪个维度导致的 OOM。5.3 负载不均衡EP 场景下的特殊问题MoE 模型的 EP 场景下负载不均衡是个特殊问题。如果某些专家被选择的频率远高于其他专家那么负责这些专家的卡就会成为瓶颈。排查的方法是统计每个专家的被选择次数。如果发现分布严重不均可以考虑调整负载均衡损失的权重或者使用专家容量Expert Capacity来限制每个专家最多处理多少 token。实操心得负载均衡损失的权重需要调。太小了起不到均衡作用太大了会影响模型效果。一般从 0.01 开始试根据专家选择的分布情况调整。5.4 通信超时NCCL 配置的坑NCCL 是 NVIDIA 的集合通信库PyTorch 的分布式训练底层用的就是它。NCCL 的配置对训练稳定性影响很大。常见的 NCCL 问题包括通信超时、通信挂起、通信错误。这些问题通常和网络配置、NCCL 版本、环境变量有关。几个关键的 NCCL 环境变量NCCL_DEBUG 可以打开调试日志NCCL_IB_DISABLE 可以禁用 InfiniBandNCCL_P2P_DISABLE 可以禁用 P2P 通信。排查问题的时候可以先打开 NCCL_DEBUGINFO看看通信初始化是否正常。注意事项NCCL 的版本要和 CUDA 版本匹配否则可能出现兼容性问题。另外不同机器上的 NCCL 版本要一致否则可能出现通信错误。6. 从单卡到多卡一个实际的配置案例6.1 场景描述与硬件条件假设我们有一个 70B 参数的模型需要在一批 8 卡 A100 80GB 的机器上训练。总共有 8 台机器64 张 GPU。序列长度 4Kbatch size 目标是 512。先算一下显存需求。70B 参数FP16 精度下参数占 140GB梯度占 140GBAdam 优化器状态占 560GBFP32加起来 840GB。64 张卡平均下来每张卡 13GB 左右看起来不多。但激活值的占用还没算而且 TP、PP 的切分也会影响每张卡的实际占用。6.2 并行策略选择与参数计算基于前面的分析我们选择 TP8PP4DP2 的配置。TP8 在单机内完成利用 NVLinkPP4 跨 4 台机器DP2 复制两份。每张卡上的参数占用70B / (8×4) 2.2B 参数FP16 下占 4.4GB。梯度同样 4.4GB。优化器状态如果不用 ZeRO每张卡要存完整的 70B 优化器状态那就爆了。所以必须用 ZeRO Stage 1把优化器状态分片到 DP 组内。DP2 的话每张卡存一半的优化器状态70B×4bytes/2 140GB还是太大。这里就看出问题了DP2 太小ZeRO 的分片效果有限。如果把 DP 调大比如 DP8那么 TP8PP1总卡数 64。但 PP1 意味着没有流水线并行所有层都在同一组卡上显存可能放不下。所以需要权衡。一种可行的配置是 TP8PP8DP1。总卡数 64。每张卡负责 70B/(8×8) 1.1B 参数FP16 下占 2.2GB。优化器状态用 ZeRO Stage 1但 DP1 意味着没有分片每张卡还是要存完整的优化器状态。这就矛盾了。实际中70B 模型的训练通常需要更多的卡或者用更激进的显存优化。比如用 ZeRO Stage 3 把参数也分片或者用 FP8 精度减少显存占用。这个案例说明并行策略的选择不是拍脑袋决定的需要根据硬件条件和模型大小仔细计算。6.3 实际训练中的调优过程假设我们最终选择了 TP8PP4DP2并且用了 ZeRO Stage 1 激活值重计算。训练启动后观察 GPU 利用率和吞吐量。如果发现 GPU 利用率只有 60% 左右说明有通信瓶颈。用 Profiler 一看发现 TP 的 AllReduce 占用了大量时间。这时候可以考虑减小 TP 度比如 TP4PP8DP2。但 PP8 意味着流水线更深气泡可能更大。需要增加微批次数量来减少气泡。调优的过程就是不断试错和权衡。每次调整后记录吞吐量和 GPU 利用率找到最优配置。这个过程可能需要多次迭代但一旦找到合适的配置训练效率会有明显提升。7. 一些踩过的坑和总结的经验7.1 TP 度选择的经验法则TP 度不是越大越好。我试过 TP16 的配置结果通信开销大到训练几乎跑不动。后来改成 TP8速度反而快了。经验法则是TP 度不要超过单机 GPU 数最好等于单机 GPU 数或者其一半。如果单机是 8 卡TP 选 4 或 8 比较合适。另外TP 度要和模型大小匹配。小模型用大 TP 度每张卡上的计算量太小通信占比过高效率反而低。大模型用小 TP 度单卡放不下也不行。所以要根据模型大小和硬件条件综合选择。7.2 PP 微批次数量的调优PP 的微批次数量直接影响流水线效率和显存占用。微批次太少气泡大效率低微批次太多显存占用高可能 OOM。我的经验是微批次数量设为 PP 度的 2 到 4 倍比较合适。比如 PP4微批次设为 8 到 16。这样气泡占比可以降到 10% 以下同时显存占用也不会太高。如果显存实在紧张可以减少微批次数量但要做好效率下降的准备。如果显存充裕可以增加微批次数量提高流水线效率。7.3 CP 和 TP 的配合技巧CP 和 TP 都涉及频繁通信配合使用的时候要注意通信组的划分。通常的做法是把 TP 和 CP 放在同一台机器内但分属不同的通信组。这样 TP 的 AllReduce 和 CP 的 AllGather 可以并行进行不会互相阻塞。另外CP 的 Ring Attention 实现要注意通信和计算的重叠。如果实现得好Ring Attention 的通信可以完全被计算掩盖几乎不增加额外开销。如果实现得不好通信会成为瓶颈。7.4 EP 负载均衡的实用技巧MoE 模型的 EP 场景下负载均衡是个持续的问题。除了调整负载均衡损失的权重还可以用一些工程手段来缓解。比如可以设置专家容量Expert Capacity限制每个专家最多处理多少 token。超出的 token 会被丢弃或者传到下一层。这样可以防止某个专家过载但可能会影响模型效果。另一个技巧是使用随机路由Random Routing在训练初期让 token 随机选择专家避免某些专家一开始就被过度选择。随着训练进行逐渐减小随机性让门控网络发挥作用。这些技巧都是我在实际训练中摸索出来的不一定适用于所有场景但可以作为参考。分布式训练没有银弹每个配置都需要根据具体情况调整。多试、多记录、多分析慢慢就能找到适合自己场景的最优配置。