Introducing Piper: A Programmable Distributed Training System | SyFI Lab
Introducing Piper: A Programmable Distributed Training System
TL;DR: New distributed training strategies and optimizations should not require new distributed runtimes.<br>Large training jobs increasingly combine multiple parallelism strategies such as pipeline, data, and expert parallelism with ZeRO-style sharding, creating placement and GPU scheduling choices that<br>current frameworks cannot express cleanly.<br>Today, ML researchers and practitioners choose between building one-off specialized systems that perform well but are hard to adapt, and general-purpose frameworks that are easier to use but expose limited control.<br>Piper is a user-controllable distributed training system for PyTorch that separates model placement and GPU scheduling from model code and runtime<br>implementation.<br>With lightweight model annotations and a small scheduling language, Piper lets users express, visualize, profile, and run high-performance training schedules such as DualPipe-style pipeline- and expert-parallel overlap.
paper: https://arxiv.org/pdf/2606.11169
code: https://github.com/uw-syfi/piper
Composed parallelism dimensions have complex communication patterns
Large model training commonly composes several parallelism dimensions such as:
Data parallelism (DP) replicates model state, runs different data on each replica, and synchronizes gradients with collective communication.
Pipeline parallelism (PP) shards layers into stages and sends activations and gradients between stages with point-to-point communication. Data batches are split into microbatches to keep the pipeline busy.
Expert parallelism (EP) shards the experts in a mixture-of-experts (MoE) layer and routes tokens to and from expert subsets with collective communication.
Tensor parallelism (TP) shards individual tensor operators, such as matrix multiplications, and uses collective communication to assemble partial results.
ZeRO-style sharding reduces redundant optimizer, gradient, and parameter state across DP ranks by sharding model state and introducing additional collective communication to gather and synchronize shards.
Composing parallelism strategies is not one-size-fits-all; the right choice depends on the model architecture, memory constraints, and network topology.<br>Figure 1 depicts a mixture-of-experts (MoE) model with pipeline parallelism across layers and data/expert parallelism within layers.
Figure 1. A PP x EP x DP placement for a mixture-of-experts model.<br>Pipeline parallelism places layers across stages, expert parallelism shards expert MLPs, and data parallelism replicates non-expert components like attention.
Composing multiple dimensions creates complex communication patterns and high communication overhead.<br>A single training step for the Figure 1 placement requires coordinating PP microbatches flowing forwards and backwards through the model, critical-path EP token routing in the forward and backward pass, and DP gradient synchronization.<br>Thus, maximizing training throughput requires carefully scheduling tensor operators on each GPU to hide communication latency while avoiding bubbles (GPU idle time).
Scheduling tensor operators onto each GPU is hard
MoE training demonstrates the need for fine-grained control over GPU scheduling.<br>EP adds critical-path all-to-all (A2A) collective communication around expert computation.<br>DeepSeek-V3 reports that token routing alone produces an approximately 1:1 computation-to-communication ratio in their setup, as experts are distributed across slow inter-node links.<br>To hide this latency, they introduced the DualPipe schedule, which overlaps expert computation with collective communication from different pipeline-parallel microbatches.<br>Figure 2 shows a DualPipe schedule variant and highlights an overlapped forward-backward microbatch pair.
Figure 2. DualPipeV schedule for 2-way PP and 4 microbatches.<br>Numbers are microbatch IDs.<br>The bolded cells are overlapped forward-backward microbatch pairs.
Scheduling operators onto the GPU is not simple – especially when DP is composed with EP, adding collective all-reduce (AR) communication in the backward pass.<br>Figure 3 shows different GPU scheduling choices, where streams represent GPU parallelism.
Figure 3. Stream scheduling choices for DP all-reduce and EP all-to-all in an overlapped microbatch pair.<br>The best choice depends on kernel running times, dispatch ordering, and critical-path dependencies.
Putting A2A and AR on separate streams (a) lets them run concurrently but risks interference on network bandwidth.<br>Putting A2A and AR on the same stream (b) avoids network interference by serializing the collectives, but can delay communication.<br>Breaking up the AR into finer-grained units (c) can reduce interference (parameter bucketing is a common example of this strategy), but it’s still hard to predict the overall effect on training throughput, as partitioning can reduce communication...