Scaling Deep Learning Training with MPMD Pipeline Parallelism

Scaling Deep Learning Training with MPMD Pipeline Parallelism

Speakers: Anxhelo Xhebraj, Sean Lee, Hanfeng Chen, Vinod Grover

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

In the rapidly evolving landscape of deep learning, the relentless growth in model size necessitates increasingly sophisticated strategies for distributed training. As models surpass the capacity of single devices or even single nodes, practitioners must leverage vast clusters of GPUs and other accelerators. While techniques like tensor parallelism and fully sharded data parallelism (FSDP) offer relatively straightforward implementations often expressible through Single Program, Multiple Data (SPMD) paradigms with collective operations, the general form of pipeline parallelism presents a far more intricate challenge. This talk, presented by Jong Lee from Nvidia at MLSys 2025, introduces Jack's PP, a novel and comprehensive solution designed to abstract away the complexities of pipeline parallelism within the JAX ecosystem.

The core problem Jack's PP addresses is the perception that implementing pipeline parallelism is "tedious, time-consuming, and error-prone," famously captured by the quote, "no code base survives pipeline parallelism." This difficulty stems from pipeline parallelism's temporal nature, requiring computation to be split into potentially non-identical stages, with complex and often non-linear data flows. Jack's PP aims to provide a general-purpose framework that allows users to express pipeline stages with minimal annotations, decouple stage definitions from their execution placement, and support arbitrary scheduling strategies, thereby making advanced distributed training techniques accessible and robust for large-scale model development.

This article delves into the technical underpinnings of Jack's PP, exploring its design principles, how it integrates with JAX's auto-differentiation and compilation pipeline, and the significant performance gains it achieves. By automating critical aspects like communication inference, placement of auxiliary computations, and support for various schedules, Jack's PP stands as a pivotal contribution towards simplifying and accelerating the training of next-generation large language models and other complex deep learning architectures.

Background

▶ Watch: Introduction: The challenge of pipeline parallelism (0:00)

The exponential growth of deep learning models, particularly in domains like natural language processing, has pushed the boundaries of what single computational devices can handle. Models with hundreds of billions or even trillions of parameters require distributed training across hundreds or thousands of GPUs. To effectively utilize these massive hardware resources, various parallelism strategies have emerged.

Data parallelism, where identical model replicas process different subsets of a batch, is the most common form. However, as model size increases, even a single replica might not fit into device memory. This led to model parallelism techniques, which partition the model itself across multiple devices. Tensor parallelism (or intra-layer parallelism) splits individual layers or tensors across devices, while fully sharded data parallelism (FSDP) shards model parameters, gradients, and optimizer states across data parallel ranks. These methods are often well-suited for SPMD programming models, where a single program runs on all devices, and communication is handled via collective operations (e.g., all-reduce, all-gather).

Pipeline parallelism, however, introduces a different dimension of complexity. Instead of replicating or sharding layers, it partitions the model's layers into sequential "stages," each executed on a different group of devices. The output of one stage becomes the input to the next, creating a computational pipeline. Initially, this was done sequentially, with each stage waiting for the previous one to complete. To improve throughput and keep GPUs busy, microbatching was introduced. The global batch is split into smaller microbatches, which are then fed into the pipeline in a staggered fashion. This allows multiple stages to operate on different microbatches concurrently, filling the "pipeline bubble" – the idle time at the beginning and end of a pipeline execution.

Early pipeline parallelism schedules, such as G-pipe (Google Pipe) introduced in 2018, primarily focused on the forward pass, with the backward pass executed in reverse order. While G-pipe significantly improved throughput over purely sequential execution, it often incurred high activation memory costs and large pipeline bubbles. Subsequent innovations, like 1F1B (one-forward-one-backward), aimed to reduce peak activation memory by scheduling backward computations as soon as their dependencies are met, effectively overlapping forward and backward passes. More recent advancements like Dual Pipe continue to refine these schedules, seeking optimal throughput and memory efficiency.

Despite these advancements, implementing pipeline parallelism manually remains notoriously difficult. The challenges are multifaceted:

1. Temporal and Non-SPMD Nature: Unlike data or tensor parallelism, pipeline stages are not identical programs running on all processes. Different stages perform different computations.
2. Complex Data Flow: Data dependencies can be intricate, including residual connections that span multiple stages, requiring non-adjacent communication.
3. Arbitrary Schedules: Supporting various schedules (G-pipe, 1F1B, interleaved 1F1B, etc.) means adapting the execution order and communication patterns for each.
4. Non-Identical Stages: Stages often have different numbers of layers or different types of operations, making uniform partitioning difficult.
5. Manual Placement of Tensors and Weights: Users must manually decide where input tensors (e.g., X, Y in an example), intermediate activations, and model weights should reside.
6. Explicit Send/Receive Operations: Implementing communication across stages requires explicit insertion of send and receive primitives, a highly error-prone task, especially with complex data flows.
7. Shared Weights: Handling weights shared between multiple stages (e.g., tied embeddings) requires careful synchronization and placement logic.
8. Pre/Post-Loop Computations: Additional computations before and after the main gradient accumulation loop (e.g., input preprocessing, optimizer state updates, gradient application) must be selectively enabled or disabled based on the pipeline rank, adding further complexity.

These challenges collectively contribute to the high barrier to entry for pipeline parallelism, making it a specialized domain often requiring deep systems-level expertise. Jack's PP emerges as a direct response to these pain points, aiming to democratize access to efficient pipeline parallelism by automating much of this manual orchestration.

Key Findings

▶ Watch: Manual complexities of general pipeline parallelism (3:00)

Jack's PP introduces a paradigm shift in how pipeline parallelism is approached within the JAX ecosystem, transforming a historically tedious task into a more manageable and performant process. The key findings and contributions of this work are:

• General-Purpose Pipeline Parallelism Solution: Jack's PP provides a robust and flexible framework for implementing pipeline parallelism that goes beyond the constrained forms typically found in SPMD-based systems. It supports arbitrary computations within stages, complex data flows, and various scheduling strategies.
• Minimal User Annotation: A core objective of Jack's PP is to minimize the boilerplate code and manual orchestration required from the user. It achieves this by introducing only two primary primitives: pipeline_yields to define stage boundaries and accumulate_grads to specify the gradient accumulation loop and desired schedule.
• Decoupled Stage Definition and Placement: Jack's PP intelligently separates the logical definition of pipeline stages from their physical placement and execution order on devices. This allows users to define their model's computational graph naturally, while the system handles the intricate details of distributed execution.
• Automatic Inference of Communication and Placement: Crucially, Jack's PP automatically infers and inserts necessary send and receive operations for inter-stage communication, even across non-adjacent ranks. It also intelligently places pre- and post-gradient accumulation loop computations on the appropriate MPMD ranks based on data dependencies.
• Support for Diverse Schedules: Unlike systems limited to a single schedule (e.g., G-pipe), Jack's PP natively supports a variety of advanced pipeline schedules, including 1F1B and its variants like interleaved 1F1B. This flexibility allows users to select schedules optimized for specific memory or throughput requirements.
• Composition with JAX Features: Jack's PP primitives are designed to compose seamlessly with existing JAX features, including automatic differentiation (autodiff) and rematerialization. This ensures that the benefits of JAX's functional programming model and high-performance compilation are preserved.
• Significant Performance Gains: Benchmarked on a large-scale GPT-3 175B model, Jack's PP demonstrates substantial performance improvements:
• It is approximately 44% faster than the constrained form of pipeline parallelism supported by SPMD in JAX, primarily due to its ability to utilize more efficient schedules (like interleaved 1F1B) and better overlap communication with computation.
• It achieves performance within 5-6% of highly optimized, model-specific frameworks like Nemo, despite being entirely model agnostic.
• Exceptional Scaling Efficiency: Jack's PP exhibits impressive scalability, achieving 98% scaling efficiency when scaling from 128 GPUs to 1024 GPUs for GPT-3 175B training. This highlights its robustness for extreme-scale distributed training.
• Open-Sourced and Pure Python: Jack's PP is implemented as a pure Python package on top of JAX, requiring no modifications to the underlying XLA compiler. It has been open-sourced, encouraging community adoption and feedback.

These findings underscore Jack's PP as a critical enabler for training increasingly massive deep learning models, significantly lowering the implementation barrier for efficient pipeline parallelism while delivering state-of-the-art performance and scalability.

Technical Deep Dive

▶ Watch: Introducing JAX PP: A general solution for pipeline parallelism (5:00)

Jack's PP integrates deeply into the JAX ecosystem, leveraging its functional programming model and advanced compilation pipeline to provide a sophisticated yet user-friendly solution for pipeline parallelism. The core of its design revolves around minimal annotations that guide automatic transformations of the computational graph.

A typical JAX training step involves defining a loss function, using JAX's grad transformation for automatic differentiation to compute gradients, and then applying these gradients to update model weights using an optimizer. This entire train_step is often compiled using jax.jit and executed on a device mesh for distributed execution, which handles data and tensor parallelism.

Jack's PP extends this paradigm by introducing two key primitives: pipeline_yields and accumulate_grads, alongside a mechanism for defining a Multi-Program, Multiple Data (MPMD) execution environment.

Defining Pipeline Stages with pipeline_yields

The pipeline_yields primitive is central to logically partitioning the model. Instead of requiring explicit stage definitions or manual layer grouping, pipeline_yields delineates stage boundaries based purely on data dependencies. For example, if a computation X = f(input) is followed by Y = g(X) and pipeline_yields is placed between them, it implicitly defines a stage boundary. The critical insight is that pipeline_yields is differentiable, meaning that corresponding pipeline_yields are automatically generated for the backward pass, ensuring that the gradient computation also adheres to the defined pipeline structure.

Consider an example where Z is defined in stage zero, used to define X (also in stage zero), and then Y is defined in stage one, depending on X. If pipeline_yields separates the definitions of X and Y, Z is implicitly part of stage zero because it's first used there. If Y later depends on Z directly, even if Z is not explicitly mentioned in the pipeline_yields for stage one, Jack's PP infers this data dependency. If stage zero and stage one are placed on different MPMD ranks by the chosen schedule, then passing Z from stage zero to stage one automatically triggers the insertion of send and receive operations. If they are on the same rank, Z is simply passed locally. This automatic inference of communication based on data flow is a significant reduction in user burden.

This dependency-based stage definition handles complex scenarios such as:

• Non-linear Data Flow: Stages are not restricted to simple sequential layers; branches and merges are naturally supported.
• Residual Connections: If an activation from an early stage is needed in a later, non-adjacent stage, Jack's PP infers the necessary communication.
• Shared Weights: If a weight (e.g., tied embeddings) is used across multiple stages, its placement and synchronization are handled automatically.

Specifying Schedules and Gradient Accumulation with accumulate_grads

The accumulate_grads primitive is a higher-order reduction function that encapsulates the logic for gradient computation and accumulation over microbatches, based on a user-specified pipeline schedule. Instead of hardcoding a schedule, users pass a schedule object (e.g., representing 1F1B or interleaved 1F1B) to accumulate_grads. This primitive then orchestrates the execution of the forward and backward passes for multiple microbatches, managing gradient accumulation and communication according to the chosen schedule.

The schedule itself is encoded, often as a sequence of (stage_id, microbatch_id, direction) tuples, allowing for fine-grained control over the execution order and mapping of tasks to MPMD ranks. This flexibility is critical for leveraging schedules like 1F1B, which significantly reduce peak activation memory by interleaving forward and backward passes, thereby requiring less rematerialization and improving throughput.

MPMD Execution and Automatic Placement

Jack's PP operates within an MPMD (Multiple Program, Multiple Data) environment. Users define a remote_mesh that specifies the number of MPMD ranks and the local device mesh (e.g., 2x2 for 4 GPUs) allocated to each rank. This means different programs (sub-JAXprs) can run on different ranks, addressing the non-SPMD nature of pipeline parallelism.

A crucial aspect of Jack's PP is its ability to automatically infer the placement of auxiliary computations that occur outside the main gradient accumulation loop. Computations before the loop (e.g., input preprocessing X) and after the loop (e.g., Y, gradient application, optimizer state updates) are analyzed based on their data dependencies. Jack's PP determines which MPMD rank should execute these computations, ensuring that all necessary inputs are available and outputs are correctly passed. For instance, if an initial computation X only feeds into stage zero, it will be placed on MPMD rank zero. Similarly, final computations that depend on outputs from the last stage will be placed on the corresponding rank.

Runtime Architecture

Jack's PP employs a single controller runtime architecture to manage the distributed execution:

1. Driver Process: A central driver process initiates the training. It obtains the JAX internal representation (JAXpr) of the entire train_step.
2. Transformation and Sub-JAXpr Generation: Jack's PP applies its transformations to this JAXpr. This involves:

• Analyzing pipeline_yields to define logical stages.
• Inferring all data dependencies across stages.
• Generating corresponding pipeline_yields for the backward pass.
• Determining communication points (send/receive) for inter-stage data flow.
• Identifying and placing pre/post-loop computations.
• Based on the specified schedule in accumulate_grads, the driver partitions the overall JAXpr into multiple sub-JAXprs, one for each MPMD rank, representing its local computation.

1. Lowering to Stable HLO: These sub-JAXprs are then lowered to Stable HLO (High-Level Optimizer) modules. Stable HLO is a portable, stable intermediate representation for ML computations, ensuring compatibility and efficient compilation.
2. Serialization and Dispatch: The Stable HLO modules are serialized and sent, along with execution instructions (e.g., execution order, communication patterns), to remote actors.
3. Remote Actor Execution: Each remote actor receives its specific Stable HLO module and instructions. It performs the final compilation using the vanilla XLA compiler (no modifications to XLA are needed) and executes the compiled modules, including handling the inferred send and receive operations.

The entire process happens transparently to the user, who interacts with Jack's PP through a few high-level primitives, oblivious to the underlying distributed dispatch and compilation. The talk also briefly mentions a recently added multi-controller runtime, indicating ongoing development to further enhance scalability and robustness, though it was not the focus of this particular presentation.

By meticulously handling the graph transformations, communication inference, and runtime orchestration, Jack's PP effectively bridges the gap between JAX's functional programming model and the demanding requirements of general pipeline parallelism, making it a powerful tool for large-scale deep learning.

Experimental Setup & Results

▶ Watch: Automatic data dependency handling with pipeline_yields (7:00)

The effectiveness of Jack's PP was rigorously evaluated through benchmarks primarily focused on training large language models, specifically the GPT-3 175B model. The experiments aimed to quantify Jack's PP's performance against existing pipeline parallelism solutions in JAX and state-of-the-art frameworks, as well as its scaling efficiency.

Comparison with SPMD Pipeline Parallelism in JAX

A significant benchmark involved comparing Jack's PP against the constrained form of pipeline parallelism offered by SPMD within JAX. This SPMD approach typically supports only the G-pipe schedule.

• SPMD's Limitations: The G-pipe schedule, while a foundational pipeline parallelism technique, has inherent limitations. It often leads to higher peak activation memory requirements because backward passes are delayed, necessitating more rematerialization (recomputing activations during backward pass to save memory). This adds computational overhead. Furthermore, SPMD's design might not fully optimize the overlap of point-to-point communication with computation.
• Jack's PP's Advantages: For this benchmark, Jack's PP leveraged the interleaved 1F1B schedule. This schedule is known to significantly reduce the pipeline bubble (idle time) and, more importantly, drastically lower activation memory demands by scheduling backward computations earlier. This reduces the need for costly rematerialization. Jack's PP's runtime is also designed to better overlap point-to-point communication with actual computation, further boosting efficiency.
• Result: Jack's PP demonstrated a substantial performance advantage, being approximately 44% faster than the SPMD pipeline parallelism in JAX for the GPT-3 175B training step. This speedup is attributed to its ability to use more memory-efficient and throughput-optimized schedules, coupled with superior communication-computation overlap, all while offering greater expressiveness.

Comparison with State-of-the-Art Nemo

To assess its competitiveness against highly optimized, domain-specific frameworks, Jack's PP was compared with Nemo, Nvidia's toolkit for conversational AI, which includes highly optimized implementations for large language model training.

• Nemo's Strengths: Nemo's model definitions often involve extensively optimized custom kernels and meticulously orchestrated communication patterns, hand-tuned for specific architectures and hardware.
• Jack's PP's Strengths: In contrast, Jack's PP is designed to be entirely model agnostic. It provides a general solution that works across arbitrary JAX computations without requiring custom kernel development or manual communication orchestration.
• Result: Despite being a general-purpose, model-agnostic solution, Jack's PP achieved performance that was remarkably close to Nemo, falling within 5-6% of Nemo's throughput. This finding is particularly impressive, highlighting that Jack's PP can deliver near state-of-the-art performance without the extensive engineering effort typically associated with highly specialized frameworks.

Scaling Efficiency

The ability to scale efficiently across a large number of devices is paramount for training truly massive models. Jack's PP was tested for its scaling efficiency on the GPT-3 175B model.

• Setup: The evaluation involved scaling the training from 128 GPUs to 1024 GPUs. While the specific GPU types were not explicitly mentioned in the talk, the context of Nvidia and large-scale training implies the use of high-performance Nvidia GPUs (e.g., A100 or H100).
• Result: Jack's PP achieved an exceptional 98% scaling efficiency when moving from 128 to 1024 GPUs. This indicates that the overhead introduced by Jack's PP's distributed orchestration and communication mechanisms remains very low even as the system scales to a significant number of accelerators, affirming its suitability for extreme-scale deep learning training. More detailed numbers regarding specific throughputs or hardware configurations are available in the associated paper.

In summary, the experimental results unequivocally demonstrate that Jack's PP not only simplifies the implementation of complex pipeline parallelism but also delivers highly competitive performance and excellent scalability, making it a compelling solution for the JAX community.

Practical Implications

▶ Watch: Differentiability and composability with JAX auto-differentiation (8:00)

Jack's PP represents a significant leap forward in making advanced distributed deep learning techniques more accessible and efficient for a broad range of practitioners. Its practical implications span model builders, infrastructure teams, and anyone involved in deploying large-scale AI systems.

For Practitioners and Model Builders

• Simplified Pipeline Parallelism: The most direct benefit is the dramatic simplification of implementing pipeline parallelism. Model builders no longer need to manually split models, insert send/receive calls, or painstakingly manage communication across stages. The pipeline_yields and accumulate_grads primitives abstract away these complexities, allowing them to focus on model design rather than distributed systems engineering.
• Flexibility with Schedules: The ability to choose and easily switch between different pipeline schedules (e.g., 1F1B, interleaved 1F1B) is a powerful tool. Practitioners can experiment with schedules that best balance memory consumption and throughput for their specific model and hardware, without rewriting large portions of their training code. This is a stark contrast to systems that might be locked into less optimal schedules like G-pipe.
• Support for Complex Models: Jack's PP's ability to handle arbitrary computations, non-linear data flows, residual connections, and shared weights makes it suitable for a wider range of modern deep learning architectures. This removes a significant barrier for researchers and developers working on novel models that might not conform to simple layer-sequential structures.
• Leveraging JAX's Ecosystem: For those already using JAX, Jack's PP seamlessly integrates with existing features like jax.jit for compilation and jax.grad for automatic differentiation. This means practitioners can continue to enjoy JAX's functional programming benefits and high-performance compilation without needing to learn an entirely new distributed framework.

For Infrastructure Teams and Deployers

• Robust Distributed Training: Jack's PP provides a robust and well-engineered solution for distributed training, reducing the likelihood of hard-to-debug distributed bugs stemming from manual communication errors. This improves the stability and reliability of large-scale training jobs.
• Efficient Resource Utilization: By enabling the use of more efficient schedules (like 1F1B) and optimizing communication-computation overlap, Jack's PP helps infrastructure teams maximize the utilization of expensive GPU resources. The demonstrated 98% scaling efficiency to 1024 GPUs underscores its capability for efficient resource scaling.
• Vanilla XLA Compatibility: The fact that Jack's PP operates as a pure Python package on top of JAX and utilizes vanilla XLA for compilation simplifies deployment and maintenance. Infrastructure teams don't need to manage custom XLA builds or patched compilers, reducing operational overhead.
• Model Agnostic Deployment: Its model-agnostic nature means that infrastructure can support a diverse range of models and training workloads using a single, consistent distributed training framework, rather than maintaining multiple specialized solutions.

Tradeoffs and Limitations

While highly beneficial, it's important to consider potential tradeoffs and limitations:

• JAX Specificity: Jack's PP is currently a JAX-specific solution. Organizations primarily using other frameworks (e.g., PyTorch, TensorFlow) would not directly benefit from it.
• Learning Curve for Primitives: While minimal, understanding how pipeline_yields and accumulate_grads interact with JAX's functional transformations and data dependencies still requires some initial learning. Misplacement of pipeline_yields could lead to sub-optimal partitioning or errors.
• Optimization Ceiling: While competitive with highly optimized frameworks like Nemo, there might still be a marginal performance gap. For extremely niche, hand-tuned models with highly specialized hardware and custom kernels, a bespoke solution might eke out a tiny bit more performance, though at a significantly higher development and maintenance cost. Jack's PP aims for "general-purpose, near-optimal" rather than "hyper-specialized, absolute optimal."
• Debugging Distributed Systems: While Jack's PP automates many aspects, debugging performance issues or logical errors in large-scale distributed training remains inherently complex. Tools and techniques for profiling and debugging JAX-based distributed systems will continue to be crucial.

In conclusion, Jack's PP significantly lowers the barrier to entry for highly efficient and scalable pipeline parallelism in JAX. It empowers practitioners to build and train larger, more complex models with greater ease, while providing infrastructure teams with a robust and performant framework for managing distributed AI workloads.

Key Takeaways

• Jack's PP simplifies pipeline parallelism: It addresses the historically complex and error-prone nature of pipeline parallelism in deep learning by providing a general-purpose solution within JAX.
• Minimal user annotation: Users define pipeline stages and schedules using only two core primitives, pipeline_yields and accumulate_grads, reducing manual orchestration significantly.
• Automatic inference of crucial details: The framework automatically infers inter-stage communication (send/receive operations), placement of auxiliary computations, and generates backward pass logic based on data dependencies.
• Superior performance and flexibility: Jack's PP is 44% faster than JAX's SPMD pipeline parallelism (G-pipe) and within 5-6% of highly optimized, model-specific frameworks like Nemo, while supporting diverse and efficient schedules (e.g., 1F1B, interleaved 1F1B).
• Exceptional scaling efficiency: It demonstrates 98% scaling efficiency when training GPT-3 175B from 128 GPUs to 1024 GPUs, proving its capability for extreme-scale distributed training.
• Open-sourced and JAX-native: Implemented as a pure Python package on top of vanilla JAX and XLA, Jack's PP is open-sourced, fostering community adoption and seamless integration into existing JAX workflows.

About the Speaker(s)

The primary presenter for this talk was Jong Lee, affiliated with Nvidia. He introduced Jack's PP as a novel solution for scaling deep learning training with MPMD pipeline parallelism. The co-authors of the paper, also acknowledged during the presentation, include Anxhelo Xhebraj, Hanfeng Chen, and Vinod Grover. All co-authors were also present at the MLSys 2025 conference, implying their collective contribution to this research, likely as part of the Nvidia team.