# Scaling Deep Learning Training with MPMD Pipeline Parallelism

arXiv: https://arxiv.org/abs/2412.14374 · HTML: https://mlsanity.com/paper/2412.14374/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 32 references that are in the map.
- First posted on arXiv 18 Dec 2024.

## Builds on (most cited first, 15 of 15)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [Efficiently Scaling Transformer Inference](https://mlsanity.com/paper/2211.05102/) (2022)
- [cuDNN: Efficient Primitives for Deep Learning](https://mlsanity.com/paper/1410.0759/) (2014)
- [The Hardware Lottery](https://mlsanity.com/paper/2009.06489/) (2020)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training](https://mlsanity.com/paper/2410.06511/) (2024)
- [PartIR: Composing SPMD Partitioning Strategies for Machine Learning](https://mlsanity.com/paper/2401.11202/) (2024)
- [Breadth-First Pipeline Parallelism](https://mlsanity.com/paper/2211.05953/) (2022)

## Built on by (most cited first, 1 of 1)

- [Piper: A Programmable Distributed Training System](https://mlsanity.com/paper/2606.11169/) (2026)

## Abstract (arXiv)

We present JaxPP, a system for efficiently scaling the training of large deep learning models with flexible pipeline parallelism. We introduce a seamless programming model that allows implementing user-defined pipeline schedules for gradient accumulation. JaxPP automatically distributes tasks, corresponding to pipeline stages, over a cluster of nodes and automatically infers the communication among them. We implement a MPMD runtime for asynchronous execution of SPMD tasks. The pipeline parallelism implementation of JaxPP improves hardware utilization by up to 1.11× with respect to the best performing SPMD configuration.