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

Abstract

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.

arXiv comment24 pages. Submitted for review

References in corpus (13)

• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
• Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks
• Language Models are Few-Shot Learners
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• REALM: Retrieval-Augmented Language Model Pre-Training
• Training Deep Nets with Sublinear Memory Cost
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• Efficient and Robust Parallel DNN Training through Model Parallelism on Multi-GPU Platform
• ZeRO-Offload: Democratizing Billion-Scale Model Training
• ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
• Mesh-TensorFlow: Deep Learning for Supercomputers
• GSPMD: General and Scalable Parallelization for ML Computation Graphs
• TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models

Cited by in corpus (3)

• Merak: An Efficient Distributed DNN Training Framework with Automated 3D Parallelism for Giant Foundation Models
• Hanayo: Harnessing Wave-like Pipeline Parallelism for Enhanced Large Model Training Efficiency
• Comparative Analysis of AWS Model Deployment Services