# Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training

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

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 32 references that are in the map.
- First posted on arXiv 10 Nov 2021.

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

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (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)
- [Neural Collaborative Filtering](https://mlsanity.com/paper/1708.05031/) (2017)
- [Regularized Evolution for Image Classifier Architecture Search](https://mlsanity.com/paper/1802.01548/) (2018)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [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)
- [TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models](https://mlsanity.com/paper/2102.07988/) (2021)
- [Efficient and Robust Parallel DNN Training through Model Parallelism on Multi-GPU Platform](https://mlsanity.com/paper/1809.02839/) (2018)

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

- [BloombergGPT: A Large Language Model for Finance](https://mlsanity.com/paper/2303.17564/) (2023)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [Recurrent Drafter for Fast Speculative Decoding in Large Language Models](https://mlsanity.com/paper/2403.09919/) (2024)
- [CASSINI: Network-Aware Job Scheduling in Machine Learning Clusters](https://mlsanity.com/paper/2308.00852/) (2023)
- [Unicron: Economizing Self-Healing LLM Training at Scale](https://mlsanity.com/paper/2401.00134/) (2023)
- [Merak: An Efficient Distributed DNN Training Framework with Automated 3D Parallelism for Giant Foundation Models](https://mlsanity.com/paper/2206.04959/) (2022)
- [OLLA: Optimizing the Lifetime and Location of Arrays to Reduce the Memory Usage of Neural Networks](https://mlsanity.com/paper/2210.12924/) (2022)
- [Pipeline MoE: A Flexible MoE Implementation with Pipeline Parallelism](https://mlsanity.com/paper/2304.11414/) (2023)
- [Baechi: Fast Device Placement of Machine Learning Graphs](https://mlsanity.com/paper/2301.08695/) (2023)
- [Partitioning Distributed Compute Jobs with Reinforcement Learning and Graph Neural Networks](https://mlsanity.com/paper/2301.13799/) (2023)
- [Comparative Analysis of AWS Model Deployment Services](https://mlsanity.com/paper/2405.08175/) (2024)

## Abstract (arXiv)

With deep learning models rapidly growing in size, systems-level solutions for large-model training are required. We present Amazon SageMaker model parallelism, a software library that integrates with PyTorch, and enables easy training of large models using model parallelism and other memory-saving features. In contrast to existing solutions, the implementation of the SageMaker library is much more generic and flexible, in that it can automatically partition and run pipeline parallelism over arbitrary model architectures with minimal code change, and also offers a general and extensible framework for tensor parallelism, which supports a wider range of use cases, and is modular enough to be easily applied to new training scripts. The library also preserves the native PyTorch user experience to a much larger degree, supporting module re-use and dynamic graphs, while giving the user full control over the details of the training step. We evaluate performance over GPT-3, RoBERTa, BERT, and neural collaborative filtering, and demonstrate competitive performance over existing solutions.