FlexPipe：把流水线调度写进十行代码，让大模型训练最高加速 2.28 倍

作者：JLJ, QXJ, MZX, LHJ from DeepLink Group @Shanghai AI Lab

流水线并行（Pipeline Parallelism）是大模型分布式训练的基石之一。从 GPipe、1F1B 到 Interleaved 1F1B、Chimera、Zero Bubble，每一代预定义高效调度（schedule）都是系统专家手工雕琢的结晶。

但模型架构一直在变：多语言模型带来了百万级词表的巨型 Embedding 层，多模态模型带来了多个需要同步的异构子模块，MoE 带来了新的通信模式……面对这些新架构，预定义的调度不再是最优解，甚至会严重拖慢训练。而无论是手工设计新调度（调度空间大到天文数字），还是手工实现新调度（Megatron-LM 中 1F1B 约 0.7k 行代码、Interleaved 1F1B 约 1.4k 行），都足以劝退绝大多数人。

能否不再手写一个个调度，而是让框架自己去"发现"和"生成"高效调度，同时将定制调度的成本极致压缩？

面向这一痛点，上海AI实验室DeepLink团队提出一种灵活可编程的流水线并行框架FlexPipe，将流水线调度写进十行代码，且能让大模型训练最高加速 2.28 倍。

先看核心思路与实测数据：

FlexPipe 包含一门简洁的 DSL、一个自动调度器、一个 Auto-tuner 和运行时。

它的核心洞察是：只要确定了 stage 放置方式，一个流水线调度就可以被拆解为一连串"调度步"，而调度步的行为完全由两个"遍历优先级"决定——先算 forward 还是先算 backward（计算类型遍历优先级），以及多个 stage 在同一设备上时先算谁（stage 遍历优先级）。 把这两个优先级参数化，主流调度（1F1B、Interleaved 1F1B、Chimera、Hanayo、ZB-H1 等）都能用几行代码表达；枚举这两个优先级，就能在秒级到分钟级的开销内自动搜索出新调度。

实验表明：在大词表 GPT 模型上，FlexPipe 相比 Megatron-LM 最高加速 2.28×，相比 SOTA 自动调度搜索框架 Tessel 最高加速 1.49×；在多模态 CLIP 模型上相比 DistMM-Pipe 最高加速 1.29×。而当 Tessel 在 8 卡规模就搜索超时失败时，FlexPipe 在 32 卡规模的搜索开销仍不到 9 分钟。

论文地址：https://arxiv.org/abs/2510.05112

背景：影响流水线并行效率的"空间"与"时间"因素

流水线并行把模型按层切分成多个 stage，放到不同设备上；再把一个 mini-batch 切成多个 micro-batch，让它们像流水线一样依次流过各个 stage，使所有设备同时工作。它只依赖轻量的点对点通信传递中间结果，因此是跨服务器扩展大模型训练的必备手段。

决定流水线并行效率的因素有两个维度：

• 空间维度——stage 放置（stage placement）：哪些 stage 放在哪个设备上。如下图所示，1F1B 采用一对一（one-to-one）放置；Interleaved 1F1B 采用环形（circular）放置，每个设备持有多个不相邻的 stage；Hanayo 采用 V 形（V-shape）放置；Chimera 采用双向（bidirectional）放置，两条方向相反的流水线共享设备。

• 时间维度——调度（schedule）：各个 micro-batch 的 forward / backward 计算以什么顺序在流水线上执行。

图1：主流的 stage 放置策略（a0-a4 表示设备，s1-s8 表示 stage）

这两个维度共同决定了流水线并行的两大性能指标：气泡（bubble，即设备因等待依赖而空闲的时间）和显存消耗。过去几年的经典方案，本质上都是在这两个维度上做文章：1F1B 用"早做 backward"压低激活显存；Interleaved 1F1B 用更细的 stage 切分缩小气泡；ZB-H1 把 backward 拆成输入梯度和权重梯度两部分、用权重梯度计算填充气泡；BitPipe 则融合了交错流水线与双向流水线。

为什么需要FlexPipe这类可编程流水线调度？

接下来我们先通过两个典型的场景，具象地分析现有预定义调度方案为什么不够用了；并进一步拆解手工实现调度要翻越的三座大山，以及现有自动化方案的能力局限。

案例一：大词表带来的负载失衡

多语言模型为了覆盖多种语言的词汇，Embedding 层的词表规模不断膨胀（Gemma 等模型的词表已达 25 万以上）。对一个 5B 参数的 GPT 模型做 profiling 会发现：随着词表增大，含 Embedding 层的 stage 与其他 stage 的负载差距急剧拉大——词表达到 1M 时，最慢的 stage 比最快的慢 5.63×，显存多耗 4.87×。

图2：5B GPT 模型在不同词表规模下的计算与显存失衡

一个自然的解法是设计新的 stage 放置策略（上文 图 1 中的"MLLM"放置），把 Embedding 层沿词表维度切分到流水线组内的多个设备上。但问题随之而来：新的放置策略套用预定义调度后，会因为新增的数据依赖（Embedding 的 AllReduce 等）产生额外气泡。

此时需要为新放置策略重新探索调度——而调度空间是天文数字：数百上千个 micro-batch、错综复杂的依赖关系、允许多少个 micro-batch 同时在流水线中……更麻烦的是，调度的优劣还与模型输入有关：batch 较小时 Interleaved 1F1B 优于 1F1B，batch 较大时气泡的影响变小、通信更少的 1F1B 反而更优。没有一个调度能包打天下。

案例二：多模态模型的同步开销

多模态模型（如 CLIP 类模型）包含处理不同模态的多个子模块。DistMM-Pipe 把不同子模块放到不同设备组上并行计算，但子模块之间需要同步交换输出状态——这是一个现有调度里根本不存在的新操作。DistMM-Pipe 的做法是在 warm-up 阶段多发射一些 forward，每隔几个 micro-batch 插入一次同步。

图3：DistMM-Pipe 的调度

进一步分析发现，同步开销约占 17%。其实还有很大的优化空间：可以用异步通信消除同步气泡；不同子模块的规模差异很大，完全可以对不同子模块采用不同类型的调度（混合调度）。但现有框架既没有高效手段去灵活调整调度，也不支持往调度里添加新操作。

手工实现调度的三座大山

就算设计出了新调度，把它写出来也是苦役：

1. 生产力（Productivity）：在 Megatron-LM 中，1F1B 的实现约 0.7k 行，Interleaved 1F1B 约 1.4k 行。每个新调度都要从头写起。

1. 可编程性（Programmability）：开发者要在分布式环境中维护不同 micro-batch、不同计算类型甚至不同 stage 之间纠缠的通信与数据依赖；新模型架构还会不断引入新操作（比如多模态的子模块同步）。

1. 可调试性（Debuggability）：手工实现极易出错（死锁、错序、精度问题），且缺乏中间调试信息，排错过程令人痛苦。

现有自动化方案的不足

• Tessel 支持在给定 stage 放置策略下自动搜索调度，但可搜索的调度类型有限——例如不支持环形（circular）放置，而 Interleaved 1F1B、BFS Pipeline、Hanayo 等一大批高效调度恰恰构建在这类放置之上；更致命的是搜索成本：基于 Z3 求解器的搜索在设备数达到 8 时就会失败（后文实验有具体数据）。

• Koala 等编译器方案提出了 DSL 来提升开发效率，但其语法是"计算图 DSL + 可执行 DSL"的两层结构，实现 1F1B 仍需约 0.2k 行代码；而且编译器方案既不能自动探索新调度，也不支持为新模型架构注册新操作。

一句话总结：我们需要一个框架，既能自动搜索多种类型的高效调度（且搜得快），又能让用户用极少的代码定制调度、注册新操作。

FlexPipe 四大组件

FlexPipe 由四个组件构成：DSL、自动调度器（Scheduler）、Auto-tuner 和运行时（Runtime）。

图4：FlexPipe 总体架构

• DSL 提供简洁的 API 表达各种调度：用户只需为不同调度类型设定几个参数；要往调度里加新操作，只需写一个函数。其余的复杂调度工作全部交给框架。

• 自动调度器 与 DSL 交互，自动编排调度。它内部把设备抽象为 actor，构建"计算调度空间表示"（CSSR，灵感来自 CPU 指令流水线），再通过 actor 感知的调度机制自动生成以指令序列表示的调度，并附带梯度分离、异步通信等优化 pass。

• Auto-tuner 面向不同的模型输入和硬件资源搜索、调优调度：profiling 组件采集计算与通信耗时（结果落盘复用，避免重复采集），离线 tuner 在调度空间中枚举并选出最优配置。

• Runtime 将调度实例化为高效的运行时执行。

核心洞察：调度 = 两个遍历优先级

FlexPipe 之所以能同时做到"可编程"和"可搜索"，源于对调度过程的三个关键观察。

观察一：调度过程可以拆解为一连串"调度步"。 把一次 micro-batch 计算（一次 forward 或一次 backward）的调度视为一个调度步。一旦 stage 放置策略确定，每个设备要执行哪些计算就确定了——剩下的事情就是：每个设备不断从"依赖已就绪"的计算中挑一个执行，直到全部完成。

观察二：决定调度形态的，是"挑哪个"的优先级。 当一个设备面前摆着多个依赖就绪的计算时，挑选规则只涉及两个维度：

观察三：高效调度普遍呈现重复周期性。 观察 1F1B、Interleaved 1F1B、Chimera 等高效调度会发现，它们都在每个设备上周期性地重复同样的计算模式——等价于说，不同调度步使用相同的遍历优先级。这个性质大幅收缩了需要探索的调度空间：搜索一个调度，不再是搜索一张巨大的时序表，而是枚举几个优先级参数的组合。

图5：(i) 各 stage 放置策略的优劣特性；(ii) 不同优先级下指令的遍历顺序示意

举例来说：1F1B = 一对一放置 + bwdfirst；Interleaved 1F1B = 环形放置 + bwdfirst + 带 interval 的 breadth-first；BFS Pipeline = 环形放置 + fwdfirst……主流调度都落在这个参数化空间里。反过来，枚举这个空间，就可能"搜出"前人没画过的新调度。

FlexPipe DSL：十行代码写出一个调度

FlexPipe DSL 内嵌于 Python，由描述 stage 的数据模型和组合调度的调度语言两部分构成。

数据模型与指令集

数据模型的核心构件是 Stage / StageSet（stage 数据及其 DAG）与 Actor / ActorMesh（设备数据及拓扑），通过 ModelConfig 统一操作。用户先定义设备拓扑和模型配置，再用 partition 和 place 把模型层切分成 stage 并映射到 actor 上。框架内置了一对一、环形、V 形、双向四种放置策略（正是因为没有一种放置在所有指标上全胜，见上文 图 5(i)），也允许用户重写 partition 自定义切分规则。

FlexPipe 把调度表示为指令序列，每条指令带有 (stage ID, micro-batch ID) 标签，运行时映射到对应操作：

下面是用 FlexPipe DSL 实现 1F1B 的完整代码——11 行（对比 Megatron-LM 的约 700 行）：

图6：1F1B 的 FlexPipe DSL 实现（前 3 行定义设备拓扑与模型配置，第 4-5 行切分并放置 stage，第 6-8 行设置两个遍历优先级，第 9-11 行构建 stage DAG、CSSR 并生成调度）

注册新指令：以多模态同步为例

面对新模型架构，FlexPipe 允许用户注册框架尚未内置的新操作。以多模态模型的子模块同步为例：

图7：为多模态模型注册同步指令

流程分三步：

1. 注册指令：register_new_inst 注册新指令并返回指令类型。通过 sched_attr 配置调度规则——check_prev_dep 指定该指令调度前必须完成哪些前置依赖指令，check_nxt_dep 指定何时可以调度后继依赖指令，sched_unit 决定一次调度多少条指令；通过 inst_attr 配置运行时属性（如同步通信组的 ranks）。调度器会在调度过程中自动调用这些规则解析数据依赖。
2. 挂接依赖：指令的依赖一般与 stage 相关（比如同步必须发生在各子模块最后一个 stage 的 forward 之后、backward 之前），所以先用 register_new_stage 注册新 stage 挂载新指令，再用 set_cssr_deps 传入依赖对集合 ((指令类型1, stage_i), (指令类型2, stage_j)) 声明先后关系。
3. 绑定操作：map_inst_to_operation 把用户定义的可调用操作映射到注册的指令，运行时执行。

更细粒度的控制

除了两个遍历优先级，DSL 还提供三个"手动挡"机制：

图8：控制 micro-batch 调度顺序的 DSL 示例

• config_inflight_micros：控制各 stage 允许的最大在途（in-flight）micro-batch 数量，直接影响峰值显存与调度形态，也可传入自定义函数设定规则；

• register_new_checkfunc：注册调度过程中的检查规则，比如自定义某类指令的依赖判定；

• register_new_priority：当内置优先级不够用时，注册全新的优先级（写一个 step 函数定义新的遍历顺序）。

用户按需取用：写惯例调度只需设参数，做深度定制才需要写函数。复杂度是渐进暴露的。

自动调度器：像 CPU 指令流水线一样调度

FlexPipe 调度器的设计灵感来自经典的 CPU 指令流水线：micro-batch 计算是"指令"，设备是执行单元，依赖就绪的指令进入"就绪队列"等待发射。整个调度过程分三个阶段：切分放置 → 构建 CSSR → actor 感知调度。

CSSR：计算调度空间表示

CSSR（Computation Schedule Space Representation）负责在每个调度步动态给出每个 actor 当前"依赖已就绪"的候选计算。它由五个部分组成：

• 指令项（Instruction Item）：调度的基本单元，即一条指令实例，标签为 (指令类型, micro-batch ID, stage ID)；

• 指令池（Instruction Pool）：容纳所有待调度的指令项，按指令类型组织成多个指令队列（默认创建 FwdPass、BwdPass 队列，用户注册的指令类型也会建队列）；

• Actor 感知的重排队列（Reorder Queues）：每个 actor 为每种指令类型维护一个重排队列，存放依赖已解决、可以调度的指令项——相当于乱序执行中的就绪队列；

• 本地缓冲（Local Buffer）：每个 actor 记录当前调度步的调度状态，调度步结束后统一提交并清空——用来模拟分布式环境中"同一时刻看不到远端状态"的现实约束，保证生成的调度在真实分布式环境中可执行；

• 数据依赖与动态依赖解析器：forward 指令间的依赖由 stage DAG 转换而来，backward 依赖是其反向，再加上用户通过 set_cssr_deps 注入的依赖；每当一条指令项被调度，解析器立即检查其后继是否就绪，将就绪者放入对应 actor 的重排队列。

Actor 感知调度

配置好优先级后，调度按五步循环执行：

1. 初始化：只有第一个 stage 的 FwdPass 指令项无依赖，进入对应 actor 的重排队列；
2. 取指令：每个调度步中，所有 actor 执行 StepFunction 取出可调度指令项——先按计算类型遍历优先级决定先查哪个重排队列，再按 stage 遍历优先级（方向 + interval）决定队列内取哪一项；若无可调度项，就产生一个气泡；
3. 记状态：各 actor 把调度状态写入本地缓冲；
4. 解依赖：动态依赖解析器把新就绪的指令项放入相应重排队列；
5. 提交：把本地缓冲的状态更新到指令池并清空缓冲；指令池全部调度完则结束，否则回到第 2 步。

值得一提的是，warm-up、稳态、cool-down 三个阶段用完全相同的方式取指令，无需特殊处理——周期性是优先级机制的自然结果，而不是人为拼接的。

图9：自动调度过程示例：两个模态采用混合调度——modality1 用环形放置 + interleaved + 带 interval 的遍历，modality2 用一对一放置 + bwdpass-first

气泡优化：梯度分离与异步通信

调度生成后，FlexPipe 还会跑优化 pass。梯度分离算法自动检测气泡并"暂存-回填"：检测到气泡时，先定位被阻塞的指令项（如某个 BwdPass），回溯已调度的 BwdPass，把其中的权重梯度计算（CompWeightGrad）暂存起来，让关键依赖链上的计算尽早执行，从而释放被阻塞的指令；暂存的权重梯度计算再被插入后续气泡中。这相当于把 Zero Bubble 的核心思想泛化成了一个对任意调度类型都适用的自动优化。

图10：梯度分离算法示例

此外，FlexPipe 还采用异步 P2P 通信进一步隐藏通信开销。

Auto-Tuner：为不同输入自动选调度

如前所述，没有一个调度能在所有模型输入和硬件资源下都最优。FlexPipe 用网格搜索在如下参数空间中寻找最优配置：

(pp, dp, mbs, 放置策略, fstp, bstp, ctp)

其中后四项（放置策略 + forward/backward 的 stage 遍历优先级 + 计算类型遍历优先级）定义了调度类型，取值来自内置值表；并配有剪枝规则：interval 仅在环形放置时启用、ctp 默认不取 fwdfirst、调度单元默认为 1（配合细粒度梯度计算减少气泡）。得益于"调度 = 优先级组合"这一参数化，整个空间小而稠密——里面几乎全是结构合理的调度，这是 FlexPipe 搜索开销远低于逐时隙求解方案的根本原因。

实验结果

实验平台：8 台节点共 64 张 NVIDIA A800 SXM4 80GB。

搜索开销：Tessel 8 卡就搜不动了

与基于 Z3 求解器的 Tessel 及其加速版 Tessel-fast 对比调度搜索开销（✗ 表示搜索失败）：

Tessel 的搜索空间随设备数爆炸式增长，8 卡即失败；Tessel-fast 撑到 16 卡也失败了。FlexPipe 的搜索开销随设备数平缓增长，32 卡仍在 9 分钟内完成，并且能为不同模型配置搜出多个高效候选。当流水线并行需要与其他并行方式联合调优时，秒级到分钟级的搜索成本是"能不能用"的分水岭。

图11：Tessel 与 FlexPipe 搜出的调度对比：Tessel 为避免稳态气泡在 warm-up 阶段发射更多 micro-batch，反而在 warm-up 和 cool-down 阶段引入更多气泡；FlexPipe 搜出的调度整体气泡更少

端到端性能：大词表 GPT

针对第 2 节的大词表场景（MLLM 放置策略 + 词表维度张量并行），对比 Megatron-LM（1F1B）、T-1F1B（1F1B 适配 MLLM 放置）和 Tessel。

5B GPT，词表 64K→1M，32 卡：FlexPipe 相比三者分别加速 1.14×–1.91×、1.20×–1.28×、1.13×–1.30×。词表越大，Megatron-LM 的失衡越严重、性能下降越多；T-1F1B 因 Embedding 张量并行的 AllReduce 依赖产生气泡；Tessel 则在 warm-up/cool-down 阶段付出额外气泡。

图12：5B GPT 模型端到端性能

16.1B GPT，词表 256K→1.2M，32 卡：FlexPipe 分别加速 1.31×–2.28×、1.27×–1.31×、1.24×–1.49×。值得注意的是，Megatron-LM 在词表达到 1M 时直接 OOM（即使开启重计算）；Tessel 由于 16 卡以上无法完成搜索、只能在 pp<16 的配置下运行，词表 1.2M 时同样 OOM。

图13：16.1B GPT 模型端到端性能（✗ 表示开启重计算后仍 OOM）

性能分解与扩展性

把每次迭代的运行时间拆解为 stage 计算时间和设备等待时间（反映气泡）：各调度的计算时间相差无几（差异主要来自 micro-batch 大小带来的计算效率区别），真正拉开差距的是等待时间——Megatron-LM 等待最久，T-1F1B 和 Tessel 次之，FlexPipe 最少，与端到端结果一致。

图14：1M 词表下 5B / 16.1B 模型的运行时间分解

弱扩展实验中（16→64 卡，dp 随卡数等比扩展），FlexPipe 在不同硬件资源规模下始终保持性能优势。

图15：5B GPT 模型弱扩展实验

多模态模型优化

回到第 2 节的第二个案例。基于 FlexPipe 的可编程机制，我们对 DistMM-Pipe 做了两处优化，全部通过 DSL 完成、无需触碰框架内核：

• 异步同步通信：把子模块间的同步注册为新指令 SyncWithGather（即第 5 节的注册示例），跟随两个模态最后一个流水线 stage 的每次 forward 异步执行，消除集中同步的等待；

• 混合调度：audio/image 模态子模块用 Interleaved 1F1B，text 模态子模块用 1F1B——不同子模块各取所需。

在 40 卡上训练 4.7B / 8.2B / 9.7B 的 CLIP 模型，相比 DistMM-Pipe 分别加速 1.26× / 1.29× / 1.18×。

图16：CLIP 模型端到端性能

总结

FlexPipe 的关键并不在于发明了某一个新调度，而在于换了一个刻画调度的坐标系：流水线调度的本质不是一张张手绘的时序图，而是"stage 放置 + 两个遍历优先级"这寥寥几个自由度。把自由度参数化之后，两件原本很难的事情同时变简单了——

• 搜索：调度空间从"逐时隙求解"（Tessel 8 卡即失败）坍缩为一个小而稠密的参数空间，32 卡搜索只需 9 分钟；

• 编程：表达一个调度从上千行代码（Megatron-LM）降到十行左右，注册一个新操作只需写一个函数，依赖、通信、warm-up/cool-down 的拼接全部由框架接管。

在大词表 GPT 上相比 Megatron-LM 最高 2.28× 加速、在多模态 CLIP 上相比 DistMM-Pipe 最高 1.29× 加速，都是这一抽象带来的自然收益。未来我们期待在 FlexPipe 中集成更多优化手段（重计算、显存卸载、更多通信优化等），覆盖更多训练场景。

论文地址：https://arxiv.org/abs/2510.05112

如果你喜欢我们的内容，请点赞❤️、收藏⭐️、关注➕我们！

也欢迎在评论区与我们互动！

你的支持是我们持续创作的动力！

参考文献

[1] Lijuan Jiang, Xingjian Qian, Zhenxiang Ma, Zan Zong, Hengjie Li, Chao Yang, Jidong Zhai. "A Flexible Programmable Pipeline Parallelism Framework for Efficient DNN Training." arXiv preprint arXiv:2510.05112 (2025).

[2] Huang et al. "GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism." NeurIPS 2019.

[3] Narayanan et al. "PipeDream: Generalized Pipeline Parallelism for DNN Training." SOSP 2019.（1F1B）

[4] Narayanan et al. "Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM." SC 2021.（Interleaved 1F1B）

[5] Li & Hoefler. "Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines." SC 2021.

[6] Liu et al. "Hanayo: Harnessing Wave-like Pipeline Parallelism for Enhanced Large Model Training Efficiency." SC 2023.

[7] Qi et al. "Zero Bubble Pipeline Parallelism." ICLR 2024.（ZB-H1）

[8] Lamy-Poirier. "Breadth-First Pipeline Parallelism." MLSys 2023.

[9] Wu et al. "BitPipe: Bidirectional Interleaved Pipeline Parallelism for Accelerating Large Models Training." 2024.

[10] Lin et al. "Tessel: Boosting Distributed Execution of Large DNN Models via Flexible Schedule Search." HPCA 2024.

[11] Huang et al. "DISTMM: Accelerating Distributed Multimodal Model Training." NSDI 2024.（DistMM-Pipe）

[12] Tang et al. "Koala: Efficient Pipeline Training through Automated Schedule Searching on Domain-Specific Language." 2025.