ARTICLE DETAIL

资讯详情

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

大模型分布式训练必知:DP、TP、PP、CP、EP并行策略全解析

大模型分布式训练必知:DP、TP、PP、CP、EP并行策略全解析 很多算法同学第一次看到分布式训练配置时内心是崩溃的。满屏的并行参数TP、DP、PP、CP、EP光看缩写就劝退一半。更尴尬的是跟硬件同事聊 DP他第一反应是显示器接口说 TP他以为你在讲路由器。这篇关于 LLM 分布式计算的文章就是要把这些缩写一次性讲明白。作为系列第二篇我不打算堆术语而是从算法同学最容易理解的视角切入一个 Transformer 模型到底是怎么被拆到多张卡上的每种并行方式解决什么问题、付出什么代价。看完之后你至少能看懂训练框架的并行配置也能在模型跑不动的时候判断出该往哪个方向调。1. 先搞清楚分布式训练到底在干什么1.1 显存画像为什么 7B 模型在单卡上跑不起来在聊 TP、DP、PP 之前必须先建立显存的概念。大模型训练和推理的显存消耗是完全不同的两本账。训练侧以 7B 参数模型、混合精度FP16/BF16训练为例显存消耗大致如下项目计算方式显存占用FP16 权重7B × 2 字节约 14 GBFP16 梯度7B × 2 字节约 14 GBFP32 优化器状态Adam 的 master weight7B × 4 字节约 28 GBAdam 的动量momentum和方差variance2 × 7B × 4 字节约 56 GB激活值activation取决于 batch size、序列长度、层数等通常数十 GB 起步光前面四项加起来就已经超过 110 GB这还没算激活值。市面上主流单卡显存是 80 GB所以 7B 模型用单卡训练基本不可能这不是算力不够的问题是根本放不下。推理侧相对简单一些主要是权重加 KV Cache。以 70B 模型为例FP16 权重就要 140 GB单卡塞不下KV Cache 在长上下文场景下更是惊人。我算过一个数一个 80 层、GQA 8 组 KV head、head_dim 128 的模型每个 token 的 KV Cache 约为 2 × 80 × 8 × 128 × 2 字节约 320 KB。如果上下文是 4K、batch 是 32KV Cache 总量约为 40 GB。这个数字意味着推理时对大 batch、长上下文的处理同样必须依赖分布式方案。1.2 性能画像算力、带宽和通信开销模型放得下只是第一步还得跑得快。分布式训练的本质是在多张卡之间交换数据而数据交换的速度直接决定伸缩效率。这里有个核心矛盾GPU 的算力增长远快于卡间通信带宽的增长。NVLink 在单机内可以提供约 900 GB/s 的带宽但跨节点的 RDMA 网络通常只有 200-400 Gb/s换算下来只有 25-50 GB/s差了几乎一个数量级。所以选择并行策略时先要回答一个问题我切完模型之后每个计算步骤需要多少个跨卡通信原语如果通信原语太频繁、数据量太大计算可能被通信拖死出现“卡越多跑得越慢”的反常现象。这里顺带提一下集合通信原语。搞过 MPI 的同学应该不陌生分布式计算框架里最核心的几个操作是 AllReduce所有卡的数据先求和再广播给所有人、AllGather把分散在各卡的数据收集起来拼成完整张量、AlltoAll每张卡把自己的一部分数据发给其他所有卡。TP、DP、PP、CP、EP 的差异本质上就是组合了不同频率、不同数据量的集合通信操作。1.3 五种并行策略怎么分类建立分类框架比死记缩写重要得多。我的理解是所有并行策略要么在切数据要么在切模型。切数据DP数据并行切的是 batch 维度CP上下文并行切的是序列长度维度。切模型TP张量并行切的是单个算子内部的计算PP流水线并行切的是模型的层EP专家并行切的是 MoE 模型里的专家。这个分类能帮你快速定位问题如果你的模型权重太大优先考虑 TP、PP、EP如果权重放得下但训练速度慢、通过率低优先考虑 DP如果是长序列跑不动优先考虑 CP。下面逐个拆解。2. DP最直观的数据并行2.1 核心逻辑一人一份模型各算各的 batch数据并行是五种策略里最容易理解的假设你有 8 张卡每张卡都放一份完整的模型副本然后把一个 batch 的数据切成 8 份分给每张卡。每张卡独立做前向和反向计算算出各自的梯度之后通过 AllReduce 把所有卡上的梯度求和取平均再统一更新参数。DP 适合什么场景模型权重加上优化器状态在单卡上刚好放得下但单卡算力不够、训练速度太慢的场景。它的优点是实现简单不需要改动模型结构数据加载和梯度同步的框架都已经非常成熟。PyTorch 的 DistributedDataParallelDDP几乎是开箱即用这也是很多算法同学接触的第一个分布式组件。但 DP 有个天然问题每张卡都存一份完整模型显存冗余非常严重。如果模型本身就接近单卡显存上限DP 就无能为力了。另外梯度 AllReduce 的通信量跟模型参数量成正比模型越大通信开销越大。我在实际项目里见过一个 70B 模型的训练任务DP 维度拉到 64 之后通信占比超过了 40%算力浪费非常严重。2.2 梯度同步背后的 AllReduce 机制DP 的灵魂在梯度同步也就是 AllReduce。这里我建议算法同学至少理解一下 Ring AllReduce 的思路它把 N 张卡排成一个环先把梯度数据切分成 N 份第一步做 ReduceScatter——每张卡把自己的一部分梯度发给下一张卡累加后继续传递经过 N-1 轮后每张卡拥有完整求和梯度的一部分第二步做 AllGather——把求和好的分片再沿环转一圈让每张卡都获得完整梯度。这种设计的巧妙之处在于通信量不随卡数线性增长而是近似等于 2 倍的单卡梯度数据量。所以理论上 DP 的扩展性相当好。但工程上要小心一个坑小张量频繁 AllReduce 的效率非常低。因为每次通信都有延迟开销张量太小的话延迟会占据主导。为了解决这个问题现代框架普遍使用梯度桶gradient bucket机制把多个小梯度张量拼接成一个大的连续缓冲区凑满一定大小再触发一次 AllReduce。PyTorch DDP 里可以通过bucket_cap_mb参数控制桶的大小默认是 25 MB。我实测下来对于百亿参数模型把桶调大到 50-100 MB 往往能进一步提升吞吐但具体最优值跟网络带宽和模型结构有关需要自己测。2.3 从 DP 到 ZeRO显存冗余的破解思路既然 DP 的问题在于每张卡都存一份完整模型状态那很自然的想法就是能不能把这份冗余去掉DeepSpeed 提出的 ZeROZero Redundancy Optimizer就是这个思路。它把 DP 中的优化器状态、梯度、甚至参数本身做分片每张卡只保存一部分。具体分三个阶段ZeRO Stage 1把优化器状态分片每张卡只负责一部分参数的优化器更新。显存需求从 16Φ 字节降到约 4Φ 12Φ/卡数。ZeRO Stage 2在 Stage 1 基础上再把梯度分片进一步降低显存。ZeRO Stage 3参数也分片前向计算时通过 AllGather 临时取回需要的参数。此时单卡已经不需要存完整模型。用 7B 模型举例标准 DP 需要约 112 GB 显存ZeRO Stage 3 在 8 卡环境下可能只需要 20-30 GB 每卡。这几乎是质变。不过 ZeRO Stage 3 也是有代价的参数 AllGather 增加了通信量训练吞吐通常会下降。所以工程里常见做法是单机内用 TP 把模型切开卡间有 NVLink带宽充足跨机之间用 ZeRO 做分片扬长避短。这个组合后面第六部分会展开。3. TP把一个算子拆开算3.1 一个线性层怎么切成两半张量并行TP是另一种完全不同的思路模型不是复制多份而是把单个计算算子切分到多张卡上。以最常见的线性层 Y XW 为例切法有两种。按列切分权重 W把 W 沿输出维度切成两块每张卡持有一部分列。每张卡独立计算 XW_i得到部分输出最后用 AllGather 把结果拼起来。因为 Y 的每一列只依赖 W 对应的列这种做法非常自然。按行切分权重 W把 W 沿输入维度切成上下两块同时把输入 X 沿特征维度切成对应的两半。每张卡算 X_iW_i得到的是部分和最后需要 AllReduce 把结果累加。为什么 TP 能把单卡放不下的模型跑起来因为权重被拆开了显存占用随之分摊。7B 模型用 TP8 切分后每卡只需要存不到 2 GB 的权重未算优化器状态效果立竿见影。3.2 通信量和带宽敏感度分析TP 是五种并行策略中对通信带宽最敏感的一种原因在于通信太频繁。在 Transformer 的每一层里注意力输出和 MLP 输出各需要一次 AllReduce也就是说每层至少两次全卡通信。模型有几十层每一层的前向和反向都要通信频率极高。这就带来一个硬约束TP 的卡必须尽量放在同一台机器内部通过 NVLink 互联。跨节点做大 TP 我基本不推荐因为网络带宽不够通信时间会远超计算时间吞吐直线下降。一个经验值是在 8 卡 A100/H100 节点内TP8 的效果通常不错一旦超过 8如果节点内没有更高密度的互联收益就会明显递减。TP 的另一个特性是内存占用均衡。因为计算是同步的——所有卡在同一个算子内协同工作所以显存消耗天然是均匀的。这跟后面要讲的 PP 很不一样PP 会出现不同设备显存水位不一致的问题。3.3 Transformer 里的 TP 布局具体到 Transformer 模型TP 的切法有成熟范式以 Megatron-LM 的经典方案为例。注意力的 QKV 投影矩阵用列并行切分因为这样可以天然地把不同注意力头分到不同卡上注意力输出投影用行并行MLP 的第一个线性层用列并行加激活函数第二个线性层用行并行。这样的布局保证每一层两个线性块之间只做一次 AllReduce通信次数被压到最低。不过 TP 在实现上有个比较隐蔽的坑LayerNorm 和 Dropout 这类按特征维度操作的层在 TP 切分后需要额外的同步逻辑。以 LayerNorm 为例它需要对最后一个维度做归一化而最后一个维度被切到多张卡上每张卡手里只有部分特征算不出完整的均值和方差。所以 Megatron 里引入了序列并行Sequence Parallel的概念在 LayerNorm 和 Dropout 前后插入通信让这类操作只在一张卡上执行而不是重复计算。这个设计是后来 CP 的基础之一理解它有助于后面读上下文并行的代码。4. PP按层切分的流水线4.1 为什么有了 TP 还要 PP既然 TP 能切算子显存问题似乎解决了为什么还要 PP核心原因是扩展性。TP 的通信太频繁一旦超过单节点范围性能迅速恶化。所以当模型大到连一个 8 卡节点都放不下时例如万亿参数模型就必须在更粗的粒度上切开——按层切。流水线并行PP的思路非常直观把 Transformer 的 L 层分成若干段每张卡负责其中一段。GPU 0 计算第 1-2 层GPU 1 计算第 3-4 层前向传播像流水线一样前一段算完把中间激活传给下一段。这样显存压力被线性分摊而且卡间通信是低频的——只在段与段之间传输激活值和梯度频率远低于 TP。代价是什么GPU 之间存在严重的不平衡等待。用生活化的类比一个工厂生产线有 4 个工位每个工位处理一个环节。第一个工件上线后前面的工位忙后面的工位闲着收尾阶段反过来。这种空闲时间称为流水线气泡bubble是 PP 最核心的敌人。4.2 气泡 bubble 怎么算理解气泡大小的公式很重要它直接决定你要不要用 PP。假设段数为 Pmicro-batch 数量为 M理想状态下气泡占比约为bubble 占比 ≈ (P - 1) / (M P - 1)举例P4M16 时气泡占比 3/19 ≈ 15.8%P8M16 时气泡占比 7/23 ≈ 30.4%。可以看出段数越多、micro-batch 越少气泡越严重。所以工程上两个原则一PP 的段数不要太大一般不超过 8二PP 必须配合足够的 micro-batch 来填满流水线。这里需要解释一下 micro-batch 的概念。PP 不是简单地把一个 batch 的样本按层切分而是把一个大的 batch 切成多个小份称为 micro-batch。每个 micro-batch 依次进入流水线前一个还在 GPU 0 上算后一个已经可以跟上流水线才能转起来。micro-batch 数量越接近段数气泡占比越高反之越多气泡越小。4.3 micro-batch 调度和 1F1B 机制早期 PP 的 naive 实现是让所有 micro-batch 先完整跑完前向再开始反向。这个方案实现简单但显存爆炸所有层的中间激活都要存下来等反向计算使用。百亿参数模型下激活值可能占掉大半显存。后来业界普遍采用 1F1Bone-forward-one-backward调度每张卡交替执行前向和反向计算做到“算一个 micro-batch 的前向就释放一份旧的激活再算下一个”。这样能显著压低峰值显存。这个调度在 Megatron-LM 和 DeepSpeed 里都是默认实现算法同学不需要手写但理解它有助于解释为什么同一个模型PP 配置不同显存表现会差很多。PP 还有一个容易忽略的坑反向传播和参数更新的时机。因为在流水线里不同 micro-batch 的梯度计算有先后如果等所有 micro-batch 都跑完再统一更新参数会浪费显存如果每个 micro-batch 跑完就更新权重会频繁变化影响收敛稳定性。主流做法仍然是累积梯度到一定步数再更新这时候 batch size 的实际含义跟 DP 的 global batch 不一样调学习率时要格外小心。5. CP 和 EP长序列与 MoE 的专用并行5.1 上下文并行 CP把长序列拆开如果说 DP 切 batch、TP 切算子、PP 切层那 CP上下文并行切的则是序列长度维度。CP 的诞生背景非常明确长序列场景。当输入序列达到 128K、1M token 时即使权重放得下激活值和 KV Cache 也会把显存撑爆。比如前面算过4K 上下文、32 batch 的 KV Cache 就要约 40 GB如果是 128K 上下文这个数字会直接翻 32 倍到 1.2 TB 以上。CP 的核心思路是把整个序列切分成多段每张卡负责一段序列的计算。但这里有个技术难点标准注意力机制中每个 token 都要跟序列中所有 token 计算注意力分数序列被切开后跨段的注意力依赖怎么处理答案是 Ring Attention。每张卡持有自己那段序列的 Q同时循环传递 K/V 块。卡在计算本地注意力分数的同时把本地 K/V 传给下一张卡再从上一张卡接收另一段的 K/V。经过 N-1 轮循环后每张卡都见过所有序列段的 K/V也就计算出了完整的注意力输出。整个过程像流水线一样持续滚动让“序列再长也不怕”从理论上变成了可能。CP 和 TP 里的 Sequence Parallel 容易混淆但两者目标不同SP 是为了解决 LayerNorm 等操作在特征维度切分后的计算问题切的是特征维度CP 切的是序列长度维度SP 是 TP 的配套技术CP 则是为超长序列准备的正交方案。实际训练中两者经常叠加使用。5.2 专家并行 EPtoken 怎么“找”专家专家并行EP是专门为 MoEMixture of Experts模型设计的。理解 EP 前先要理解 MoE 模型的基本结构每一层不再是单一 MLP而是多个并行的 MLP专家外加一个路由器router。每个 token 经过路由器计算后只会被发往 top-k 个专家比如 k2表示每个 token 由两个专家处理。问题来了如果所有专家都放在一张卡上显存压力巨大。Mixtral 8x7B 这种 8 个专家的模型单层 MLP 就占掉了大部分显存普通卡根本放不下。EP 的做法是把不同专家分布到不同设备上token 通过 AlltoAll 通信被发送到目标专家所在的设备。专家计算完毕后再通过一次 AlltoAll 把结果送回原设备。EP 的优势是显存分散和计算并行每个专家只在一张卡上按需激活闲置专家不占算力。缺点是通信代价极高。AlltoAll 是集合通信里开销最大的原语之一因为它不是“广播给所有人”或“收集整合”而是“每张卡都有一份数据要发给其他所有卡”通信量随卡数呈指数级别增长。另一个隐患是负载均衡如果一批 token 恰好大量路由到同一个专家那张卡就会成为瓶颈。所以 MoE 训练通常需要额外的负载均衡 loss变相迫使 router 把 token 分配得更均匀。5.3 两种并行的适用场景小结CP 和 EP 的应用场景高度明确CP 用于长序列训练的显存优化和推理时的 KV Cache 分片EP 用于 MoE 模型的训练和推理。对于普通稠密模型CP 和 EP 不需要也不应该强上。我见过有人给一个 13B 的稠密模型硬配置了 EP8结果是路由逻辑白跑、通信白开销、性能不升反降。在工程里没有无代价的并行策略只有适合当前场景的策略。6. 实际工程里怎么组合和选择6.1 训练场景的组合Megatron 3D 并行 ZeRO真实的大模型训练极少只用一种并行策略而是把多种策略组合起来。经典组合是 Megatron 提出的 3D 并行TP × PP × DP。假设你有 32 张卡可以这样规划先按 TP8 切分模型因为 8 卡在单节点内NVLink 带宽充足通信开销可控再按 PP2 把模型按层分成两段降低跨节点通信频率剩下 16 张卡上做 DP也就是每个模型副本处理不同 batch。总卡数 TP × PP × DP 8 × 2 × 2 32。这里有个简单的产品逻辑先决定 TP因为它决定单节点的模型能否放下再决定 PP因为它决定跨节点切分粒度最后决定 DP因为你希望尽量把数据并行维度做大来提高整体吞吐。同时可以叠加 ZeRO Stage 1 或 2 来削减优化器状态的冗余。我自己在实际训练里的体会是先跑一个小规模配置用 profiler 看通信和计算的占比再决定是否调整。如果不看数据直接抄大厂的配置很容易在自家集群上翻车。每个集群的拓扑、带宽、卡型都不同最合适的并行比例也完全不同。6.2 推理场景的组合TP PP CP/EP推理时的显存结构和训练完全不同没有优化器状态但有巨大的 KV Cache。权重用 TP 切分后可以用多卡分摊显存KV Cache 则需要用 CP 按序列维度分片这样长上下文场景下才能撑住更大的 batch。推理性价比最高的方式是权重做 TP 分片KV Cache 做 CP 分片。对于 MoE 模型推理侧 EP 也很常见因为推理时专家计算量小但数量多EP 可以把不同专家放在不同卡上让每次推理只激活需要的专家减少整体显存占用。要注意的是推理侧对延迟的敏感度远超训练。训练关心的是吞吐每秒处理多少 token推理关心的是首 token 延迟和单 token 延迟。TP 虽然可以减少单次计算的显存但每层两次 AllReduce 的延迟会直接影响推理时延PP 的流水线气泡在推理时也不能忽略。所以推理框架如 vLLM、TensorRT-LLM的并行策略通常比训练更保守很多场景下 TP8 就是上限极少看到 PP 在在线推理里拉得太高。6.3 并行策略选择速查表并行策略切分维度核心通信原语显存效果通信频率适用场景DPbatchAllReduce梯度不降权重显存每步一次总量随模型增大单卡放得下模型吞吐不足TP算子内部AllReduce / AllGather权重分摊每层 2 次以上单卡放不下节点内 NVLinkPP模型层间P2P Send/Recv权重分摊层与层之间低频超大模型跨节点切分CP序列长度P2P 循环传递激活和 KV Cache 分摊序列轮转超长序列训练/推理EP专家AlltoAll专家权重分摊每次 token 路由MoE 模型训练和推理选型时我一般先问三个问题模型多大单卡能放下吗卡间是什么互联如果模型权重加优化器状态在单卡上能放下但速度慢DP 是最优解放不下但同位卡数不多用 TP模型大到单节点装不下加 PP序列极长加 CP模型是 MoE加 EP。这并不是什么高深理论而是取舍逻辑每引入一种并行策略都意味着引入一种通信开销你要确保收益大于代价。6.4 给算法同学的具体建议作为算法同学我对 Infra 的态度一直是“够用就行但必须能定位问题”。这里分享几个实际项目里的心得。第一先会估算显存再谈配置。跑任务之前把前面那张显存画像表拉出来手动估算一遍权重、梯度、优化器状态、激活值/KV Cache 分别占多少心里就有底了。很多显存爆掉的问题根本不需要加卡把 batch size 调小或者打开梯度检查点activation checkpointing就能解决。第二善用 profiler 而不是瞎猜。训练任务跑起来后用 PyTorch Profiler 或者 NVIDIA 的 Nsight Systems 看一下通信占比。如果通信占比超过 30%优先怀疑并行配置是否合理如果 GPU 利用率低但通信占比不高可能问题在数据加载或者算子实现上。第三从“小模型小并行”开始验证。不要一上来就 64 卡训练大模型先在 8 卡上用小配置跑通流程逐步加并行维度。每次只动一个变量对比吞吐变化。这比什么都改、最后出了问题不知道是哪一步导致的要高效得多。第四多关注集群的物理拓扑。同样的 TP8在 8 卡 NVLink 全互联的节点上和 2 台 4 卡机器用高速网络连起来的场景里性能差距可能是倍数级的。配置并行策略前先了解卡和卡之间的实际带宽这个信息在算法同学日常工作中经常被忽略但恰恰是分布式计算性能的底层决定因素。最后再分享一个小技巧打印一下每张卡的显存峰值。跑完同一个任务如果发现各卡显存占用严重不均说明并行策略里的负载均衡有问题——可能是 PP 分段不均也可能是 MoE 的 token 路由倾斜。这种情况优先调并行配置而不是简单地加卡。多做几次这样的分析你对 TP、DP、PP、CP、EP 的理解会比只看文档深刻得多。
返回列表