mlsanityData to 14 Oct 2026

Scaling Deep Learning Training with MPMD Pipeline Parallelism

arXiv 2412.14374 ↗First posted 18 Dec 2024cs.DC cs.LG cs.PLcited by 1

Not on the map. Why?

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.

Key facts. Cited by 1 paper 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

15 of 32 references.

References · 15

  1. 2024The Llama 3 Herd of Modelscited by 14,715
  2. 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 13,028
  3. 2022PaLM: Scaling Language Modeling with Pathwayscited by 5,326
  4. 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsitycited by 3,073
  5. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,868
  6. 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,631
  7. 2016Training Deep Nets with Sublinear Memory Costcited by 848
  8. 2022Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Modelcited by 559
  9. 2022Efficiently Scaling Transformer Inferencecited by 539
  10. 2014cuDNN: Efficient Primitives for Deep Learningcited by 245
  11. 2020The Hardware Lotterycited by 160
  12. 2021GSPMD: General and Scalable Parallelization for ML Computation Graphscited by 111
3 more
  1. 2024TorchTitan: One-stop PyTorch native solution for production ready LLM pre-trainingcited by 85
  2. 2024PartIR: Composing SPMD Partitioning Strategies for Machine Learningcited by 10
  3. 2022Breadth-First Pipeline Parallelismcited by 3

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 1.

1 paper.

  1. 2026Piper: A Programmable Distributed Training Systemcited by 0