🤖 AI

Piper: A Programmable Distributed Training System

Piper 是一个可编程分布式训练系统，它通过统一的中间表示将高层并行策略与底层运行时实现解耦，在实现 ZeRO 和 DeepSeek-V3 的 DualPipe 等多样化策略灵活组合的同时，保持了性能对等并提升了效率。

原作者： Megan Frisella, Shubham Tiwari, Andy Ruan, Yi Pan, Parker Gustafson, Mat Jacob, Gilbert Bernstein, Stephanie Wang

发布于 2026-06-10

📖 1 分钟阅读☕ 轻松阅读

原作者： Megan Frisella, Shubham Tiwari, Andy Ruan, Yi Pan, Parker Gustafson, Mat Jacob, Gilbert Bernstein, Stephanie Wang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图烘焙一个巨大且复杂的蛋糕（一个巨型 AI 模型），它大到无法装进单个厨房。你必须雇佣数百名厨师（计算机芯片），并将工作分配给他们。

问题在于，协调这数百名厨师是非常困难的。如果由于你告诉厨师 A 去搅拌面糊，而让厨师 B 等待鸡蛋，就会浪费时间。如果厨师 A 和厨师 B 同时试图抢夺同一个碗，他们就会发生碰撞。

目前的 AI 训练系统大多像僵化的食谱。它们只有几种预先写好的分工方式（比如“每个人搅拌自己的面糊”或者“把碗沿着队伍传下去”）。如果你想尝试一种新的、更高级的烹饪方法来节省时间或空间，你必须聘请一名专家从头开始重写整个食谱。这既缓慢又昂贵，而且难以更改。

Piper 是一个全新的“智能厨房经理”，它改变了游戏规则。Piper 不再提供僵化的食谱，而是给你一套简单的指令（比如“把搅拌过程放在左边的柜台上”或“在烤箱加热时传递碗”）和一个通用的蓝图。随后，无论你的指令多么复杂，Piper 都能算出如何高效地完成任务。

以下是 Piper 的工作原理，通过简单的概念进行拆解：

1. 蓝图（即“IR”）

想象一下，蛋糕的食谱不仅仅是一份步骤清单，而是一个巨大的 3D 流向图，展示了每位厨师需要执行的每一个动作，以及他们彼此之间的依赖关系。Piper 将其称为 Global Training DAG（一个关于所有工作的地图，这是一个高级术语）。

旧方式： 系统只知道几种固定的流向图。
Piper 方式： 你可以使用简单的标签绘制自己的流向图。Piper 会将你的绘图转化为一份所有厨师都可以遵循的总计划。

2. 指令（即“API”）

Piper 提供了一种简单的语言，让你无需了解硬件的底层细节即可告诉系统你的需求。

注解（Annotations）： 你可以为模型的不同部分打上标签。例如，你可以说：“这部分蛋糕仅供‘专家’厨师使用”，或者“这部分需要拆分到 4 个不同的柜台上”。
指令（Directives）： 这些是你的命令。你可以说：“将这大批量的面团分成更小的份量，以便我们能同时进行处理”，或者“确保在烘焙冷却的同时进行搅拌”。
神奇之处： 你不需要告诉计算机如何移动碗，也不需要告诉它使用哪把特定的勺子。你只需要说明需要发生什么，Piper 就会搞定物流细节。

3. 经理（即“Runtime”）

一旦你向 Piper 提供了指令，它就会充当中央交通控制器。

中央调度器（Central Scheduler）： 它查看你的蓝图，并将其分解为每个厨师的具体微小任务。它决定谁在何时做什么。
工人（Workers）： 每位厨师（计算机芯片）都会得到一份属于自己的任务清单。他们不需要互相交流来制定计划；他们只需遵循 Piper 给出的清单即可。
避免交通拥堵： AI 训练中的一个主要问题是“气泡（bubbles）”——即厨师因为等待一个碗而闲置的时间。Piper 非常聪明，擅长让任务重叠。它可能会告诉厨师 A 在厨师 B 还在清洗上一个碗的时候，就开始搅拌下一批面糊。它确保通信（传递碗）与烹饪（搅拌）在同一时间发生，因此没有人会无故等待。

为什么这很重要？

论文表明，Piper 可以实现两件事，而目前的系统在这些方面表现挣扎：

它能达到专家的速度： 对于标准的、常见的训练方式，Piper 的速度与目前顶尖公司使用的顶级系统不相上下。
它解锁了新的、复杂的策略： 因为 Piper 具有灵活性，它可以处理其他系统无法处理的“混合”策略。
- 例子： 论文提到了一个叫做 DualPipe 的策略。想象一个厨房，你正在同时烘焙两种不同类型的蛋糕，但你必须以一种非常特定、复杂的节奏在烤箱之间交换食材。旧系统会被这种节奏搞混，要么崩溃，要么运行得非常缓慢。Piper 能完美处理这种节奏，让厨师的工作效率提高 6–30%，并且能在不耗尽柜台空间（内存）的情况下，装下更大的蛋糕。

核心结论

Piper 就像是从一份僵化的、预印好的说明书升级到了一个智能、灵活的项目经理。它让 AI 研究人员能够设计出自定义的、训练巨型模型的方法，而无需成为硬件工程师。它让系统既能满足当下的速度需求，又具备应对未来那些疯狂、复杂策略的灵活性。

技术摘要：Piper —— 一种可编程分布式训练系统

问题陈述

随着机器学习模型向需要数百或数千个加速器的基础模型规模演进，设计分布式训练策略的复杂性已成为一个瓶颈。现有方法面临着一种两极分化的局面：

部署系统： 高性能系统（例如用于 DeepSeek-V3 的系统）依赖人类专家来手动协同设计高层并行策略（如流水线并行 (PP)、专家并行 (EP)）与底层执行策略（如 GPU 内资源调度、流管理）。这导致系统过于僵化，难以适配新的策略或组合并行维度。
通用框架： Megatron-LM、DeepSpeed 和 TorchTitan 等框架提供了灵活的接口，但通常将并行维度视为独立的。它们采用贪婪调度（eager dispatch）模式，这使得在组合策略中（例如重叠 EP all-to-alls 与 DP all-reduces）进行通信与计算的联合调度变得非常困难。此外，由于层钩子（layer hooks）与执行调度之间的相互作用，它们往往无法完全支持复杂的组合，例如 PP 与 ZeRO-2 或 ZeRO-3 的结合。

核心挑战在于可扩展性：即创建一个能够最小化指定和实现任意分布式训练策略所需工作量的系统，同时保持高性能和高显存效率。

方法论

Piper 通过一个统一的中间表示（IR）将策略规范与运行时实现解耦，从而解决了上述挑战。该系统由三个主要部分组成：

1. 用户调度 API

Piper 提供了一个声明式 API，允许用户在无需编写自定义调度器的情况下，控制高层放置（placement）和底层执行。

注解 (Annotations)： 用户使用上下文管理器为 PyTorch 模型中有语义意义的区域（例如特定的层、专家块）打上标签。这些标签定义了并行维度的边界（例如 PP、EP、DP）。
调度指令 (Scheduling Directives)： 用户应用指令来转换系统的 IR。关键指令包括：
- 放置 (Placement)： Place、Replicate 和 Shard 定义了设备映射，并引入必要的通信节点（P2P、All-Reduce、All-Gather、All-to-All）。
- 粒度 (Granularity)： Split 通过复制节点来创建微批次（microbatches），从而实现流水线重叠。
- 顺序 (Ordering)： Order 强制执行子图之间的时序依赖关系，允许用户指定复杂的重叠模式（例如，重叠前向微批次 $i$ 与反向微批次 $i-1$ ）以及资源分配（例如，将特定的逻辑流分配给通信）。
资源控制： 用户可以将逻辑 GPU 流分配给特定的指令，从而实现对通信与计算重叠的精细控制，而无需直接管理物理 CUDA 流。

2. 统一中间表示 (Unified IR)

Piper 的核心是一个全局训练 DAG（有向无环图）。

结构： 该 DAG 表示所有的计算和通信操作。节点要么是块 (Chunks)（粗粒度计算单元），要么是 Comm 节点（通信操作）。
转换： 编译器从模型代码中提取单设备 DAG 开始。用户指令作为图重写操作，将单设备 DAG 转换为包含显式通信节点、微批次处理和设备放置的分布式 DAG。
安全性： 系统保证转换过程遵循模型及用户指令定义的各种数据和时序依赖关系。

3. 与策略无关的分布式运行时

Piper 使用一个集中式调度器和一个基于 Worker 的执行模型（基于 Ray 实现）。

集中式调度器： 它接收转换后的 DAG，并将其分解为每个流水线并行 (PP) 秩 (rank) 独有的子 DAG。它解析分配给同一流的独立任务之间的偏序关系，从而为每个流生成全序关系。
Worker 执行： Workers 加载各自的模型分片并根据调度器的计划执行任务。
资源管理：
- 流 (Streams)： 任务被分发到逻辑流，逻辑流映射到物理 CUDA 流。仅当依赖关系跨越不同流时，才会插入同步（通过 CUDA 事件）。
- 通信： 为了避免死锁和气泡（bubbles），Piper 为点对点收发操作分配了独立的通信器和流，从而确保在不需要全局运行时了解执行时间的情况下，实现一致的顺序。
- 内存： Piper 管理模型状态（参数、梯度、优化器状态）和激活值。它通过在计算前实例化状态并在计算后重新分片（resharding）来实现 ZeRO 式的分片，从而确保在微批次间的正确显存使用。

核心贡献

用户调度接口： 一个用于控制设备间与设备内并行策略的灵活 API，允许用户通过指令指定高层放置和底层执行顺序。
统一 IR (全局训练 DAG)： 一种能够对组合并行策略（如 PP + EP + ZeRO）引入的通信与计算操作进行联合调度的表示形式。
与策略无关的运行时： 一个高效的分布式运行时，能够执行由 IR 定义的任意策略，并自动处理资源竞争与同步问题。
端到端评估： 证明了 Piper 在标准策略上可以达到与通用框架相当的性能，同时能够支持其他框架无法支持的新型高性能策略。

结果

作者在配备 NVIDIA A100 GPU 的 AWS EC2 实例上，将 Piper 与 Megatron-LM、DeepSpeed 和 TorchTitan 进行了对比评估。

常见策略： 在标准的 1F1B 和交错式 (Interleaved) 1F1B 调度下，Piper 的性能与 Megatron-LM 和 TorchTitan 持平。
PP x ZeRO 兼容性：
- 虽然 Megatron、DeepSpeed 和 TorchTitan 在支持 PP 与 ZeRO-2（参数分片）及 ZeRO-3（权重分片）的结合方面存在困难，但 Piper 支持所有此类组合。
- 显存效率： 在 PP x ZeRO-2 和 PP x ZeRO-3 场景下，相比于因无法正确在微批次间重分片权重和梯度而导致显存溢出 (OOM) 的 TorchTitan，Piper 允许使用 3–8 倍更大的 Batch Size。
新颖策略 (DualPipeV)：
- Piper 成功实现了 DualPipeV 调度（结合了 PP 和 EP 以及微批次重叠）。
- 吞吐量提升： 在 1B 模型上，Piper-DualPipeV 比 Piper 的基准 1F1B 调度吞吐量提高了 13%。在 9B 模型上，它比自身的交错式调度提高了 10%，比 Megatron 的交联式调度提高了 6%。
- 对比： TorchTitan 实现的 DualPipeV 相比其基准仅有微乎其微的提升 (3%)，这归因于由基于线程的分发导致的非预期序列化。

重要性

本文声称，Piper 的重要性在于其能够将策略与实现解耦。通过提供统一的 IR 和灵活的 API，Piper 降低了实现复杂的组合并行策略（如 DualPipe）的门槛，而这类策略目前需要大量的人工工程投入。它弥合了通用框架的灵活性与专用定制系统的高性能之间的鸿沟。此外，由于 Piper 可以执行任意策略而不依赖于硬编码的优化，因此它可以作为一个可靠的运行时，服务于未来的自动化搜索系统，从而实现对更广阔的分布式训练配置空间的探索。