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

By Can Bülent Karakus, R. Huilgol, Fei Wu, Anirudh Subramanian, Cade Daniel, D. Çavdar, Teng Xu, Haohan Chen, Arash Rahnama, L. Quintela. 2021

32 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Hardware and Architecture, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Hardware and Architecture

In Efficiency & Compression

Citations in the corpus

Reducing Activation Recomputation in Large Transformer Models (2022, methodology): “Along with layer-norm communication, Sagemaker [8] does four reduce-scatter communications per transformer layer in comparison to just two all-reduce communications in the Megatron-LM approach.”
BloombergGPT: A Large Language Model for Finance (2023, methodology): “We utilize the proprietary SageMaker Model Parallelism (SMP) library from AWS, which enables the automatic distribution of large models across multiple GPU devices and instances (Karakus et al., 2021).”
OLLA: Optimizing the Lifetime and Location of Arrays to Reduce the Memory Usage of Neural Networks (2022, background): “However, this approach, known as model parallelism [46], significantly increases the financial cost of training deep neural networks since it requires access to additional expensive compute accelerators and fast networks.”