# Automatic Discovery of Composite SPMD Partitioning Strategies in PartIR

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 38 references that are in the map.
- First posted on arXiv 7 Oct 2022.

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

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [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)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Pre-training via Denoising for Molecular Property Prediction](https://mlsanity.com/paper/2206.00133/) (2022)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [PipeMare: Asynchronous Pipeline Parallel DNN Training](https://mlsanity.com/paper/1910.05124/) (2019)
- [Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training](https://mlsanity.com/paper/2004.13336/) (2020)
- [Automap: Towards Ergonomic Automated Parallelism for ML Models](https://mlsanity.com/paper/2112.02958/) (2021)
- [Auto-MAP: A DQN Framework for Exploring Distributed Execution Plans for DNN Workloads](https://mlsanity.com/paper/2007.04069/) (2020)
- [A Transferable Approach for Partitioning Machine Learning Models on Multi-Chip-Modules](https://mlsanity.com/paper/2112.04041/) (2021)

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

- [PartIR: Composing SPMD Partitioning Strategies for Machine Learning](https://mlsanity.com/paper/2401.11202/) (2024)
- [TOAST: Fast and scalable auto-partitioning based on principled static analysis](https://mlsanity.com/paper/2508.15010/) (2025)

## Abstract (arXiv)

Large neural network models are commonly trained through a combination of advanced parallelism strategies in a single program, multiple data (SPMD) paradigm. For example, training large transformer models requires combining data, model, and pipeline partitioning; and optimizer sharding techniques. However, identifying efficient combinations for many model architectures and accelerator systems requires significant manual analysis. In this work, we present an automatic partitioner that identifies these combinations through a goal-oriented search. Our key findings are that a Monte Carlo Tree Search-based partitioner leveraging partition-specific compiler analysis directly into the search and guided goals matches expert-level strategies for various models.