# PartIR: Composing SPMD Partitioning Strategies for Machine Learning

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

## Key facts

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

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [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)
- [Scaling Language Models: Methods, Analysis & Insights from Training Gopher](https://mlsanity.com/paper/2112.11446/) (2021)
- [Learning to Simulate Complex Physics with Graph Networks](https://mlsanity.com/paper/2002.09405/) (2020)
- [Galactica: A Large Language Model for Science](https://mlsanity.com/paper/2211.09085/) (2022)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Efficiently Scaling Transformer Inference](https://mlsanity.com/paper/2211.05102/) (2022)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [TAPIR: Tracking Any Point with per-frame Initialization and temporal Refinement](https://mlsanity.com/paper/2306.08637/) (2023)
- [TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings](https://mlsanity.com/paper/2304.01433/) (2023)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [Memory-Efficient Pipeline-Parallel DNN Training](https://mlsanity.com/paper/2006.09503/) (2020)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training](https://mlsanity.com/paper/2004.13336/) (2020)
- [Measuring Progress in Fine-grained Vision-and-Language Understanding](https://mlsanity.com/paper/2305.07558/) (2023)
- [Stateful Dataflow Multigraphs: A Data-Centric Model for Performance Portability on Heterogeneous Architectures](https://mlsanity.com/paper/1902.10345/) (2019)
- [Automap: Towards Ergonomic Automated Parallelism for ML Models](https://mlsanity.com/paper/2112.02958/) (2021)
- [Getting to the Point. Index Sets and Parallelism-Preserving Autodiff for Pointful Array Programming](https://mlsanity.com/paper/2104.05372/) (2021)
- [High Resolution Medical Image Analysis with Spatial Partitioning](https://mlsanity.com/paper/1909.03108/) (2019)
- [DISTAL: The Distributed Tensor Algebra Compiler](https://mlsanity.com/paper/2203.08069/) (2022)
- [Weakly-Supervised Learning of Visual Relations in Multimodal Pretraining](https://mlsanity.com/paper/2305.14281/) (2023)

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

- [Scaling Deep Learning Computation over the Inter-Core Connected Intelligence Processor with T10](https://mlsanity.com/paper/2408.04808/) (2024)
- [Magellan: Autonomous Discovery of Novel Compiler Optimization Heuristics with AlphaEvolve](https://mlsanity.com/paper/2601.21096/) (2026)
- [Scalable Multi-Domain Adaptation of Language Models using Modular Experts](https://mlsanity.com/paper/2410.10181/) (2024)
- [Tilus: A Tile-Level GPGPU Programming Language for Low-Precision Computation](https://mlsanity.com/paper/2504.12984/) (2025)
- [Kareus: Joint Reduction of Dynamic and Static Energy in Large Model Training](https://mlsanity.com/paper/2601.17654/) (2026)
- [Compiling to linear neurons](https://mlsanity.com/paper/2511.13769/) (2025)
- [Scaling Deep Learning Training with MPMD Pipeline Parallelism](https://mlsanity.com/paper/2412.14374/) (2024)
- [Large-Area Fabrication-Aware Computational Diffractive Optics](https://mlsanity.com/paper/2505.22313/) (2025)
- [TOAST: Fast and scalable auto-partitioning based on principled static analysis](https://mlsanity.com/paper/2508.15010/) (2025)
- [AoiZora: Topology-Aware Auto-Parallel Optimization for Inference of Diffusion Transformers](https://mlsanity.com/paper/2606.17566/) (2026)

## 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..