A Flexible Programmable Pipeline Parallelism Framework for Efficient DNN Training

2025-09-27

AI Summary

Existing pipeline parallelism approaches rely on predefined scheduling policies, limiting adaptability to emerging DNN architectures, incurring high manual coding overhead, and lacking automated schedule exploration capability. This paper proposes FlexPipe, a flexible, programmable pipeline parallelism framework that enables declarative specification of micro-batch execution order via a domain-specific language (DSL) and integrates a lightweight, multi-objective automated scheduler for low-overhead schedule search and dynamic operator extension. Its core innovation lies in empowering users with full scheduling control while simultaneously achieving high automation—thereby significantly improving both development efficiency and runtime performance without sacrificing flexibility. Experiments demonstrate that FlexPipe achieves up to 2.28× speedup over Megatron-LM and 1.49× higher training efficiency than the state-of-the-art automatic scheduling framework across diverse models.

Technology Category

• Planning, Routing, and Scheduling: Planning with Language Models
• Machine Learning: Deep Neural Architectures and Foundation Models
• Natural Language Processing: Code Generation / Program Synthesis from Natural Language

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

Abstract

Pipeline parallelism is an essential distributed parallelism method. Increasingly complex and diverse DNN models necessitate meticulously customized pipeline schedules for performance. However, existing practices typically rely on predefined schedules, each with strengths, but fail to adapt automatically to the emerging model architectures. Exploring novel high-efficiency schedules is daunting due to the enormous and varying schedule space. Besides, manually implementing schedules can be challenging due to the onerous coding burdens and constantly changing needs. Unfortunately, existing frameworks have limitations in automated schedule exploration and lack flexibility and controllability. This paper presents FlexPipe, a programmable pipeline parallelism framework with enhanced productivity, programmability, debuggability, and ease of tuning. FlexPipe has two main components: a succinct domain-specific language (DSL) and an automated scheduler. FlexPipe enables automated schedule exploration for various parallel scenarios within a broad spectrum of schedule types at a small search cost. Besides, users can swiftly develop and customize schedules using the FlexPipe DSL, which embodies flexible controllability in the pipeline order of micro-batch computations over stages. It also provides convenient mechanisms to include new operations in schedules to meet changing demands. Our evaluation results demonstrate that FlexPipe achieves up to 2.28X performance speedup compared to the popular large-scale parallel framework Megtron-LM, and gains up to 1.49X performance speedup compared to the state-of-the-art automated pipeline parallelism framework.

Problem

• Automating pipeline schedule exploration for diverse DNN models
• Reducing manual coding burdens in customizing pipeline parallelism
• Enhancing flexibility and controllability in pipeline scheduling systems

Innovation

• Uses a domain-specific language for flexible pipeline control
• Automates schedule exploration across diverse parallel scenarios
• Enables rapid customization of micro-batch computation ordering