[Paper Review] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

Mohammad Shoeybi, Mostofa Patwary+4 more|arXiv (Cornell University)|Sep 17, 2019

Topic Modeling46 references821 citations

TL;DR

Megatron-LM demonstrates intra-layer model parallelism to train transformer models with billions of parameters in PyTorch, achieving 8.3B parameters on 512 GPUs with strong scaling and state-of-the-art results on several NLP benchmarks.

ABSTRACT

Recent work in language modeling demonstrates that training large transformer models advances the state of the art in Natural Language Processing applications. However, very large models can be quite difficult to train due to memory constraints. In this work, we present our techniques for training very large transformer models and implement a simple, efficient intra-layer model parallel approach that enables training transformer models with billions of parameters. Our approach does not require a new compiler or library changes, is orthogonal and complimentary to pipeline model parallelism, and can be fully implemented with the insertion of a few communication operations in native PyTorch. We illustrate this approach by converging transformer based models up to 8.3 billion parameters using 512 GPUs. We sustain 15.1 PetaFLOPs across the entire application with 76% scaling efficiency when compared to a strong single GPU baseline that sustains 39 TeraFLOPs, which is 30% of peak FLOPs. To demonstrate that large language models can further advance the state of the art (SOTA), we train an 8.3 billion parameter transformer language model similar to GPT-2 and a 3.9 billion parameter model similar to BERT. We show that careful attention to the placement of layer normalization in BERT-like models is critical to achieving increased performance as the model size grows. Using the GPT-2 model we achieve SOTA results on the WikiText103 (10.8 compared to SOTA perplexity of 15.8) and LAMBADA (66.5% compared to SOTA accuracy of 63.2%) datasets. Our BERT model achieves SOTA results on the RACE dataset (90.9% compared to SOTA accuracy of 89.4%).

Motivation & Objective

• Motivate training of multi-billion parameter language models beyond single-GPU memory limits.
• Develop a simple, efficient intra-layer model parallel approach that fits in PyTorch with minimal changes.
• Evaluate scaling efficiency and performance on GPT-2 and BERT-like architectures.
• Demonstrate state-of-the-art results on language modeling and downstream tasks while releasing open-source code.

Proposed method

• Implement intra-layer model parallelism by partitioning GEMMs in MLP and self-attention across GPUs with minimal synchronization.
• Use two forward and two backward all-reduce communications per transformer layer to enable efficient scaling.
• Parallelize the input and output embedding matrices along vocabulary dimensions to reduce cross-GPU communication.
• Maintain duplicated layer normalization and residual computations on each GPU to avoid extra communication.
• Train mixed-precision Transformer models with dynamic loss scaling and activation checkpointing for memory efficiency.

Experimental results

Research questions

• RQ1 Can intra-layer model parallelism in PyTorch scale transformer models to multi-billion parameters without custom compilers?
• RQ2 How does model size affect GPT-2 and BERT-like model performance on standard NLP benchmarks when scaled to billions of parameters?
• RQ3 What architectural adjustments (e.g., layer normalization placement) are required to maintain or improve performance as models grow?
• RQ4 What are the practical scalability limits (FLOPs, throughput, efficiency) when training up to billions of parameters on hundreds of GPUs?

Key findings

• Trains up to 8.3B parameters on 512 GPUs using 8-way model parallelism with up to 15.1 PetaFLOPs sustained and 76% scaling efficiency versus a strong single-GPU baseline (39 TeraFLOPs).
• Model parallelism results in strong weak-scaling for both model-only and model+data configurations, achieving 74%–77% scaling efficiency for the largest setups.
• Careful placement of layer normalization in BERT-like models enables monotonic performance gains with increasing model size.
• GPT-2 style models reach state-of-the-art perplexity on WikiText103 (10.81) and strong LAMBADA accuracy (66.51%) with 8.3B parameters; RACE accuracy also improves to 90.9% with larger BERT/Megatron setups.
• BERT-style models up to 3.9B parameters achieve state-of-the-art development results on several GLUE-like tasks and RACE when using the proposed architecture and training regimen.