🤖 AI

Piper: A Programmable Distributed Training System

Piper is a programmable distributed training system that decouples high-level parallelism strategies from low-level runtime implementations through a unified intermediate representation, enabling flexible composition of diverse strategies like ZeRO and DeepSeek-V3's DualPipe while maintaining performance parity and improving efficiency.

Original authors: Megan Frisella, Shubham Tiwari, Andy Ruan, Yi Pan, Parker Gustafson, Mat Jacob, Gilbert Bernstein, Stephanie Wang

Published 2026-06-10

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Megan Frisella, Shubham Tiwari, Andy Ruan, Yi Pan, Parker Gustafson, Mat Jacob, Gilbert Bernstein, Stephanie Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to bake a massive, complex cake (a giant AI model) that is too big to fit in a single kitchen. You have to hire hundreds of chefs (computer chips) and split the work among them.

The problem is that coordinating these hundreds of chefs is incredibly difficult. If you tell Chef A to mix the batter while Chef B waits for the eggs, you waste time. If Chef A and Chef B both try to grab the same bowl at the same time, they crash.

Currently, most AI training systems are like rigid recipe books. They have a few pre-written ways to split the work (like "everyone mixes their own batter" or "everyone passes the bowl down a line"). If you want to try a new, fancy way of cooking that saves time or space, you have to hire a human expert to rewrite the entire recipe book from scratch. It's slow, expensive, and hard to change.

Piper is a new "smart kitchen manager" that changes the game. Instead of a rigid recipe, Piper gives you a set of simple instructions (like "put the mixing on the left counter" or "pass the bowl while the oven is heating") and a universal blueprint. Piper then figures out exactly how to get the job done efficiently, no matter how complex your instructions are.

Here is how Piper works, broken down into simple concepts:

1. The Blueprint (The "IR")

Imagine the cake recipe isn't just a list of steps, but a giant, 3D flowchart showing every single action every chef needs to take, and how they depend on each other. Piper calls this the Global Training DAG (a fancy term for a map of all the work).

Old way: The system only knew a few fixed flowcharts.
Piper way: You draw your own flowchart using simple tags. Piper takes your drawing and turns it into a master plan that every chef can follow.

2. The Instructions (The "API")

Piper gives you a simple language to tell the system what you want, without needing to know the nitty-gritty details of the hardware.

Annotations: You can tag parts of the model. For example, you can say, "This part of the cake is for the 'Expert' chefs only," or "This part needs to be split across 4 different counters."
Directives: These are your commands. You can say, "Split this big batch into smaller pieces so we can work on them at the same time," or "Make sure the mixing happens while the baking is cooling down."
The Magic: You don't have to tell the computer how to move the bowls or which specific spoon to use. You just say what needs to happen, and Piper figures out the logistics.

3. The Manager (The "Runtime")

Once you give Piper your instructions, it acts as a central traffic controller.

The Central Scheduler: It looks at your blueprint and breaks it down into tiny, specific tasks for each chef. It decides who does what and when.
The Workers: Each chef (computer chip) gets their specific list of tasks. They don't need to talk to each other to figure out the plan; they just follow the list Piper gave them.
Avoiding Traffic Jams: A major problem in AI training is "bubbles"—times when a chef sits idle waiting for a bowl. Piper is smart about overlapping tasks. It might tell Chef A to start mixing the next batch while Chef B is still washing the bowl from the previous one. It ensures that communication (passing bowls) happens at the exact same time as cooking (mixing), so no one is ever waiting around.

Why is this a big deal?

The paper shows that Piper can do two things that current systems struggle with:

It matches the speed of the experts: For standard, common ways of training, Piper is just as fast as the top-tier systems used by big companies today.
It unlocks new, complex strategies: Because Piper is flexible, it can handle "hybrid" strategies that other systems can't.
- Example: The paper mentions a strategy called DualPipe. Imagine a kitchen where you are baking two different types of cakes at the same time, but you have to swap ingredients between the ovens in a very specific, tricky rhythm. Old systems get confused by this and either crash or run very slowly. Piper handles this rhythm perfectly, allowing the chefs to work 6–30% faster and fit much larger cakes into the kitchen without running out of counter space (memory).

The Bottom Line

Piper is like upgrading from a rigid, pre-printed instruction manual to a smart, adaptable project manager. It lets AI researchers design their own custom ways to train giant models without having to become hardware engineers. It makes the system fast enough for today's needs but flexible enough to handle the crazy, complex strategies of tomorrow.

Technical Summary: Piper - A Programmable Distributed Training System

Problem Statement

As machine learning models scale to foundation models requiring hundreds or thousands of accelerators, the complexity of designing distributed training strategies has become a bottleneck. Current approaches face a dichotomy:

Deployed Systems: High-performance systems (e.g., those used for DeepSeek-V3) rely on human experts to manually co-design high-level parallelism strategies (e.g., Pipeline Parallelism (PP), Expert Parallelism (EP)) with low-level execution strategies (e.g., intra-GPU resource scheduling, stream management). This results in rigid systems that are difficult to adapt to new strategies or composed parallelism dimensions.
General-Purpose Frameworks: Frameworks like Megatron-LM, DeepSpeed, and TorchTitan offer flexible interfaces but often treat parallelism dimensions as independent. They eagerly dispatch operations, making it challenging to jointly schedule communication and computation across composed strategies (e.g., overlapping EP all-to-alls with DP all-reduces). Furthermore, they often fail to fully support complex combinations, such as PP with ZeRO-2 or ZeRO-3, due to interactions between layer hooks and execution schedules.

The core challenge is extensibility: creating a system that minimizes the effort to specify and implement arbitrary distributed training strategies while maintaining high performance and memory efficiency.

Methodology

Piper addresses these challenges by decoupling the strategy specification from the runtime implementation through a unified Intermediate Representation (IR). The system consists of three main components:

1. User Scheduling API

Piper provides a declarative API that allows users to control both high-level placement and low-level execution without writing custom schedulers.

Annotations: Users tag semantically meaningful regions of a PyTorch model (e.g., specific layers, expert blocks) using context managers. These tags define the boundaries for parallelism dimensions (e.g., PP, EP, DP).
Scheduling Directives: Users apply directives to transform the system's IR. Key directives include:
- Placement: Place, Replicate, and Shard define device mapping and introduce necessary communication nodes (P2P, All-Reduce, All-Gather, All-to-All).
- Granularity: Split replicates nodes to create microbatches, enabling pipeline overlapping.
- Ordering: Order enforces temporal dependencies between sub-graphs, allowing users to specify complex overlapping patterns (e.g., overlapping forward microbatch $i$ with backward microbatch $i-1$ ) and resource allocation (e.g., assigning specific logical streams to communication).
Resource Control: Users can assign logical GPU streams to specific directives, enabling fine-grained control over the overlap of communication and computation without managing physical CUDA streams directly.

2. Unified Intermediate Representation (IR)

The core of Piper is a Global Training DAG (Directed Acyclic Graph).

Structure: The DAG represents all computation and communication operations. Nodes are either Chunks (coarse-grained compute units) or Comm nodes (communication operations).
Transformation: The compiler starts with a single-device DAG extracted from the model code. User directives act as graph rewrites, transforming the single-device DAG into a distributed DAG with explicit communication nodes, microbatching, and device placements.
Safety: The system guarantees that transformations respect data and temporal dependencies defined by the model and user directives.

3. Strategy-Agnostic Distributed Runtime

Piper uses a centralized scheduler and a worker-based execution model (implemented on Ray).

Centralized Scheduler: It takes the transformed DAG and decomposes it into unique sub-DAGs for each Pipeline Parallelism (PP) rank. It resolves partial orderings between independent tasks assigned to the same stream, generating a total order per stream.
Worker Execution: Workers load their specific model shards and execute tasks according to the scheduler's plan.
Resource Management:
- Streams: Tasks are dispatched to logical streams, which map to physical CUDA streams. Synchronization (via CUDA events) is inserted only when dependencies cross streams.
- Communication: To avoid deadlocks and bubbles, Piper allocates separate communicators and streams for point-to-point send/receive operations, ensuring consistent ordering without requiring global runtime knowledge of execution times.
- Memory: Piper manages memory for model states (parameters, gradients, optimizer states) and activations. It supports ZeRO-style sharding by materializing states before computation and resharding them afterward, ensuring correct memory usage across microbatches.

Key Contributions

User Scheduling Interface: A flexible API for controlling inter- and intra-device parallelism strategies, allowing users to specify high-level placement and low-level execution ordering via directives.
Unified IR (Global Training DAG): A representation that enables the joint scheduling of communication and computation operations introduced by composed parallelism strategies (e.g., PP + EP + ZeRO).
Strategy-Agnostic Runtime: An efficient distributed runtime capable of executing arbitrary strategies defined by the IR, handling resource contention and synchronization automatically.
End-to-End Evaluation: Demonstration that Piper matches the performance of general-purpose frameworks on standard strategies while enabling novel, high-performance strategies that other frameworks cannot support.

Results

The authors evaluated Piper against Megatron-LM, DeepSpeed, and TorchTitan on AWS EC2 instances with NVIDIA A100 GPUs.

Common Strategies: Piper achieves performance parity with Megatron-LM and TorchTitan on standard 1F1B and Interleaved 1F1B schedules.
PP x ZeRO Compatibility:
- While Megatron, DeepSpeed, and TorchTitan struggle to support PP combined with ZeRO-2 (parameter sharding) and ZeRO-3 (weight sharding), Piper supports all combinations.
- Memory Efficiency: In PP x ZeRO-2 and PP x ZeRO-3 scenarios, Piper allows for 3–8x larger batch sizes compared to TorchTitan, which fails to properly reshard weights and gradients between microbatches, leading to Out-Of-Memory (OOM) errors.
Novel Strategies (DualPipeV):
- Piper successfully implements the DualPipeV schedule (combining PP and EP with microbatch overlapping).
- Throughput Gains: On a 1B model, Piper-DualPipeV improves throughput by 13% over Piper's baseline 1F1B schedule. On a 9B model, it improves by 10% over its own interleaved schedule and 6% over Megatron's interleaved schedule.
- Comparison: TorchTitan's implementation of DualPipeV showed negligible improvement (3%) over its baseline, attributed to unintended serialization caused by thread-based dispatching.

Significance

The paper claims that Piper's significance lies in its ability to decouple strategy from implementation. By providing a unified IR and a flexible API, Piper lowers the barrier for implementing complex, composed parallelism strategies (like DualPipe) that currently require extensive human engineering. It bridges the gap between the flexibility of general-purpose frameworks and the performance of specialized, hand-tuned systems. Furthermore, Piper serves as a viable runtime for future automated search systems, as it can execute arbitrary strategies without hard-coding specific optimizations, thereby enabling the exploration of a much larger space of distributed training configurations.