mlsanityData to 14 Oct 2026

Automap: Towards Ergonomic Automated Parallelism for ML Models

arXiv 2112.02958 ↗First posted 6 Dec 2021cs.LG cs.DCcited by 12

Not on the map. Why?

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.

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

19 of 37 references.

References · 19

  1. 2017Attention Is All You Needcited by 60,361
  2. 2020Language Models are Few-Shot Learnerscited by 32,419
  3. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
  4. 2021Evaluating Large Language Models Trained on Codecited by 8,320
  5. 2020Scaling Laws for Neural Language Modelscited by 6,586
  6. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,868
  7. 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,631
  8. 2017In-Datacenter Performance Analysis of a Tensor Processing Unitcited by 618
  9. 2016Interaction Networks for Learning about Objects, Relations and Physicscited by 509
  10. 2021ZeRO-Offload: Democratizing Billion-Scale Model Trainingcited by 341
  11. 2018Beyond Data and Model Parallelism for Deep Neural Networkscited by 179
  12. 2020Memory-Efficient Pipeline-Parallel DNN Trainingcited by 134
7 more
  1. 2021GSPMD: General and Scalable Parallelization for ML Computation Graphscited by 111
  2. 2017Device Placement Optimization with Reinforcement Learningcited by 86
  3. 2019PipeMare: Asynchronous Pipeline Parallel DNN Trainingcited by 55
  4. 2018Supporting Very Large Models using Automatic Dataflow Graph Partitioningcited by 43
  5. 2020Transferable Graph Optimizers for ML Compilerscited by 29
  6. 2021Equality Saturation for Tensor Graph Superoptimizationcited by 27
  7. 2019Reinforced Genetic Algorithm Learning for Optimizing Computation Graphscited by 17

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2022: 5; 2023: 3; 2024: 2; 2025: 1; 2026: 1.

12 papers.

  1. 2022Respecting causality is all you need for training physics-informed neural networkscited by 156
  2. 2022Pre-training via Denoising for Molecular Property Predictioncited by 124
  3. 2023Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Tracescited by 13
  4. 2024PartIR: Composing SPMD Partitioning Strategies for Machine Learningcited by 10
  5. 2022Learned Force Fields Are Ready For Ground State Catalyst Discoverycited by 8
  6. 2023AutoParLLM: GNN-guided Context Generation for Zero-Shot Code Parallelization using LLMscited by 5
  7. 2022Pex: Memory-efficient Microcontroller Deep Learning through Partial Executioncited by 3
  8. 2022Automatic Discovery of Composite SPMD Partitioning Strategies in PartIRcited by 2
  9. 2023UniAP: Unifying Inter- and Intra-Layer Automatic Parallelism by Mixed Integer Quadratic Programmingcited by 2
  10. 2024HETHUB: A Distributed Training System with Heterogeneous Cluster for Large-Scale Modelscited by 2
  11. 2025TOAST: Fast and scalable auto-partitioning based on principled static analysiscited by 0
  12. 2026MLCommons Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Tracescited by 0