# Automap: Towards Ergonomic Automated Parallelism for ML Models

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 37 references that are in the map.
- First posted on arXiv 6 Dec 2021.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [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)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Interaction Networks for Learning about Objects, Relations and Physics](https://mlsanity.com/paper/1612.00222/) (2016)
- [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)
- [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)
- [Device Placement Optimization with Reinforcement Learning](https://mlsanity.com/paper/1706.04972/) (2017)
- [PipeMare: Asynchronous Pipeline Parallel DNN Training](https://mlsanity.com/paper/1910.05124/) (2019)
- [Supporting Very Large Models using Automatic Dataflow Graph Partitioning](https://mlsanity.com/paper/1807.08887/) (2018)
- [Transferable Graph Optimizers for ML Compilers](https://mlsanity.com/paper/2010.12438/) (2020)
- [Equality Saturation for Tensor Graph Superoptimization](https://mlsanity.com/paper/2101.01332/) (2021)
- [Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs](https://mlsanity.com/paper/1905.02494/) (2019)

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

- [Respecting causality is all you need for training physics-informed neural networks](https://mlsanity.com/paper/2203.07404/) (2022)
- [Pre-training via Denoising for Molecular Property Prediction](https://mlsanity.com/paper/2206.00133/) (2022)
- [Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Traces](https://mlsanity.com/paper/2305.14516/) (2023)
- [PartIR: Composing SPMD Partitioning Strategies for Machine Learning](https://mlsanity.com/paper/2401.11202/) (2024)
- [Learned Force Fields Are Ready For Ground State Catalyst Discovery](https://mlsanity.com/paper/2209.12466/) (2022)
- [AutoParLLM: GNN-guided Context Generation for Zero-Shot Code Parallelization using LLMs](https://mlsanity.com/paper/2310.04047/) (2023)
- [Pex: Memory-efficient Microcontroller Deep Learning through Partial Execution](https://mlsanity.com/paper/2211.17246/) (2022)
- [Automatic Discovery of Composite SPMD Partitioning Strategies in PartIR](https://mlsanity.com/paper/2210.06352/) (2022)
- [UniAP: Unifying Inter- and Intra-Layer Automatic Parallelism by Mixed Integer Quadratic Programming](https://mlsanity.com/paper/2307.16375/) (2023)
- [HETHUB: A Distributed Training System with Heterogeneous Cluster for Large-Scale Models](https://mlsanity.com/paper/2405.16256/) (2024)
- [TOAST: Fast and scalable auto-partitioning based on principled static analysis](https://mlsanity.com/paper/2508.15010/) (2025)
- [MLCommons Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Traces](https://mlsanity.com/paper/2605.11333/) (2026)

## Abstract (arXiv)

The rapid rise in demand for training large neural network architectures has brought into focus the need for partitioning strategies, for example by using data, model, or pipeline parallelism. Implementing these methods is increasingly supported through program primitives, but identifying efficient partitioning strategies requires expensive experimentation and expertise. We present the prototype of an automated partitioner that seamlessly integrates into existing compilers and existing user workflows. Our partitioner enables SPMD-style parallelism that encompasses data parallelism and parameter/activation sharding. Through a combination of inductive tactics and search in a platform-independent partitioning IR, automap can recover expert partitioning strategies such as Megatron sharding for transformer layers.