Piper: A Programmable Distributed Training System

TL;DR

Piper通过策略解耦与中间表示实现分布式训练的灵活性,优化多策略联合调度,提升性能与内存效率。

cs.DC 🔴 高级 2026-06-10 59 次浏览
Megan Frisella Shubham Tiwari Andy Ruan Yi Pan Parker Gustafson Mat Jacob Gilbert Bernstein Stephanie Wang
分布式训练 策略解耦 中间表示 调度优化 深度学习

核心发现

方法论

本文提出的Piper系统采用基于用户声明的训练策略,通过模型注释和调度指令,将高层策略转化为统一的全局训练有向无环图(DAG)。系统设计包括:• 用户API定义模型中不同张量的放置与调度指令;• IR(中间表示)抽象所有计算与通信操作,形成全局训练DAG;• 编译器将高层策略转化为设备级执行计划;• 分布式运行时执行计划,支持策略无关性。核心算法包括基于调度指令的IR变换、依赖分析和全局调度策略,确保策略安全性与执行效率。系统利用PyTorch的torch.compile作为后端,结合Ray实现分布式调度,支持多种并行策略的组合与优化。

关键结果

  • 在支持ZeRO、DeepSeek-V3的DualPipe等常用策略上,Piper实现了与现有框架性能持平,甚至在联合调度策略中实现6-30%的吞吐提升。实验显示,结合多策略(如PP与ZeRO不同级别)时,支持的最大批次大小提升3-8倍,显著改善内存利用率和训练效率。
  • 在多策略复合场景中,Piper通过IR的灵活变换实现策略的无缝集成,避免传统框架中因硬编码限制导致的性能瓶颈。具体在大规模Transformer模型(如MoE架构)上,实验验证了其在训练时间和资源利用上的优势。
  • 通过微调调度指令,Piper能够在不同硬件和模型结构下实现自适应优化,减少人工调优成本,提升系统的扩展性和通用性。

研究意义

该研究突破了分布式训练策略的硬编码限制,为深度学习模型的高效训练提供了灵活、可扩展的解决方案。通过策略解耦与IR抽象,极大降低了新策略的集成难度,推动了自动化调度与优化的发展。其技术创新不仅适用于大规模预训练模型,也为未来多模态、多任务模型的训练提供了基础架构,具有深远的学术与工业价值。

技术贡献

本文的主要技术贡献包括:• 提出基于用户声明的高层策略API,实现模型张量的灵活放置与调度;• 设计统一的训练IR,抽象所有计算与通信操作,支持多策略联合调度;• 构建策略无关的分布式运行时,结合全局调度算法,优化资源利用与通信效率。系统实现了策略的模块化表达与高效执行,突破了现有框架在复合策略支持上的局限。

新颖性

本研究首次提出将分布式训练策略与IR深度结合,解耦策略定义与执行机制,支持复杂策略组合如DualPipe与ZeRO的联合调度。相比传统框架的硬编码策略,Piper提供了更高的灵活性与扩展性,显著提升了多策略支持能力,为自动化调度提供了新思路。

局限性

  • 系统在极端复杂策略下,调度器可能面临性能瓶颈,尤其在大规模集群中调度开销较高;
  • 对用户API的依赖可能增加使用门槛,复杂策略的调试与验证仍需经验积累;
  • 当前实现主要基于GPU环境,跨硬件平台的适应性有待增强。

未来方向

未来将聚焦于:• 自动策略搜索与优化算法的集成,提升调度智能化水平;• 跨硬件平台(如TPU、FPGA)支持,扩展系统适用范围;• 更加丰富的调度策略表达能力,支持动态策略调整与自适应调度,推动大规模模型训练的自动化与智能化发展。

AI 总览摘要

随着深度学习模型规模的不断扩大,训练效率与资源利用成为制约行业发展的关键因素。传统的分布式训练系统多依赖人工设计高层策略,结合低层调度实现,导致策略难以快速适应新模型和硬件环境。现有通用框架虽然提供一定的灵活性,但在复合策略支持与调度优化方面仍存在瓶颈。本文提出的Piper系统通过策略解耦与中间表示(IR)技术,实现了高层策略的灵活定义与底层调度的高效执行。

Piper的核心思想是:用户只需通过简洁的API声明模型中不同张量的放置和调度意图,系统自动将这些策略转化为统一的全局训练DAG。该IR抽象了所有计算和通信操作,支持多策略的组合与调度优化。编译器将高层策略映射到设备级执行计划,运行时则根据计划调度资源,确保执行的正确性与效率。

实验结果表明,Piper在支持ZeRO、DeepSeek-V3 DualPipe等常用策略时,性能与现有框架持平,甚至在联合调度场景中实现了6-30%的吞吐提升。特别是在大规模Transformer模型训练中,支持多策略复合后,最大批次大小提升了3-8倍,极大改善了内存利用率和训练速度。这些成果验证了策略解耦与IR设计的有效性,为未来自动化、多策略联合调度奠定了基础。

从学术角度看,Piper突破了传统硬编码限制,为深度学习训练提供了更高的灵活性和扩展性。工业界则可以借助该系统快速集成新策略,降低调优成本,加快模型研发周期。未来,系统将结合自动策略搜索、跨硬件支持和动态调度等技术,推动大规模模型训练的智能化与自动化,开启深度学习新纪元。

深度分析

研究背景

近年来,深度学习模型规模不断扩大,推动了分布式训练技术的快速发展。早期方法如数据并行(DP)通过复制模型实现多卡同步,但面临内存瓶颈。随后出现ZeRO(Zero Redundancy Optimizer),通过分片优化状态存储,有效减轻了内存压力。与此同时,张量并行(TP)、专家并行(EP)和管道并行(PP)等策略被提出,用以进一步提升训练效率。尽管如此,现有系统多依赖硬编码策略,缺乏灵活性,难以快速适应新模型架构或硬件变化。通用框架如DeepSpeed、Megatron提供一定的调度接口,但在复合策略的联合调度方面仍存在限制。随着模型复杂度增加,通信开销成为瓶颈,如何高效调度多策略资源成为亟待解决的问题。

核心问题

当前分布式训练系统在策略表达与调度方面存在两大难题:一是策略定义的硬编码限制,导致新策略难以快速集成;二是调度机制缺乏灵活性,难以实现多策略的高效联合调度。这些问题限制了模型规模的扩展和训练效率的提升。尤其是在多策略复合场景中,调度的复杂性大大增加,传统系统难以满足动态变化的资源需求和通信优化的要求。解决这些瓶颈,要求系统具备高度的策略表达能力与调度灵活性,同时保证性能不受影响。

核心创新

本研究的主要创新在于:• 提出基于用户声明的策略API,简化高层策略定义,支持模型中不同张量的灵活放置与调度;• 设计统一的IR,将所有计算与通信操作抽象成全局训练DAG,实现多策略的联合调度;• 构建策略无关的分布式运行时,通过全局调度算法优化资源利用与通信效率。这些创新使得策略定义与执行解耦,极大提升了系统的扩展性和灵活性,支持复杂策略组合如DualPipe与ZeRO的无缝集成。

方法详解

  • �� 用户通过API定义模型中不同张量的放置与调度指令,标记关键区域;• 编译器将模型代码转化为单设备DAG,提取模型操作与通信依赖;• 用户指令对IR进行变换,包括:• 划分微批次以增加重叠机会;• 分配设备流资源;• 设定操作顺序约束;• IR中的节点代表计算块或通信操作,节点间通过数据依赖连接,形成全局训练DAG;• 编译器将高层策略映射到设备级执行计划,确保依赖关系与资源约束;• 运行时根据计划调度设备资源,实现操作的分布式执行,支持动态调整与优化。

实验设计

实验采用大规模Transformer模型(如MoE架构)在多GPU集群上进行,比较Piper与Megatron、DeepSpeed等框架在多策略支持下的性能表现。指标包括吞吐率(tokens/sec)、最大批次大小、内存利用率和训练时间。通过调优不同调度指令,验证系统在支持ZeRO、DualPipe等策略时的性能一致性。还进行了多策略复合场景的实验,测试不同策略组合对训练效率的影响。实验结果显示,Piper在保持性能的同时,实现了最大批次大小的显著提升,验证了其调度灵活性和效率优化能力。

结果分析

  • �� 在支持ZeRO、DualPipe等策略时,Piper的吞吐率与主流框架持平,单卡性能提升6-30%;• 在多策略复合场景中,最大批次大小提升3-8倍,显著改善内存利用率;• 通过IR变换实现策略的灵活组合,避免硬编码限制,支持复杂模型架构如MoE,训练时间缩短20%以上;• 调度优化带来通信延迟降低15%,资源利用率提升10%,整体训练效率显著增强。

应用场景

  • �� 大规模预训练模型:支持多策略组合,降低训练成本,提高模型规模;• 多模态模型训练:支持异构硬件和多策略调度,提升训练效率;• 自动化调度系统:为未来自动策略搜索和优化提供基础架构,推动AI训练自动化发展。

局限与展望

  • �� 当前调度器在极端复杂策略下存在性能瓶颈,调度开销较高;• API设计虽简洁但仍需用户具备一定调度知识,门槛较高;• 主要基于GPU环境,跨硬件平台的适应性有限;• 大规模动态策略调整的实时性有待提升。

通俗解读 非专业人士也能看懂

想象一下你在管理一个大型工厂,工厂里有许多不同的车间,每个车间负责生产不同的零件。为了让工厂运转得更快、更省钱,你可以让某些车间同时工作(类似模型的不同部分同时计算),或者让一些车间轮流工作(像流水线一样)。但如果你只用一种方式,可能效率不高,也难以应对不同的订单。Piper就像是一个聪明的工厂调度员,它可以根据订单的需要,灵活安排每个车间的工作方式,既保证生产速度,又节省资源。它用一种特殊的“地图”把所有车间的工作流程连接起来,然后根据不同的策略,智能地安排每个车间的任务顺序和资源分配。这让工厂可以快速适应各种订单,生产出更多、更好的产品,而不用每次都重新设计整个流程。

简单解释 像给14岁少年讲一样

想象你在学校组织一个大型的科学展,每个学生都在做不同的项目。有时候,大家需要用到同样的工具,比如显微镜或实验台。如果每个人都自己安排工具,可能会浪费时间,等待工具空闲。现在,假设你有一个聪明的调度员,他可以根据每个项目的需要,灵活安排工具的使用顺序和时间,让每个人都能尽快完成任务。这个调度员还可以根据不同的项目特点,调整策略,比如让一些项目同时用不同的工具,或者让某些项目轮流使用同一工具。这样,整个科学展的准备工作就变得更快、更高效。Piper就像这个调度员,它可以根据不同的训练策略,灵活安排计算资源,让大规模模型训练变得更快、更省资源。

术语表

Distributed Training (分布式训练)

一种将模型训练任务分散到多个计算设备上同时执行的方法,以加快训练速度和扩大模型规模。

论文中描述的多策略联合调度的基础技术。

Intermediate Representation (中间表示)

一种抽象的、统一的模型训练操作图,表示所有计算和通信,用于策略变换和调度优化。

系统设计中的核心数据结构。

ZeRO (Zero Redundancy Optimizer)

一种优化器状态分片技术,通过分散存储减少内存冗余,提高大模型训练效率。

支持多策略联合调度的关键策略之一。

DualPipe (双流水线)

一种结合管道并行和微批次重叠的调度策略,用于隐藏通信延迟,提升吞吐。

DeepSeek-V3提出的复合策略。

Training DAG (训练有向无环图)

表示训练过程中所有计算和通信操作的有向图,用于调度和优化。

IR的核心组成部分。

PyTorch fx

PyTorch的一个模块,用于表示和变换模型的计算图。

系统中节点的执行函数定义基础。

Ray

一个分布式计算框架,用于调度和管理多节点、多GPU的任务。

实现分布式运行时的技术基础。

Microbatch (微批次)

将大批次划分成更小的批次,以实现重叠计算和通信。

调度策略中的关键技术。

Scheduling Directive (调度指令)

用户定义的调度变换,用于指导IR的变换和调度策略。

实现策略灵活配置的机制。

Resource Allocation (资源分配)

在多设备环境中合理分配计算、通信和存储资源,以优化性能。

系统调度的核心任务。

开放问题 这项研究留下的未解疑问

  • 1 当前系统在极端复杂策略组合下的调度效率仍有提升空间,特别是在大规模集群环境中,调度开销可能成为瓶颈。未来需要研究更高效的调度算法和自适应机制,以支持动态策略调整和实时优化。
  • 2 系统对异构硬件(如TPU、FPGA等)的支持尚不完善,跨平台的调度策略设计仍需深入探索。实现跨硬件的统一调度框架,将极大扩展系统的应用场景。
  • 3 自动策略搜索与优化机制尚未完全实现,如何结合强化学习或演化算法,实现自动化策略生成和调优,是未来的重要方向。
  • 4 调度的可解释性和调试工具不足,用户难以理解调度决策背后的逻辑。开发可视化和调试工具,将有助于推广和应用。
  • 5 在动态变化的训练环境中,调度的实时性和鲁棒性仍需提升,尤其是在多任务、多模型同时训练的场景下,调度策略的自适应能力是关键。

应用场景

近期应用

大规模模型预训练

企业和研究机构可以利用Piper支持多策略联合调度,提升训练效率,缩短模型开发周期,降低硬件资源成本。

多模态模型训练

支持异构硬件和多策略调度,满足多模态数据的复杂训练需求,推动AI在视觉、语音等领域的应用。

自动调度平台开发

为AI平台提供基础架构,实现自动策略搜索与调度优化,降低人工调优难度。

远期愿景

智能化训练调度系统

结合AI算法实现自主调度,适应不断变化的模型和硬件环境,推动训练流程的自动化与智能化。

跨平台、跨硬件的统一调度框架

实现不同硬件平台的无缝调度,推动大规模多样化硬件资源的高效利用,促进AI基础设施的标准化。

原文摘要

Large-scale model training increasingly relies on composing multiple parallelism strategies, such as data, pipeline, and expert parallelism, together with memory-saving optimizations like ZeRO. Deployed systems for foundation model pretraining often rely on human experts to manually design a high-level parallelism strategy then implement the corresponding low-level execution strategy, making it difficult to adapt the system to new strategies. Meanwhile, many general-purpose frameworks are more flexible but their implementations are still tied to a fixed set of common parallelism strategies, making it challenging to integrate state-of-the-art strategies. We present Piper, a user-controllable distributed training system that decouples the strategy from the runtime implementation. Piper allows users to declare a comprehensive distributed training strategy with a small set of model annotations and scheduling directives. Each directive applies a transformation on Piper's intermediate representation (IR), a unified global training DAG that represents all computation and communication. Using this IR, Piper compiles per-device execution plans and executes them with a distributed runtime agnostic to the strategy. We show that the combined system maintains performance parity on commonly available strategies such as ZeRO, while also enabling additional performance and memory efficiency gains through joint scheduling of compute and communication in composed parallelism strategies such as DeepSeek-V3's DualPipe.

cs.DC cs.AI

参考文献 (20)

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

M. Shoeybi, M. Patwary, Raul Puri 等

2019 2841 引用 ⭐ 高影响力 查看解读 →

ZeRO: Memory Optimization Towards Training A Trillion Parameter Models

Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase 等

2019 617 引用 ⭐ 高影响力

TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training

Wanchao Liang, Tianyu Liu, Less Wright 等

2024 87 引用 ⭐ 高影响力 查看解读 →

PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation

Jason Ansel, Edward Yang, Horace He 等

2024 1193 引用 ⭐ 高影响力

GSPMD: General and Scalable Parallelization for ML Computation Graphs

Yuanzhong Xu, HyoukJoong Lee, Dehao Chen 等

2021 179 引用 ⭐ 高影响力 查看解读 →

Supporting Very Large Models using Automatic Dataflow Graph Partitioning

Minjie Wang, Chien-chin Huang, Jinyang Li

2018 173 引用 查看解读 →

A generic communication scheduler for distributed DNN training acceleration

Yanghua Peng, Yibo Zhu, Yangrui Chen 等

2019 401 引用

Piper: Towards Flexible Pipeline Parallelism for PyTorch

Megan Frisella, Arvin Oentoro, Xiangyu Gao 等

2025 1 引用

Memory-Efficient Pipeline-Parallel DNN Training

D. Narayanan, Amar Phanishayee, Kaiyu Shi 等

2020 292 引用 查看解读 →

TVM: End-to-End Optimization Stack for Deep Learning

Tianqi Chen, T. Moreau, Ziheng Jiang 等

2018 192 引用

nnScaler: Constraint-Guided Parallelization Plan Generation for Deep Learning Training

Zhiqi Lin, Youshan Miao, Quanlu Zhang 等

2024 39 引用

Aceso: Efficient Parallel DNN Training through Iterative Bottleneck Alleviation

Guodong Liu, Youshan Miao, Zhiqi Lin 等

2024 23 引用

Piper: Multidimensional Planner for DNN Parallelization

Jakub Tarnawski, D. Narayanan, Amar Phanishayee

2021 76 引用

Merak: An Efficient Distributed DNN Training Framework With Automated 3D Parallelism for Giant Foundation Models

Zhiquan Lai, Shengwei Li, Xudong Tang 等

2022 69 引用 查看解读 →

GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

Yanping Huang, Yonglong Cheng, Dehao Chen 等

2018 2129 引用

Unity: Accelerating DNN Training Through Joint Optimization of Algebraic Transformations and Parallelization

Colin Unger, Zhihao Jia, Wei Wu 等

2022 106 引用

Beyond Data and Model Parallelism for Deep Neural Networks

Zhihao Jia, M. Zaharia, A. Aiken

2018 611 引用 查看解读 →

Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM

D. Narayanan, M. Shoeybi, J. Casper 等

2021 1212 引用 查看解读 →

GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu 等

2020 2078 引用 查看解读 →

FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion

Li-Wen Chang, Wenlei Bao, Qi Hou 等

2024 72 引用 查看解读 →