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 优化器,优化器状态(一阶矩和二阶矩)还要占 28GB(FP32),加起来就是 56GB。这还没算激活值。一张 80GB 的 A100 勉强能放下,但如果你想把 batch size 调大一点,或者模型再大一点,单卡就放不下了。
所以 DP 的局限性很明显:它解决的是计算吞吐的问题,不解决单卡显存放不下模型的问题。当模型大到单卡装不下的时候,你就需要模型并行来帮忙了。
不过 DP 本身也在进化。现在常说的 DP 其实有两种形态:一种是传统的 DP,每张卡独立计算梯度然后 AllReduce;另一种是 ZeRO(Zero 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 和 1F1B(One 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。
具体来说,假设序列长度为 S,CP 度为 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,专家并行。它是专门为 MoE(Mixture 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 的模型。一种可能的配置是:TP=8,PP=4,DP=2。这样总卡数就是 8×4×2=64。TP=8 在单机内完成,PP=4 跨 4 台机器,DP=2 复制两份。
这种配置的逻辑是:TP 负责单机内的模型切分,利用 NVLink 的高带宽;PP 负责跨机的模型切分,通信量相对较小;DP 负责数据并行,提高吞吐量。三者各司其职,共同支撑起大模型的训练。
配置的时候有个经验法则:TP 度不要超过单机 GPU 数,PP 度不要超过机器数,DP 度尽量大。因为 DP 的通信效率最高(只需要 AllReduce 梯度),TP 的通信效率最低(每层都要 AllReduce),PP 居中。
3.2 长上下文场景:CP 的引入
当序列长度超过 8K 时,激活值的显存占用会急剧增加。这时候就需要引入 CP 来切分序列维度。
一个典型的配置是:TP=4,CP=2,PP=4,DP=2。总卡数 4×2×4×2=64。CP=2 把序列切成两半,每张卡负责一半的序列,注意力计算时通过 Ring Attention 交换 K、V。
CP 的引入会改变通信模式。原本 TP 的 AllReduce 是在 TP 组内进行的,现在 CP 组内也需要通信。实际配置中,通常会把 TP 和 CP 放在同一台机器内,因为它们的通信都比较频繁,需要高带宽支持。
3.3 MoE 模型:EP 的加入
MoE 模型的并行策略更复杂,因为多了 EP 这个维度。一个典型的 MoE 训练配置可能是:TP=4,EP=8,PP=2,DP=2。总卡数 4×8×2×2=128。
EP=8 意味着有 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 可以禁用 InfiniBand,NCCL_P2P_DISABLE 可以禁用 P2P 通信。排查问题的时候,可以先打开 NCCL_DEBUG=INFO,看看通信初始化是否正常。
注意事项:NCCL 的版本要和 CUDA 版本匹配,否则可能出现兼容性问题。另外,不同机器上的 NCCL 版本要一致,否则可能出现通信错误。
6. 从单卡到多卡:一个实际的配置案例
6.1 场景描述与硬件条件
假设我们有一个 70B 参数的模型,需要在一批 8 卡 A100 80GB 的机器上训练。总共有 8 台机器,64 张 GPU。序列长度 4K,batch size 目标是 512。
先算一下显存需求。70B 参数,FP16 精度下参数占 140GB,梯度占 140GB,Adam 优化器状态占 560GB(FP32),加起来 840GB。64 张卡平均下来每张卡 13GB 左右,看起来不多。但激活值的占用还没算,而且 TP、PP 的切分也会影响每张卡的实际占用。
6.2 并行策略选择与参数计算
基于前面的分析,我们选择 TP=8,PP=4,DP=2 的配置。TP=8 在单机内完成,利用 NVLink;PP=4 跨 4 台机器;DP=2 复制两份。
每张卡上的参数占用:70B / (8×4) = 2.2B 参数,FP16 下占 4.4GB。梯度同样 4.4GB。优化器状态如果不用 ZeRO,每张卡要存完整的 70B 优化器状态,那就爆了。所以必须用 ZeRO Stage 1,把优化器状态分片到 DP 组内。DP=2 的话,每张卡存一半的优化器状态,70B×4bytes/2 = 140GB,还是太大。
这里就看出问题了:DP=2 太小,ZeRO 的分片效果有限。如果把 DP 调大,比如 DP=8,那么 TP=8,PP=1,总卡数 64。但 PP=1 意味着没有流水线并行,所有层都在同一组卡上,显存可能放不下。
所以需要权衡。一种可行的配置是 TP=8,PP=8,DP=1。总卡数 64。每张卡负责 70B/(8×8) = 1.1B 参数,FP16 下占 2.2GB。优化器状态用 ZeRO Stage 1,但 DP=1 意味着没有分片,每张卡还是要存完整的优化器状态。这就矛盾了。
实际中,70B 模型的训练通常需要更多的卡,或者用更激进的显存优化。比如用 ZeRO Stage 3 把参数也分片,或者用 FP8 精度减少显存占用。这个案例说明,并行策略的选择不是拍脑袋决定的,需要根据硬件条件和模型大小仔细计算。
6.3 实际训练中的调优过程
假设我们最终选择了 TP=8,PP=4,DP=2,并且用了 ZeRO Stage 1 + 激活值重计算。训练启动后,观察 GPU 利用率和吞吐量。
如果发现 GPU 利用率只有 60% 左右,说明有通信瓶颈。用 Profiler 一看,发现 TP 的 AllReduce 占用了大量时间。这时候可以考虑减小 TP 度,比如 TP=4,PP=8,DP=2。但 PP=8 意味着流水线更深,气泡可能更大。需要增加微批次数量来减少气泡。
调优的过程就是不断试错和权衡。每次调整后,记录吞吐量和 GPU 利用率,找到最优配置。这个过程可能需要多次迭代,但一旦找到合适的配置,训练效率会有明显提升。
7. 一些踩过的坑和总结的经验
7.1 TP 度选择的经验法则
TP 度不是越大越好。我试过 TP=16 的配置,结果通信开销大到训练几乎跑不动。后来改成 TP=8,速度反而快了。经验法则是:TP 度不要超过单机 GPU 数,最好等于单机 GPU 数或者其一半。如果单机是 8 卡,TP 选 4 或 8 比较合适。
另外,TP 度要和模型大小匹配。小模型用大 TP 度,每张卡上的计算量太小,通信占比过高,效率反而低。大模型用小 TP 度,单卡放不下,也不行。所以要根据模型大小和硬件条件综合选择。
7.2 PP 微批次数量的调优
PP 的微批次数量直接影响流水线效率和显存占用。微批次太少,气泡大,效率低;微批次太多,显存占用高,可能 OOM。
我的经验是:微批次数量设为 PP 度的 2 到 4 倍比较合适。比如 PP=4,微批次设为 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 随机选择专家,避免某些专家一开始就被过度选择。随着训练进行,逐渐减小随机性,让门控网络发挥作用。
这些技巧都是我在实际训练中摸索出来的,不一定适用于所有场景,但可以作为参考。分布式训练没有银弹,每个配置都需要根据具体情况调整。多试、多记录、多分析,慢慢就能找到适合自己场景的最优配置。