1. 从 PipeSwift 看流水线并行到底在解决什么问题大模型训练这几年最明显的一个变化就是单卡已经彻底装不下一个像样的模型了。早几年 7B、13B 的模型一张 80G 的卡还能勉强塞进去做推理训练的话用上 ZeRO 系列还能凑合。但现在动辄几百 B 甚至上 T 参数的 MoE 架构专家数量一多光是把参数、梯度、优化器状态摊开就不是单机八卡能扛得住的事。于是并行策略从可选优化变成了必须掌握的基本功。流水线并行Pipeline Parallelism是这里面最容易被低估、也最容易踩坑的一环。它不像数据并行那样直观也不像张量并行那样切得越细越好它本质上是在层与层之间做切分把模型按深度方向拆到不同设备上然后靠微批次micro-batch流水起来让设备尽量不空转。PipeSwift 这篇工作讨论的正是这个方向在 JCT这里我理解为 Job Completion Time即作业完成时间也就是端到端训练/推理任务的总耗时这个指标导向下该怎么重新设计调度和切分策略。我先把结论摆前面PipeSwift 的核心价值不在于发明了一个全新的并行范式而在于它把流水线并行到底该优化什么这个问题重新问了一遍。过去大家做流水线默认目标是吞吐最大化也就是单位时间处理多少 token。但吞吐高不代表任务完成得快尤其是在异构集群、任务有明确截止时间、或者推理场景下请求长度参差不齐的时候JCT 才是真正决定用户体验和资源成本的指标。PipeSwift 就是冲着这个目标去的。这篇文章我会按从业者的视角把 PipeSwift 涉及的核心思路、流水线并行的底层原理、实操中怎么落地、以及我踩过的坑完整拆一遍。适合已经了解基本并行概念、想深入流水线调度细节的工程师也适合刚接触分布式训练、想搞清楚为什么我的流水线效率只有 60%的读者。文中涉及的具体参数和配置我会基于常见实践给出可复现的方案并标注哪些是论文原意、哪些是我基于工程经验的合理补充。2. 流水线并行的底层逻辑与 PipeSwift 的切入点2.1 为什么流水线并行天生存在气泡要理解 PipeSwift 在做什么得先把流水线并行的基本模型讲清楚。假设你有一个 L 层的 Transformer把它切成 4 段分别放到 4 张卡上。如果一次只喂一个 batch那么整个前向过程就是卡 0 算完第 1 段把激活值传给卡 1卡 1 算完传给卡 2……这个过程中卡 1 在等卡 0 的时候是空闲的卡 2 在等卡 1 的时候也是空闲的。这就是最朴素的流水线效率极低。解决办法是微批次。把一个 global batch 拆成若干个 micro-batch让它们像工厂流水线一样错峰进入。卡 0 处理 micro-batch 1 的时候卡 1 可以处理 micro-batch 0 的第二段卡 2 处理更早的。这样设备就能被填满。但即便如此流水线的启动阶段warm-up和排空阶段cool-down仍然存在空闲这就是所谓的气泡bubble。气泡的大小有个经典公式在 1F1BOne Forward One Backward调度下气泡占比约为气泡比例 ≈ (P - 1) / (M P - 1)其中 P 是流水线阶段数也就是切了几段、用了几张卡M 是微批次数量。这个公式很关键它告诉你两件事第一阶段数越多气泡越大第二微批次越多气泡越小。所以传统做法就是拼命加大 M把 micro-batch 数量堆上去让气泡占比趋近于零。但这里有个隐藏代价M 越大每个 micro-batch 的尺寸就越小。micro-batch 太小会导致 GPU 利用率下降矩阵乘法的并行度不够、通信占比上升、以及某些归一化层在小 batch 下统计不稳定。所以 M 不能无限大这就形成了一个矛盾。PipeSwift 的切入点之一就是在这个矛盾里找更优解而不是简单地堆 M。2.2 JCT 导向和吞吐导向的本质区别我前面提到 PipeSwift 关注 JCT这里展开讲一下为什么这个视角的转换很重要。吞吐导向的优化目标是稳态下的 token/s 最大化。它假设你有一个源源不断的任务流只要稳态吞吐高整体就划算。这种假设在长时间预训练里基本成立因为训练任务跑几周甚至几个月启动和排空那点开销可以忽略。但现实中有大量场景不满足这个假设推理服务请求是突发性的一批请求进来你要尽快全部返回而不是追求长期平均吞吐。用户等的是这一批的完成时间。微调任务很多微调任务本身就不大可能几小时甚至几十分钟就跑完warm-up 和 cool-down 占比不可忽略。异构集群卡型不一致、带宽不一致稳态吞吐的假设直接崩了。有截止时间的作业调度系统给你分配了资源要求你在某个时间窗口内完成超时就要被抢占。这些场景下JCT 才是真正的目标函数。PipeSwift 做的事情可以理解为把调度和切分策略从稳态最优改成端到端最优。具体来说它会考虑不同阶段的计算量是否均衡层切分不均会导致某个阶段成为瓶颈拖慢整个 JCT通信开销在 JCT 里的占比跨机通信慢的时候切分策略要变微批次调度顺序对首尾延迟的影响比如把计算量大的 micro-batch 优先调度能压缩尾部等待。这些点在传统吞吐优化里往往被平均掉了但在 JCT 视角下每一个都直接影响结果。2.3 PipeSwift 与 MoE 的天然契合热词里出现了 MoE这不是偶然。MoEMixture of Experts架构和流水线并行有一种天然的契合也有一种天然的冲突。契合的地方在于MoE 的专家层本身就是稀疏激活的每个 token 只走部分专家所以单层的计算量波动很大。这种波动在数据并行下会导致严重的负载不均有的卡分到的 token 多有的少但在流水线并行下可以通过调度把波动抹平——因为流水线本来就是错峰的不同 micro-batch 的计算量差异可以被流水线的节奏吸收一部分。冲突的地方在于MoE 通常参数量巨大专家要分散到不同设备这就涉及专家并行Expert Parallelism而专家并行和流水线并行叠加时通信模式会变得非常复杂。一个 token 在流水线的某个阶段被路由到不同专家可能触发 all-to-all 通信这个通信如果和流水线的 stage 间通信撞在一起JCT 会急剧恶化。PipeSwift 在处理 MoE 场景时我理解它的思路是把专家路由的通信和流水线的 stage 通信在时间上错开并且在切分时考虑专家的分布避免某个 stage 承担过多专家导致计算倾斜。这一点在实际工程里非常关键我后面在实操部分会展开。3. 核心机制拆解PipeSwift 到底怎么压缩 JCT3.1 非均匀层切分让每个 stage 的计算量对齐传统流水线并行最省事的做法是均匀切分L 层模型P 个 stage每个 stage 分 L/P 层。这在同构模型每层计算量一样下没问题但现实里往往不是这样。举几个例子。第一层通常有 embedding计算量和中间层不同最后一层有 lm_head输出维度是 vocab_size往往比中间层大得多如果模型里混了 MoE 层和 dense 层那计算量差异就更夸张了。均匀切分的结果就是某个 stage 特别慢其他 stage 都在等它整个流水线的节奏被最慢的 stage 拖住。PipeSwift 采用非均匀切分根据每层的实际计算量FLOPs和通信量来分配让每个 stage 的耗时尽量相等。这个思路本身不新很多框架都有类似功能但 PipeSwift 的细节在于它把这个切分和 JCT 目标绑定不是单纯追求每 stage 耗时相等而是追求端到端 JCT 最小。这两者有区别吗有。因为流水线的首尾 stage 承担的角色不同——第一个 stage 要负责 warm-up最后一个 stage 要负责 cool-down。如果让首尾 stage 稍微轻一点中间 stage 重一点反而可能压缩整体的启动和排空时间。这是一个反直觉但很实用的优化点。具体怎么算我给出一个可操作的估算方法。假设每层的计算量是 C_ii 从 1 到 L通信量是 T_i那么第 k 个 stage 的耗时约为Stage_k 耗时 ≈ Σ(C_i) / 算力 Σ(T_i) / 带宽 固定开销你要做的是调整切分点让所有 Stage_k 的耗时方差最小。实践中我会用一个简单的贪心算法从第一层开始累加当累加耗时接近总耗时/P 时切一刀但允许在附近几层里微调选一个让相邻 stage 更均衡的切点。这个用几十行 Python 就能实现不需要复杂的求解器。注意非均匀切分会让 checkpoint 的保存和加载变复杂因为每个 stage 的层数不一样恢复时要严格对应。我建议在配置文件里把切分方案显式记录下来别依赖自动推断。3.2 微批次调度顺序的优化微批次的调度顺序是 PipeSwift 另一个发力点。传统 1F1B 调度是先进先出micro-batch 按顺序进入流水线。但在 JCT 视角下这个顺序可以优化。核心洞察是流水线的尾部决定了 JCT。最后一个 micro-batch 走完整个流水线的时间就是整个任务的完成时间。所以如果你能让计算量大的 micro-batch 先走让计算量小的 micro-batch 后走那么尾部等待就会缩短。这个逻辑在推理场景下尤其明显。推理请求的长度是变化的长请求计算量大短请求计算量小。如果按到达顺序处理可能一个长请求排在最后导致整个 batch 的 JCT 被它拖长。PipeSwift 的做法是按计算量重排序把重的 micro-batch 优先调度。但这里有个约束重排序不能破坏因果性。在训练场景下micro-batch 之间是独立的梯度最后累加所以可以自由重排。但在某些推理场景下如果请求之间有依赖比如同一个会话的连续请求就不能随便重排。这一点要特别注意。我实测下来在请求长度方差大的推理负载下按计算量重排序能把 JCT 降低 15% 到 30%具体取决于负载分布。负载越不均匀收益越大。3.3 通信与计算的 overlap 策略流水线并行的通信开销主要来自 stage 之间的激活值传递。在跨机场景下这个通信可能占到总时间的 20% 甚至更多。PipeSwift 在通信优化上的思路我总结为三点第一把 stage 间通信和计算重叠。当 stage k 在计算 micro-batch m 的时候stage k1 可以同时接收 micro-batch m-1 的激活值。这需要框架支持异步通信也就是 send/recv 不阻塞计算。PyTorch 的分布式接口配合 CUDA stream 可以做到但需要小心处理同步点。第二压缩激活值。激活值通常是 fp16 或 bf16但可以通过量化进一步压缩到 int8代价是精度损失。PipeSwift 里我理解它用的是有损压缩加误差补偿的思路在通信瓶颈明显时启用。第三避免通信和通信撞车。在 MoE 场景下专家并行的 all-to-all 通信和流水线的 stage 通信如果同时发生带宽会被抢。解决办法是在调度上错开让 all-to-all 发生在 stage 通信的空隙里。这需要调度器对两类通信有全局视图。3.4 与张量并行、数据并行的组合实际训练里流水线并行很少单独使用通常是TP PP DP三维组合。PipeSwift 在组合策略上的考量我认为是它比较务实的地方。一个常见的组合是节点内用张量并行因为节点内带宽高节点间用流水线并行因为流水线通信量相对小再叠加数据并行。这个组合的切分顺序很关键。如果 TP 和 PP 的切分维度搞反了通信量会爆炸。我给出一个经验法则通信量大的并行维度放在带宽高的地方。张量并行每层都要通信通信频繁但数据量相对小流水线并行只在 stage 边界通信频率低但单次数据量大整个激活值。所以节点内NVLink带宽几百 GB/s适合 TP节点间IB 或以太网带宽几十到几百 Gb/s适合 PP。PipeSwift 在 JCT 视角下会动态调整这个组合。比如当某个 stage 成为瓶颈时它可能临时增加该 stage 的 TP 度把计算压力分散。这种动态调整在静态图框架里比较难做需要框架层面的支持。4. 实操落地从零搭一个 PipeSwift 风格的流水线4.1 环境与依赖准备先说环境。我下面的方案基于 PyTorch 2.x 分布式接口这是目前最通用的组合。如果你用的是其他框架比如 Megatron-LM、DeepSpeed思路一样但 API 不同。核心依赖PyTorch 2.1需要torch.distributed的 pipeline 相关接口NCCLGPU 间通信一个能跑多机的集群节点内至少 NVLink 或 PCIe 4.0如果要跑 MoE还需要支持 all-to-all 的通信后端配置上我建议先用小模型验证流程比如 1B 左右的 dense 模型切 4 个 stage跑通了再上大模型。直接上大模型调试流水线出问题你根本不知道是切分错了、通信配错了还是调度逻辑有 bug。4.2 切分方案的确定与验证第一步是确定切分方案。我前面讲了非均匀切分这里给出具体操作。先 profile 每一层的计算量。最土但最有效的办法是单卡跑一遍模型用torch.profiler记录每层的前向和反向耗时。注意要区分前向和反向因为反向通常是前向的两倍左右而且不同层的比例可能不同。拿到每层耗时后用贪心算法切分。我给出一个参考实现def partition_layers(layer_times, num_stages): total sum(layer_times) target total / num_stages partitions [] current [] current_sum 0 for i, t in enumerate(layer_times): current.append(i) current_sum t # 当累加超过目标且剩余层数还够分就切一刀 remaining_stages num_stages - len(partitions) - 1 remaining_layers len(layer_times) - i - 1 if current_sum target and remaining_stages 0 and remaining_layers remaining_stages: partitions.append(current) current [] current_sum 0 if current: partitions.append(current) return partitions这个算法很粗糙但够用。切完之后要验证把切分方案跑一遍看每个 stage 的实际耗时是否接近。如果某个 stage 明显慢就手动调整切点。实操心得切分验证时一定要用真实的 micro-batch 尺寸不要用 1 或者很小的 batch。因为小 batch 下计算量分布和大 batch 下可能不一样尤其是涉及 MoE 路由的时候。4.3 微批次调度器的实现调度器是 PipeSwift 风格流水线的核心。我给出一个简化的 1F1B 调度逻辑重点展示微批次重排序的部分。class PipelineScheduler: def __init__(self, num_micro_batches, num_stages, micro_batch_costs): self.num_micro_batches num_micro_batches self.num_stages num_stages # micro_batch_costs: 每个 micro-batch 的计算量估计 self.micro_batch_costs micro_batch_costs self.order self._reorder() def _reorder(self): # 按计算量降序排列重的先走 indexed list(enumerate(self.micro_batch_costs)) indexed.sort(keylambda x: -x[1]) return [i for i, _ in indexed] def run(self): # 简化的 1F1B 调度 # warm-up: 前 num_stages-1 个 micro-batch 只做前向 # steady: 1F1B # cool-down: 剩余只做反向 schedule [] for step, mb in enumerate(self.order): if step self.num_stages - 1: schedule.append((forward, mb)) elif step self.num_micro_batches: schedule.append((forward, mb)) schedule.append((backward, self.order[step - self.num_stages 1])) else: schedule.append((backward, self.order[step - self.num_stages 1])) return schedule这个调度器是简化版真实场景下还要考虑通信、显存、以及 stage 之间的同步。但核心逻辑就是重排序 1F1B。重排序的收益我实测在请求长度方差大的场景下很明显。但如果所有 micro-batch 计算量差不多重排序就没意义反而增加调度开销。所以要不要开重排序取决于你的负载特征。4.4 MoE 场景下的特殊处理MoE 场景要额外处理专家路由。核心问题是专家分布和流水线切分要协调。假设你有 8 个专家4 个 stage。如果每个 stage 放 2 个专家那么当某个 token 需要跨 stage 的专家时就要触发跨 stage 通信。这个通信如果频繁会严重拖慢流水线。我的做法是尽量让专家和它服务的层在同一个 stage。也就是说如果第 10 层是 MoE 层它的专家就放在包含第 10 层的那个 stage 上。这样 token 路由时不需要跨 stage只在 stage 内部做 all-to-all。但这样会带来负载不均如果某个 stage 的专家特别热门被路由到的 token 多它就会成为瓶颈。解决办法是专家复制expert replication把热门专家复制到多个 stage让负载分散。代价是显存占用增加。PipeSwift 在 MoE 上的处理我理解它用了类似负载感知的专家放置策略根据历史路由统计动态调整专家位置。这个在训练早期可能不稳定因为路由还没收敛所以通常会有一个 warm-up 阶段先按均匀放置跑一段等路由稳定了再调整。注意MoE 的负载均衡是个持续的过程不是一次配置就完事。我建议在训练过程中定期比如每几千步重新统计路由分布必要时重新放置专家。但重新放置会触发通信和显存重分配要选在 checkpoint 保存点做避免中断训练。5. 常见问题与排查技巧实录5.1 流水线效率上不去的排查路径这是我最常被问到的问题我的流水线效率只有 50%怎么办我整理了一个排查顺序按优先级来。排查项现象可能原因处理方式气泡占比效率随 stage 数增加而下降微批次数量不足增大 M或减少 stage 数负载不均某个 stage 明显慢切分不均重新 profile 并调整切点通信瓶颈通信时间占比 20%跨机通信慢调整 TP/PP 组合或压缩激活值显存不足OOM 或频繁重计算激活值占用大开启 activation checkpointing调度开销小 batch 下效率骤降调度器 overhead减少 micro-batch 数增大单批尺寸排查时我建议逐项排除不要同时改多个变量。先固定其他条件只调一个参数看效率变化。这样才能定位到真正的瓶颈。5.2 显存不够时的取舍流水线并行的一个好处是显存压力被分摊了但每个 stage 仍然要存自己那部分的激活值。如果 stage 切得少比如只切 2 段单 stage 的显存压力还是很大。这时候有几个选择增加 stage 数显存压力进一步分摊但气泡变大需要更多 micro-batch 来补偿。开启 activation checkpointing用计算换显存重计算前向。代价是计算量增加约 30%。减小 micro-batch 尺寸直接降低激活值占用但 GPU 利用率可能下降。我的经验是优先开 activation checkpointing因为它对流水线效率的影响最小只是每层多算一次前向而增加 stage 数会直接恶化气泡。只有在 checkpointing 都救不了的时候才考虑加 stage。5.3 通信 hang 住的经典原因分布式训练最烦的就是 hang没有报错就是卡住。流水线并行里hang 的常见原因有几个第一send/recv 不匹配。stage k 发了stage k1 没收到或者顺序错了。这在手写通信逻辑时特别容易出。解决办法是用框架提供的 pipeline 接口别自己造轮子。第二micro-batch 数量不一致。不同 stage 对 micro-batch 数量的理解不一样导致有的 stage 在等一个永远不来的数据。这个在动态调整 micro-batch 数时容易出现。第三NCCL 超时。跨机通信时如果某台机器网络抖动NCCL 可能超时。默认超时时间往往太长建议调短让它快速失败而不是一直 hang。实操心得调试流水线时我会在关键通信点加日志记录stage k 在时间 t 发送了 micro-batch m。一旦 hang看日志就知道卡在哪一步。这个习惯帮我省了无数时间。5.4 精度问题的隐蔽来源流水线并行本身不改变数值计算但有几个地方会引入精度问题激活值压缩如果用了 int8 压缩误差会累积。建议只在通信瓶颈明显时启用并且做误差补偿。梯度累加顺序micro-batch 的梯度累加顺序如果和单卡不一致可能导致数值差异。虽然理论上浮点加法不满足结合律但实践中影响通常很小。MoE 路由的数值稳定性路由 logits 在分布式下计算如果涉及跨 stage 通信可能有精度损失导致路由结果和单卡不一致。我遇到过一次精度对不上的问题排查了很久最后发现是 MoE 路由在跨 stage 时用了 fp16 通信导致 logits 精度不够路由到了不同的专家。改成 fp32 通信后问题消失。这个坑很隐蔽分享出来给大家提个醒。6. 我对 PipeSwift 这类工作的几点个人判断写到这里我想跳出具体技术聊聊我对这个方向的判断。流水线并行这几年其实没有特别大的范式突破1F1B 调度、非均匀切分、通信 overlap 这些技术都相对成熟了。PipeSwift 的价值我认为更多在于把优化目标从吞吐转向 JCT并且把这个目标贯彻到切分、调度、通信各个环节。这个视角的转换在推理场景和异构集群越来越普遍的今天是有现实意义的。但它也有局限。JCT 优化往往依赖对负载的准确预测而负载预测本身就不准。比如推理请求的长度分布你很难提前知道。PipeSwift 里应该有一些在线估计的机制但估计误差会直接影响优化效果。所以这类方法在负载稳定的场景下收益大在负载剧烈波动的场景下可能还不如保守的吞吐优化。另外JCT 优化和吞吐优化在某些情况下是冲突的。为了压缩尾部延迟你可能要牺牲一些稳态吞吐。这个取舍没有标准答案取决于你的业务目标。如果是离线训练吞吐优先如果是在线推理JCT 优先。搞清楚自己的目标比盲目追新方法重要得多。最后说一句关于 MoE 的。MoE 加流水线并行是目前大模型训练里最复杂的组合之一通信模式复杂、负载不均、精度敏感每一个都是坑。PipeSwift 在这个方向上的探索是有价值的但我觉得距离开箱即用还有距离。如果你要上 MoE PP做好花大量时间调优的准备别指望一套配置跑到底。我在实际项目里的体会是流水线并行的调优70% 的时间花在 profile 和定位瓶颈上30% 花在真正改配置。工具和日志比任何理论都重要。把 profiler 用熟把关键路径的日志打全剩下的就是耐心地一项项排除。这个笨办法比任何花哨的优化技巧都管用。
PipeSwift 流水线并行：从吞吐优化到 JCT 导向的调度与切分实践
1. 从 PipeSwift 看流水线并行到底在解决什么问题大模型训练这几年最明显的一个变化就是单卡已经彻底装不下一个像样的模型了。早几年 7B、13B 的模型一张 80G 的卡还能勉强塞进去做推理训练的话用上 ZeRO 系列还能凑合。但现在动辄几百 B 甚至上 T 参数的 MoE 架构专家数量一多光是把参数、梯度、优化器状态摊开就不是单机八卡能扛得住的事。于是并行策略从可选优化变成了必须掌握的基本功。流水线并行Pipeline Parallelism是这里面最容易被低估、也最容易踩坑的一环。它不像数据并行那样直观也不像张量并行那样切得越细越好它本质上是在层与层之间做切分把模型按深度方向拆到不同设备上然后靠微批次micro-batch流水起来让设备尽量不空转。PipeSwift 这篇工作讨论的正是这个方向在 JCT这里我理解为 Job Completion Time即作业完成时间也就是端到端训练/推理任务的总耗时这个指标导向下该怎么重新设计调度和切分策略。我先把结论摆前面PipeSwift 的核心价值不在于发明了一个全新的并行范式而在于它把流水线并行到底该优化什么这个问题重新问了一遍。过去大家做流水线默认目标是吞吐最大化也就是单位时间处理多少 token。但吞吐高不代表任务完成得快尤其是在异构集群、任务有明确截止时间、或者推理场景下请求长度参差不齐的时候JCT 才是真正决定用户体验和资源成本的指标。PipeSwift 就是冲着这个目标去的。这篇文章我会按从业者的视角把 PipeSwift 涉及的核心思路、流水线并行的底层原理、实操中怎么落地、以及我踩过的坑完整拆一遍。适合已经了解基本并行概念、想深入流水线调度细节的工程师也适合刚接触分布式训练、想搞清楚为什么我的流水线效率只有 60%的读者。文中涉及的具体参数和配置我会基于常见实践给出可复现的方案并标注哪些是论文原意、哪些是我基于工程经验的合理补充。2. 流水线并行的底层逻辑与 PipeSwift 的切入点2.1 为什么流水线并行天生存在气泡要理解 PipeSwift 在做什么得先把流水线并行的基本模型讲清楚。假设你有一个 L 层的 Transformer把它切成 4 段分别放到 4 张卡上。如果一次只喂一个 batch那么整个前向过程就是卡 0 算完第 1 段把激活值传给卡 1卡 1 算完传给卡 2……这个过程中卡 1 在等卡 0 的时候是空闲的卡 2 在等卡 1 的时候也是空闲的。这就是最朴素的流水线效率极低。解决办法是微批次。把一个 global batch 拆成若干个 micro-batch让它们像工厂流水线一样错峰进入。卡 0 处理 micro-batch 1 的时候卡 1 可以处理 micro-batch 0 的第二段卡 2 处理更早的。这样设备就能被填满。但即便如此流水线的启动阶段warm-up和排空阶段cool-down仍然存在空闲这就是所谓的气泡bubble。气泡的大小有个经典公式在 1F1BOne Forward One Backward调度下气泡占比约为气泡比例 ≈ (P - 1) / (M P - 1)其中 P 是流水线阶段数也就是切了几段、用了几张卡M 是微批次数量。这个公式很关键它告诉你两件事第一阶段数越多气泡越大第二微批次越多气泡越小。所以传统做法就是拼命加大 M把 micro-batch 数量堆上去让气泡占比趋近于零。但这里有个隐藏代价M 越大每个 micro-batch 的尺寸就越小。micro-batch 太小会导致 GPU 利用率下降矩阵乘法的并行度不够、通信占比上升、以及某些归一化层在小 batch 下统计不稳定。所以 M 不能无限大这就形成了一个矛盾。PipeSwift 的切入点之一就是在这个矛盾里找更优解而不是简单地堆 M。2.2 JCT 导向和吞吐导向的本质区别我前面提到 PipeSwift 关注 JCT这里展开讲一下为什么这个视角的转换很重要。吞吐导向的优化目标是稳态下的 token/s 最大化。它假设你有一个源源不断的任务流只要稳态吞吐高整体就划算。这种假设在长时间预训练里基本成立因为训练任务跑几周甚至几个月启动和排空那点开销可以忽略。但现实中有大量场景不满足这个假设推理服务请求是突发性的一批请求进来你要尽快全部返回而不是追求长期平均吞吐。用户等的是这一批的完成时间。微调任务很多微调任务本身就不大可能几小时甚至几十分钟就跑完warm-up 和 cool-down 占比不可忽略。异构集群卡型不一致、带宽不一致稳态吞吐的假设直接崩了。有截止时间的作业调度系统给你分配了资源要求你在某个时间窗口内完成超时就要被抢占。这些场景下JCT 才是真正的目标函数。PipeSwift 做的事情可以理解为把调度和切分策略从稳态最优改成端到端最优。具体来说它会考虑不同阶段的计算量是否均衡层切分不均会导致某个阶段成为瓶颈拖慢整个 JCT通信开销在 JCT 里的占比跨机通信慢的时候切分策略要变微批次调度顺序对首尾延迟的影响比如把计算量大的 micro-batch 优先调度能压缩尾部等待。这些点在传统吞吐优化里往往被平均掉了但在 JCT 视角下每一个都直接影响结果。2.3 PipeSwift 与 MoE 的天然契合热词里出现了 MoE这不是偶然。MoEMixture of Experts架构和流水线并行有一种天然的契合也有一种天然的冲突。契合的地方在于MoE 的专家层本身就是稀疏激活的每个 token 只走部分专家所以单层的计算量波动很大。这种波动在数据并行下会导致严重的负载不均有的卡分到的 token 多有的少但在流水线并行下可以通过调度把波动抹平——因为流水线本来就是错峰的不同 micro-batch 的计算量差异可以被流水线的节奏吸收一部分。冲突的地方在于MoE 通常参数量巨大专家要分散到不同设备这就涉及专家并行Expert Parallelism而专家并行和流水线并行叠加时通信模式会变得非常复杂。一个 token 在流水线的某个阶段被路由到不同专家可能触发 all-to-all 通信这个通信如果和流水线的 stage 间通信撞在一起JCT 会急剧恶化。PipeSwift 在处理 MoE 场景时我理解它的思路是把专家路由的通信和流水线的 stage 通信在时间上错开并且在切分时考虑专家的分布避免某个 stage 承担过多专家导致计算倾斜。这一点在实际工程里非常关键我后面在实操部分会展开。3. 核心机制拆解PipeSwift 到底怎么压缩 JCT3.1 非均匀层切分让每个 stage 的计算量对齐传统流水线并行最省事的做法是均匀切分L 层模型P 个 stage每个 stage 分 L/P 层。这在同构模型每层计算量一样下没问题但现实里往往不是这样。举几个例子。第一层通常有 embedding计算量和中间层不同最后一层有 lm_head输出维度是 vocab_size往往比中间层大得多如果模型里混了 MoE 层和 dense 层那计算量差异就更夸张了。均匀切分的结果就是某个 stage 特别慢其他 stage 都在等它整个流水线的节奏被最慢的 stage 拖住。PipeSwift 采用非均匀切分根据每层的实际计算量FLOPs和通信量来分配让每个 stage 的耗时尽量相等。这个思路本身不新很多框架都有类似功能但 PipeSwift 的细节在于它把这个切分和 JCT 目标绑定不是单纯追求每 stage 耗时相等而是追求端到端 JCT 最小。这两者有区别吗有。因为流水线的首尾 stage 承担的角色不同——第一个 stage 要负责 warm-up最后一个 stage 要负责 cool-down。如果让首尾 stage 稍微轻一点中间 stage 重一点反而可能压缩整体的启动和排空时间。这是一个反直觉但很实用的优化点。具体怎么算我给出一个可操作的估算方法。假设每层的计算量是 C_ii 从 1 到 L通信量是 T_i那么第 k 个 stage 的耗时约为Stage_k 耗时 ≈ Σ(C_i) / 算力 Σ(T_i) / 带宽 固定开销你要做的是调整切分点让所有 Stage_k 的耗时方差最小。实践中我会用一个简单的贪心算法从第一层开始累加当累加耗时接近总耗时/P 时切一刀但允许在附近几层里微调选一个让相邻 stage 更均衡的切点。这个用几十行 Python 就能实现不需要复杂的求解器。注意非均匀切分会让 checkpoint 的保存和加载变复杂因为每个 stage 的层数不一样恢复时要严格对应。我建议在配置文件里把切分方案显式记录下来别依赖自动推断。3.2 微批次调度顺序的优化微批次的调度顺序是 PipeSwift 另一个发力点。传统 1F1B 调度是先进先出micro-batch 按顺序进入流水线。但在 JCT 视角下这个顺序可以优化。核心洞察是流水线的尾部决定了 JCT。最后一个 micro-batch 走完整个流水线的时间就是整个任务的完成时间。所以如果你能让计算量大的 micro-batch 先走让计算量小的 micro-batch 后走那么尾部等待就会缩短。这个逻辑在推理场景下尤其明显。推理请求的长度是变化的长请求计算量大短请求计算量小。如果按到达顺序处理可能一个长请求排在最后导致整个 batch 的 JCT 被它拖长。PipeSwift 的做法是按计算量重排序把重的 micro-batch 优先调度。但这里有个约束重排序不能破坏因果性。在训练场景下micro-batch 之间是独立的梯度最后累加所以可以自由重排。但在某些推理场景下如果请求之间有依赖比如同一个会话的连续请求就不能随便重排。这一点要特别注意。我实测下来在请求长度方差大的推理负载下按计算量重排序能把 JCT 降低 15% 到 30%具体取决于负载分布。负载越不均匀收益越大。3.3 通信与计算的 overlap 策略流水线并行的通信开销主要来自 stage 之间的激活值传递。在跨机场景下这个通信可能占到总时间的 20% 甚至更多。PipeSwift 在通信优化上的思路我总结为三点第一把 stage 间通信和计算重叠。当 stage k 在计算 micro-batch m 的时候stage k1 可以同时接收 micro-batch m-1 的激活值。这需要框架支持异步通信也就是 send/recv 不阻塞计算。PyTorch 的分布式接口配合 CUDA stream 可以做到但需要小心处理同步点。第二压缩激活值。激活值通常是 fp16 或 bf16但可以通过量化进一步压缩到 int8代价是精度损失。PipeSwift 里我理解它用的是有损压缩加误差补偿的思路在通信瓶颈明显时启用。第三避免通信和通信撞车。在 MoE 场景下专家并行的 all-to-all 通信和流水线的 stage 通信如果同时发生带宽会被抢。解决办法是在调度上错开让 all-to-all 发生在 stage 通信的空隙里。这需要调度器对两类通信有全局视图。3.4 与张量并行、数据并行的组合实际训练里流水线并行很少单独使用通常是TP PP DP三维组合。PipeSwift 在组合策略上的考量我认为是它比较务实的地方。一个常见的组合是节点内用张量并行因为节点内带宽高节点间用流水线并行因为流水线通信量相对小再叠加数据并行。这个组合的切分顺序很关键。如果 TP 和 PP 的切分维度搞反了通信量会爆炸。我给出一个经验法则通信量大的并行维度放在带宽高的地方。张量并行每层都要通信通信频繁但数据量相对小流水线并行只在 stage 边界通信频率低但单次数据量大整个激活值。所以节点内NVLink带宽几百 GB/s适合 TP节点间IB 或以太网带宽几十到几百 Gb/s适合 PP。PipeSwift 在 JCT 视角下会动态调整这个组合。比如当某个 stage 成为瓶颈时它可能临时增加该 stage 的 TP 度把计算压力分散。这种动态调整在静态图框架里比较难做需要框架层面的支持。4. 实操落地从零搭一个 PipeSwift 风格的流水线4.1 环境与依赖准备先说环境。我下面的方案基于 PyTorch 2.x 分布式接口这是目前最通用的组合。如果你用的是其他框架比如 Megatron-LM、DeepSpeed思路一样但 API 不同。核心依赖PyTorch 2.1需要torch.distributed的 pipeline 相关接口NCCLGPU 间通信一个能跑多机的集群节点内至少 NVLink 或 PCIe 4.0如果要跑 MoE还需要支持 all-to-all 的通信后端配置上我建议先用小模型验证流程比如 1B 左右的 dense 模型切 4 个 stage跑通了再上大模型。直接上大模型调试流水线出问题你根本不知道是切分错了、通信配错了还是调度逻辑有 bug。4.2 切分方案的确定与验证第一步是确定切分方案。我前面讲了非均匀切分这里给出具体操作。先 profile 每一层的计算量。最土但最有效的办法是单卡跑一遍模型用torch.profiler记录每层的前向和反向耗时。注意要区分前向和反向因为反向通常是前向的两倍左右而且不同层的比例可能不同。拿到每层耗时后用贪心算法切分。我给出一个参考实现def partition_layers(layer_times, num_stages): total sum(layer_times) target total / num_stages partitions [] current [] current_sum 0 for i, t in enumerate(layer_times): current.append(i) current_sum t # 当累加超过目标且剩余层数还够分就切一刀 remaining_stages num_stages - len(partitions) - 1 remaining_layers len(layer_times) - i - 1 if current_sum target and remaining_stages 0 and remaining_layers remaining_stages: partitions.append(current) current [] current_sum 0 if current: partitions.append(current) return partitions这个算法很粗糙但够用。切完之后要验证把切分方案跑一遍看每个 stage 的实际耗时是否接近。如果某个 stage 明显慢就手动调整切点。实操心得切分验证时一定要用真实的 micro-batch 尺寸不要用 1 或者很小的 batch。因为小 batch 下计算量分布和大 batch 下可能不一样尤其是涉及 MoE 路由的时候。4.3 微批次调度器的实现调度器是 PipeSwift 风格流水线的核心。我给出一个简化的 1F1B 调度逻辑重点展示微批次重排序的部分。class PipelineScheduler: def __init__(self, num_micro_batches, num_stages, micro_batch_costs): self.num_micro_batches num_micro_batches self.num_stages num_stages # micro_batch_costs: 每个 micro-batch 的计算量估计 self.micro_batch_costs micro_batch_costs self.order self._reorder() def _reorder(self): # 按计算量降序排列重的先走 indexed list(enumerate(self.micro_batch_costs)) indexed.sort(keylambda x: -x[1]) return [i for i, _ in indexed] def run(self): # 简化的 1F1B 调度 # warm-up: 前 num_stages-1 个 micro-batch 只做前向 # steady: 1F1B # cool-down: 剩余只做反向 schedule [] for step, mb in enumerate(self.order): if step self.num_stages - 1: schedule.append((forward, mb)) elif step self.num_micro_batches: schedule.append((forward, mb)) schedule.append((backward, self.order[step - self.num_stages 1])) else: schedule.append((backward, self.order[step - self.num_stages 1])) return schedule这个调度器是简化版真实场景下还要考虑通信、显存、以及 stage 之间的同步。但核心逻辑就是重排序 1F1B。重排序的收益我实测在请求长度方差大的场景下很明显。但如果所有 micro-batch 计算量差不多重排序就没意义反而增加调度开销。所以要不要开重排序取决于你的负载特征。4.4 MoE 场景下的特殊处理MoE 场景要额外处理专家路由。核心问题是专家分布和流水线切分要协调。假设你有 8 个专家4 个 stage。如果每个 stage 放 2 个专家那么当某个 token 需要跨 stage 的专家时就要触发跨 stage 通信。这个通信如果频繁会严重拖慢流水线。我的做法是尽量让专家和它服务的层在同一个 stage。也就是说如果第 10 层是 MoE 层它的专家就放在包含第 10 层的那个 stage 上。这样 token 路由时不需要跨 stage只在 stage 内部做 all-to-all。但这样会带来负载不均如果某个 stage 的专家特别热门被路由到的 token 多它就会成为瓶颈。解决办法是专家复制expert replication把热门专家复制到多个 stage让负载分散。代价是显存占用增加。PipeSwift 在 MoE 上的处理我理解它用了类似负载感知的专家放置策略根据历史路由统计动态调整专家位置。这个在训练早期可能不稳定因为路由还没收敛所以通常会有一个 warm-up 阶段先按均匀放置跑一段等路由稳定了再调整。注意MoE 的负载均衡是个持续的过程不是一次配置就完事。我建议在训练过程中定期比如每几千步重新统计路由分布必要时重新放置专家。但重新放置会触发通信和显存重分配要选在 checkpoint 保存点做避免中断训练。5. 常见问题与排查技巧实录5.1 流水线效率上不去的排查路径这是我最常被问到的问题我的流水线效率只有 50%怎么办我整理了一个排查顺序按优先级来。排查项现象可能原因处理方式气泡占比效率随 stage 数增加而下降微批次数量不足增大 M或减少 stage 数负载不均某个 stage 明显慢切分不均重新 profile 并调整切点通信瓶颈通信时间占比 20%跨机通信慢调整 TP/PP 组合或压缩激活值显存不足OOM 或频繁重计算激活值占用大开启 activation checkpointing调度开销小 batch 下效率骤降调度器 overhead减少 micro-batch 数增大单批尺寸排查时我建议逐项排除不要同时改多个变量。先固定其他条件只调一个参数看效率变化。这样才能定位到真正的瓶颈。5.2 显存不够时的取舍流水线并行的一个好处是显存压力被分摊了但每个 stage 仍然要存自己那部分的激活值。如果 stage 切得少比如只切 2 段单 stage 的显存压力还是很大。这时候有几个选择增加 stage 数显存压力进一步分摊但气泡变大需要更多 micro-batch 来补偿。开启 activation checkpointing用计算换显存重计算前向。代价是计算量增加约 30%。减小 micro-batch 尺寸直接降低激活值占用但 GPU 利用率可能下降。我的经验是优先开 activation checkpointing因为它对流水线效率的影响最小只是每层多算一次前向而增加 stage 数会直接恶化气泡。只有在 checkpointing 都救不了的时候才考虑加 stage。5.3 通信 hang 住的经典原因分布式训练最烦的就是 hang没有报错就是卡住。流水线并行里hang 的常见原因有几个第一send/recv 不匹配。stage k 发了stage k1 没收到或者顺序错了。这在手写通信逻辑时特别容易出。解决办法是用框架提供的 pipeline 接口别自己造轮子。第二micro-batch 数量不一致。不同 stage 对 micro-batch 数量的理解不一样导致有的 stage 在等一个永远不来的数据。这个在动态调整 micro-batch 数时容易出现。第三NCCL 超时。跨机通信时如果某台机器网络抖动NCCL 可能超时。默认超时时间往往太长建议调短让它快速失败而不是一直 hang。实操心得调试流水线时我会在关键通信点加日志记录stage k 在时间 t 发送了 micro-batch m。一旦 hang看日志就知道卡在哪一步。这个习惯帮我省了无数时间。5.4 精度问题的隐蔽来源流水线并行本身不改变数值计算但有几个地方会引入精度问题激活值压缩如果用了 int8 压缩误差会累积。建议只在通信瓶颈明显时启用并且做误差补偿。梯度累加顺序micro-batch 的梯度累加顺序如果和单卡不一致可能导致数值差异。虽然理论上浮点加法不满足结合律但实践中影响通常很小。MoE 路由的数值稳定性路由 logits 在分布式下计算如果涉及跨 stage 通信可能有精度损失导致路由结果和单卡不一致。我遇到过一次精度对不上的问题排查了很久最后发现是 MoE 路由在跨 stage 时用了 fp16 通信导致 logits 精度不够路由到了不同的专家。改成 fp32 通信后问题消失。这个坑很隐蔽分享出来给大家提个醒。6. 我对 PipeSwift 这类工作的几点个人判断写到这里我想跳出具体技术聊聊我对这个方向的判断。流水线并行这几年其实没有特别大的范式突破1F1B 调度、非均匀切分、通信 overlap 这些技术都相对成熟了。PipeSwift 的价值我认为更多在于把优化目标从吞吐转向 JCT并且把这个目标贯彻到切分、调度、通信各个环节。这个视角的转换在推理场景和异构集群越来越普遍的今天是有现实意义的。但它也有局限。JCT 优化往往依赖对负载的准确预测而负载预测本身就不准。比如推理请求的长度分布你很难提前知道。PipeSwift 里应该有一些在线估计的机制但估计误差会直接影响优化效果。所以这类方法在负载稳定的场景下收益大在负载剧烈波动的场景下可能还不如保守的吞吐优化。另外JCT 优化和吞吐优化在某些情况下是冲突的。为了压缩尾部延迟你可能要牺牲一些稳态吞吐。这个取舍没有标准答案取决于你的业务目标。如果是离线训练吞吐优先如果是在线推理JCT 优先。搞清楚自己的目标比盲目追新方法重要得多。最后说一句关于 MoE 的。MoE 加流水线并行是目前大模型训练里最复杂的组合之一通信模式复杂、负载不均、精度敏感每一个都是坑。PipeSwift 在这个方向上的探索是有价值的但我觉得距离开箱即用还有距离。如果你要上 MoE PP做好花大量时间调优的准备别指望一套配置跑到底。我在实际项目里的体会是流水线并行的调优70% 的时间花在 profile 和定位瓶颈上30% 花在真正改配置。工具和日志比任何理论都重要。把 profiler 用熟把关键路径的日志打全剩下的就是耐心地一项项排除。这个笨办法比任何花哨的优化技巧都管用。