mlsanityData to 7 Oct 2026

PartIR: Composing SPMD Partitioning Strategies for Machine Learning

arXiv 2401.11202 ↗First posted 20 Jan 2024cs.LG cs.DC cs.PLcited by 10

Not on the map. Why?

Abstract (arXiv)

Training of modern large neural networks (NN) requires a combination of parallelization strategies encompassing data, model, or optimizer sharding. When strategies increase in complexity, it becomes necessary for partitioning tools to be 1) expressive, allowing the composition of simpler strategies, and 2) predictable to estimate performance analytically. We present PartIR, our design for a NN partitioning system. PartIR is focused on an incremental approach to rewriting and is hardware-and-runtime agnostic. We present a simple but powerful API for composing sharding strategies and a simulator to validate them. The process is driven by high-level programmer-issued partitioning tactics, which can be both manual and automatic. Importantly, the tactics are specified separately from the model code, making them easy to change. We evaluate PartIR on several different models to demonstrate its predictability, expressibility, and ability to reach peak performance..

Key facts. Cited by 10 papers in the map (data through 7 Oct 2026). Builds on 32 of its 84 references that are in the map. First posted on arXiv 20 Jan 2024.

Builds on

32 of 84 references.

References · 32

  1. 2017Attention Is All You Needcited by 59,932
  2. 2014Adam: A Method for Stochastic Optimizationcited by 36,960
  3. 2020Language Models are Few-Shot Learnerscited by 32,199
  4. 2020Denoising Diffusion Probabilistic Modelscited by 20,851
  5. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
  6. 2023LLaMA: Open and Efficient Foundation Language Modelscited by 13,958
  7. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,829
  8. 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,601
  9. 2021Scaling Language Models: Methods, Analysis & Insights from Training Gophercited by 1,199
  10. 2020Learning to Simulate Complex Physics with Graph Networkscited by 867
  11. 2022Galactica: A Large Language Model for Sciencecited by 736
  12. 2023PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallelcited by 641
20 more
  1. 2017In-Datacenter Performance Analysis of a Tensor Processing Unitcited by 610
  2. 2022Efficiently Scaling Transformer Inferencecited by 529
  3. 2021ZeRO-Offload: Democratizing Billion-Scale Model Trainingcited by 339
  4. 2022Reducing Activation Recomputation in Large Transformer Modelscited by 300
  5. 2023TAPIR: Tracking Any Point with per-frame Initialization and temporal Refinementcited by 281
  6. 2023TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddingscited by 198
  7. 2018Beyond Data and Model Parallelism for Deep Neural Networkscited by 177
  8. 2018Mesh-TensorFlow: Deep Learning for Supercomputerscited by 168
  9. 2020Memory-Efficient Pipeline-Parallel DNN Trainingcited by 133
  10. 2021GSPMD: General and Scalable Parallelization for ML Computation Graphscited by 109
  11. 2020Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Trainingcited by 30
  12. 2023Measuring Progress in Fine-grained Vision-and-Language Understandingcited by 25
  13. 2019Stateful Dataflow Multigraphs: A Data-Centric Model for Performance Portability on Heterogeneous Architecturescited by 21
  14. 2021Automap: Towards Ergonomic Automated Parallelism for ML Modelscited by 12
  15. 2021Getting to the Point. Index Sets and Parallelism-Preserving Autodiff for Pointful Array Programmingcited by 10
  16. 2019High Resolution Medical Image Analysis with Spatial Partitioningcited by 9
  17. 2022DISTAL: The Distributed Tensor Algebra Compilercited by 6
  18. 2023Weakly-Supervised Learning of Visual Relations in Multimodal Pretrainingcited by 6
  19. 2020Auto-MAP: A DQN Framework for Exploring Distributed Execution Plans for DNN Workloadscited by 4
  20. 2022Automatic Discovery of Composite SPMD Partitioning Strategies in PartIRcited by 2

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 3; 2025: 4; 2026: 3.

10 papers.

  1. 2024Scaling Deep Learning Computation over the Inter-Core Connected Intelligence Processor with T10cited by 5
  2. 2026Magellan: Autonomous Discovery of Novel Compiler Optimization Heuristics with AlphaEvolvecited by 5
  3. 2024Scalable Multi-Domain Adaptation of Language Models using Modular Expertscited by 3
  4. 2025Tilus: A Tile-Level GPGPU Programming Language for Low-Precision Computationcited by 3
  5. 2026Kareus: Joint Reduction of Dynamic and Static Energy in Large Model Trainingcited by 3
  6. 2024Scaling Deep Learning Training with MPMD Pipeline Parallelismcited by 1
  7. 2025Compiling to linear neuronscited by 1
  8. 2025Large-Area Fabrication-Aware Computational Diffractive Opticscited by 0
  9. 2025TOAST: Fast and scalable auto-partitioning based on principled static analysiscited by 0
  10. 2026AoiZora: Topology-Aware Auto-Parallel Optimization for Inference of Diffusion Transformerscited by 0