Primus-Pipeline：更灵活、可扩展的流水线并行实现

Primus-Pipeline：更灵活、可扩展的流水线并行实现

原文作者：Cheng Yao, Lihuan Zhang, Wen Xie, Xiaoming Peng, Yuankai Chen, Anshu Raina, Chaojun Hou, Yao Fu, Liz Li, Zhenyu Gu.

本文介绍Primus Megatron-LM 后端中的一种灵活的流水线并行（Pipeline Parallelism, PP）实现，完整支持 zero-bubble 系列算法（zerobubble / zbv / v-half / v-min），以及传统的 1f1b 与 1f1b-interleaved 调度。我们还提供从理论到实操的模拟与性能分析工具，便于对各类流水线调度算法进行设计、验证与对比。

背景

流水线并行（PP）是大语言模型训练的高效策略。对于需要跨多节点的模型，PP 通过模型分片并以邻接点对点（p2p）激活传输为主（替代大规模全局通信），通常能提升扩展通信的可扩展性与可预测性。

业界常用的1f1b 与 1f1b-interleaved 调度已在 Megatron-LM、DeepSpeed 等框架广泛应用。近年来也出现了多种先进方法，例如 Sea AI Lab 的 zero-bubble 算法 [1] 与 DeepSeek 团队的 DualPipe/DualPipe-V [2]，但这些方法尚未融入主流开源 LLM 框架，主要原因是现有调度代码逻辑固定、难以修改。

以Megatron-LM 的 1f1b-interleaved 为例，其 PP 调度由三段主循环组成：warm-up（仅 forward）、steady（按序执行一个 forward 和一个 backward）、cool down（处理最后若干 mini-batch 的 backward）。通信操作在循环内部插入。如果要加入新的调度（如 zero-bubble 或 dual-pipe），通常需要复制并修改这三段循环的细节，代码会变得又长又复杂。

Primus-Pipeline

Primus-Pipeline 的核心思路是将流水线调度逻辑与训练执行解耦，便于验证与实现新的 PP 调度算法。主要贡献如下：

• 提供适用于PP 调度算法的灵活抽象，并完整支持多种先进算法，包括 1f1b-interleaved 与 zero-bubble 系列。

• 通过重新定义Primus-Turbo [5] 的算子，实现 GEMM 与 Grouped GEMM 的 Input Gradient / Weight Gradient 分离。

• 为每类PP 算法提供理论与实践层面的模拟工具，清晰评估指定配置下的 bubble rate（bubble rate）与显存占用。

调度设计

Primus-Pipeline 对 Megatron-LM 的 megatron.core.pipeline_parallel.get_forward_backward_func进行补丁替换。调度逻辑的入口位于PrimusPipelineParallelLauncher [3]。

在Primus-Pipeline 中定义并运行一个 PP 算法的步骤如下：

1)创建包含ScheduleNodes的ScheduleTable

多数PP 算法可以在给定 PP world size、每 rank 的虚拟流水线块（VPP）数、mini-batch 数量的条件下，构造包含调度节点的表。

• Schedule Nodes：每一步执行抽象为一个节点，包含计算类（FORWARD / BACKWARD / WGRAD（计算权重梯度））与通信类（RECEIVE_FORWARD / SEND_FORWARD 等）。

• PP Algorithms：pp-algorithms [4]

2)为调度节点绑定执行函数与参数

将调度表中的每个节点绑定到其处理函数与运行时参数。例如，Primus 在此处绑定 Megatron-LM 后端函数：Megatron-LM backend binding [6]

3)启动ScheduleRunner

节点绑定完成后，通过ScheduleRunner [7] 启动执行。

模拟器与性能

为验证与对比不同PP 算法的性能，我们提供不依赖硬件的投影工具，以及真实训练运行时的观测与分析工具。

模拟器

在给定ScheduleTable 的情况下，模拟器可在无需硬件的前提下，根据配置估算bubble rate（bubble rate）、显存占用与时序。这支持在大规模实验前快速对比不同调度与迭代开发。

下面的示例展示如何使用pp_simulation.yaml [8] 中的配置，验证不同 PP 算法。我们还做了大量基于 Transformer 层性能指标与多种并行策略的细粒度投影，详见 projection readme [9]。

执行以下命令运行模拟程序。工具会在控制台直接输出不同算法的性能指标。模拟器生成的数据可用下一节的可视化工具进行展示。

python3 primus/core/projection/performance_projection/simulator.py --config=primus/core/projection/configs/pp_simulation.yaml

下图为在8 个流水线阶段下，不同算法的 simulator.py 输出示例。

性能与可视化工具

我们提供两种模式的调度分析与可视化：基于模拟器输出进行离线分析，或基于真实训练运行的数据进行分析。

• Simulation perf（基于模拟数据）

运行simulator.py后，可以使用相同的YAML 配置绘制投影调度图。可视化脚本会生成各流水线阶段的时间轴图，便于直观对比不同算法的 bubble 模式与阶段利用率。使用以下命令生成图表：

python3 tools/visualization/pp_vis/vis.py --config=primus/core/projection/configs/pp_simulation.yaml

• Real training perf（基于真实训练数据）

在Primus 中开启 dump_pp_config: true，会生成各阶段的性能数据目录（默认：output/pp_data）。你也可以通过环境变量DUMP_PP_DIR 修改输出目录。随后，在vis.py [10] 的 draw_from_task_list() 中设置你的数据目录，运行以下脚本即可生成真实训练的流水线调度图：

python3 tools/visualization/pp_vis/vis.py

与Megatron-LM Backend 集成运行

使用Primus-Pipeline 的配置主要集中在 primus_pipeline.yaml。其中两个关键项为：

- patch_primus_pipeline：启用Primus-Pipeline，实现对 Megatron 原有调度逻辑的替换。

- pp_algorithm：选择具体的PP 调度算法。

此外，以下配置与Primus-Pipeline 存在冲突，需按如下方式设置：

overlap_grad_reduce: falseoverlap_param_gather: falseno_persist_layer_norm: truecreate_attention_mask_in_dataloader: falsegradient_accumulation_fusion: true

PP调度算法对比

我们的实现重点针对Sea AI Lab 提出的 zero-bubble 系列（zerobubble / zbv / v-half / v-min）。需要说明的是，post-validation 暂不支持，因为它需要对优化器进行较大改动，且跨 Megatron-LM 版本维护成本较高。

下表在假设forward / backward / weight-grad 三类操作的执行时间相等的情况下，对不同算法进行对比：

说明：V-half 与 V-min 采用贪心算法，因此没有闭式的 bubble-rate 公式。可使用模拟器估算其 bubble rate（bubble rate）。

记号：

• p：流水线阶段数

• m：mini-batch 数量

• x：常数项

实验

我们进行了两组实验以验证Primus-Pipeline 的性能：在 1 节点上使用 PP8 训练 Llama2-7B [12]；在 4 节点上使用 PP4 与 EP8（expert parallel）训练 Qwen3-235B [13] 的 MoE（Mixture of Experts）模型。

Llama2-7B 验证

• Setup：AMD GPU集群，1 节点，PP8，Llama2-7B 模型

下表对比了在8 个流水线阶段下，不同算法的吞吐与显存表现。作为致密（非 MoE）模型，Llama2-7B 可较直观展示各算法在吞吐（TFLOPS、tokens/s/device）与显存（max_memory、max_mem_percent）上的取舍。

下图展示了Llama2-7B 上不同算法的显存与吞吐对比：

Qwen3-235B 验证

• Setup：AMD GPU 集群，4 节点，PP4，EP8，Qwen3-235B 模型Setup：AMD GPU 集群，4 节点，PP4，EP8，Qwen3-235B 模型

该组实验展示了大规模MoE 训练的实际情况。在此配置下，zero-bubble 类算法并非总是吞吐最高的选择，但在降低显存占用方面具有优势。

下图为Qwen3-235B 的性能对比柱状图：

最佳实践指南

结合上述结果，我们得到以下结论：

• Llama 系列这类致密模型的吞吐与显存收益相对更高，因为整体网络中的 GEMM 占比更大、激活内存更高。

• 对于大规模MoE 实践场景，1f1b-interleaved 通常可获得更高的吞吐上限，但在降低显存方面受限；在显存受限时，v-half 是较为稳妥的选择。

基于此，我们建议在以下场景优先使用zero-bubble 系列算法（zero-bubble / zbv / v-half / v-min），而非 1f1b-interleaved：

• 模型中GEMM / Grouped GEMM 占比高：对于致密层较多的模型，拆分 weight-grad 与 input-grad 计算更有收益；两阶段时序不平衡会引入额外气泡。

• 显存是瓶颈：1f1b-interleaved 降低显存的手段有限；v-half 与 v-min 在显存紧张时更具实用性。

• 通信效率偏低：1f1b-interleaved 以额外通信换更低的气泡率，且更大的 VPP 会增加通信量。多数场景可通过与计算重叠来隐藏 p2p 通信，但在无 AINIC 或 RDMA 的条件下，zbv / v-half / v-min 的优势更明显。

• 极限分片受限：当模型无法将VPP 拆分到超过 2 的 rank 时，zbv / v-half / v-min 通常优于 1f1b-interleaved。

总结

本文介绍了Primus 中的 Primus-Pipeline，这是一套更灵活的流水线并行框架，帮助你更高效地研究与验证 PP 算法。以下是正在进行与计划探索的方向，欢迎社区贡献与反馈：

1) CPU offloading：基于调度节点的设计，可精细控制不同 mini-batch 与模型层的 offload/reload 时序。我们正在向 zbv / v-half / v-min 加入 offload 逻辑。

2) 更多算法：实现诸如 Dual-Pipe-V 等先进 PP 调度，并探索更高效的 PP 策略，欢迎贡献。

3) 细粒度重叠：如 DeepSeek-V3 报告所述，PP 调度将不同 mini-batch 的 forward 与 backward 交织，并重叠计算与通信。我们计划探索类似的细粒度重叠策略。

致谢

感谢SeaAI Lab [14] 团队及各位同仁的支持与合作，他们的专业与帮助对本项目的推进至关重要。

免责声明

本文档仅用于介绍和讨论特定技术方案及其工程实现思路，目的在于促进技术交流与研究参考，不构成任何形式的商业推广、性能承诺或产品保证。文中涉及的性能数据、实验结果及对比分析，均基于特定测试环境、配置条件及假设前提，不同场景下的实际效果可能存在差异，不构成任何明示或暗示的承诺。对其他开源项目、算法或实现方式的提及，仅作为背景说明，相关技术方案的适用性应结合具体业务需求、系统架构及运行环境进行独立评估，亦不构成采购决策、商业判断或投资决策的依据。

第三方内容由该第三方直接授权给你，AMD 不对其进行授权。所有链接的第三方内容均按“原样”提供，不作任何形式的保证。使用该等第三方内容的风险由你自行承担，AMD 在任何情况下均不对你使用第三方内容所造成的损失负责。你应自行承担使用第三方内容可能产生的所有风险与损害。

参考链接

1. zero-bubble algorithm：https://github.com/sail-sg/zero-bubble-pipeline-parallelism

2. DualPipe/DualPipe-V：https://github.com/deepseek-ai/DualPipe

3. PrimusPipelineParallelLauncher：https://github.com/AMD-AGI/Primus/blob/dev/yc/primus-pipe-blog/primus/backends/megatron/core/pipeline_parallel/primuspipe/pipeline_launcher.py

4. pp-algorithms：https://github.com/AMD-AGI/Primus/tree/main/primus/core/pipeline_parallel/scheduler/algorithms

5. Primus-Turbo：https://github.com/AMD-AGI/Primus-Turbo

6. Megatron-LM backend binding：https://github.com/AMD-AGI/Primus/blob/main/primus/backends/megatron/core/pipeline_parallel/primuspipe/pipeline_launcher.py#L239-L272

7. ScheduleRunner：https://github.com/AMD-AGI/Primus/blob/main/primus/core/pipeline_parallel/scheduler/scheduler.py

8. pp_simulation.yaml：https://github.com/AMD-AGI/Primus/blob/main/primus/core/projection/configs/pp_simulation.yaml

9. projection readme：https://github.com/AMD-AGI/Primus/blob/main/docs/projection.md

10. vis.py：https://github.com/AMD-AGI/Primus/blob/main/tools/visualization/pp_vis/vis.py

11. Llama-2-7B：https://modelscope.cn/models/Charlesnii/meta-llama-Llama-2-7b

12. Qwen3-235B：https://modelscope.cn/models/Qwen/Qwen3-235B-A22B

13. SeaAI Lab：https://sail.sea.com/

14. Megatron-LM：https://github.com/NVIDIA/Megatron-LM