# AMP: Automatically Finding Model Parallel Strategies with Heterogeneity Awareness

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

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 32 references that are in the map.
- First posted on arXiv 13 Oct 2022.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Flamingo: a Visual Language Model for Few-Shot Learning](https://mlsanity.com/paper/2204.14198/) (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)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Horovod: fast and easy distributed deep learning in TensorFlow](https://mlsanity.com/paper/1802.05799/) (2018)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [TransGAN: Two Pure Transformers Can Make One Strong GAN, and That Can Scale Up](https://mlsanity.com/paper/2102.07074/) (2021)
- [DAPPLE: A Pipelined Data Parallel Approach for Training Large Models](https://mlsanity.com/paper/2007.01045/) (2020)
- [Device Placement Optimization with Reinforcement Learning](https://mlsanity.com/paper/1706.04972/) (2017)
- [TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models](https://mlsanity.com/paper/2102.07988/) (2021)
- [Poseidon: An Efficient Communication Architecture for Distributed Deep Learning on GPU Clusters](https://mlsanity.com/paper/1706.03292/) (2017)
- [Distributed Balanced Partitioning via Linear Embedding](https://mlsanity.com/paper/1512.02727/) (2015)

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

- [Multi-Layer Transformers Gradient Can be Approximated in Almost Linear Time](https://mlsanity.com/paper/2408.13233/) (2024)
- [FlexSP: Accelerating Large Language Model Training via Flexible Sequence Parallelism](https://mlsanity.com/paper/2412.01523/) (2024)
- [Does compressing activations help model parallel training?](https://mlsanity.com/paper/2301.02654/) (2023)
- [Spindle: Efficient Distributed Training of Multi-Task Large Models via Wavefront Scheduling](https://mlsanity.com/paper/2409.03365/) (2024)
- [Mist: Efficient Distributed Training of Large Language Models via Memory-Parallelism Co-Optimization](https://mlsanity.com/paper/2503.19050/) (2025)
- [OctoPipe: Reducing Pipeline Bubbles for Heterogeneous Models via Co-Optimizing Partitioning, Placement, and Scheduling](https://mlsanity.com/paper/2509.23722/) (2025)
- [Pipette: Automatic Fine-grained Large Language Model Training Configurator for Real-World Clusters](https://mlsanity.com/paper/2405.18093/) (2024)
- [HetCCL: Accelerating LLM Training with Heterogeneous GPUs](https://mlsanity.com/paper/2601.22585/) (2026)
- [AutoScout: Structured Optimization for Automating ML System Configuration](https://mlsanity.com/paper/2603.11603/) (2026)
- [Rethinking Dynamic Networks and Heterogeneous Computing with Automatic Parallelization](https://mlsanity.com/paper/2506.02787/) (2025)
- [ELK: Exploring the Efficiency of Inter-core Connected AI Chips with Deep Learning Compiler Techniques](https://mlsanity.com/paper/2507.11506/) (2025)
- [Parallelism Strategy Chaining for Fast Training Convergence](https://mlsanity.com/paper/2609.07236/) (2026)
- [HAPMoE: Heterogeneity-Aware Automatic Parallelism Planning for Mixture-of-Experts Models Training](https://mlsanity.com/paper/2609.39350/) (2026)

## Abstract (arXiv)

Scaling up model sizes can lead to fundamentally new capabilities in many machine learning (ML) tasks. However, training big models requires strong distributed system expertise to carefully design model-parallel execution strategies that suit the model architectures and cluster setups. In this paper, we develop AMP, a framework that automatically derives such strategies. AMP identifies a valid space of model parallelism strategies and efficiently searches the space for high-performed strategies, by leveraging a cost model designed to capture the heterogeneity of the model and cluster specifications. Unlike existing methods, AMP is specifically tailored to support complex models composed of uneven layers and cluster setups with more heterogeneous accelerators and bandwidth. We evaluate AMP on popular models and cluster setups from public clouds and show that AMP returns parallel strategies that match the expert-tuned strategies on typical cluster setups. On heterogeneous clusters or models with heterogeneous architectures, AMP finds strategies with 1.54x and 1.77x higher throughput than state-of-the-art model-parallel systems, respectively.