Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

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%).

References in corpus (21)

• Adam: A Method for Stochastic Optimization
• Decoupled Weight Decay Regularization
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• XLNet: Generalized Autoregressive Pretraining for Language Understanding
• Mixed Precision Training
• SQuAD: 100,000+ Questions for Machine Comprehension of Text
• Training Deep Nets with Sublinear Memory Cost
• Large Batch Training of Convolutional Networks
• Pointer Sentinel Mixture Models
• Learning to Generate Reviews and Discovering Sentiment
• A Simple Method for Commonsense Reasoning
• Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books
• Multi-Task Deep Neural Networks for Natural Language Understanding
• Large Batch Optimization for Deep Learning: Training BERT in 76 minutes
• Beyond Data and Model Parallelism for Deep Neural Networks
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training
• Defending Against Neural Fake News
• Efficient and Robust Parallel DNN Training through Model Parallelism on Multi-GPU Platform
• Generalization through Memorization: Nearest Neighbor Language Models
• Dynamic Evaluation of Transformer Language Models
• Unsupervised Pretraining for Sequence to Sequence Learning

Cited by in corpus (209)

• HuggingFace's Transformers: State-of-the-art Natural Language Processing
• Language Models are Few-Shot Learners
• LoRA: Low-Rank Adaptation of Large Language Models
• On the Opportunities and Risks of Foundation Models
• Pre-trained Models for Natural Language Processing: A Survey
• ALBERT: A Lite BERT for Self-supervised Learning of Language Representations
• The Pile: An 800GB Dataset of Diverse Text for Language Modeling
• Beyond English-Centric Multilingual Machine Translation
• DeBERTa: Decoding-enhanced BERT with Disentangled Attention
• DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharing
• Self-Supervised Speech Representation Learning: A Review
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• Large language models in medicine: the potentials and pitfalls
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• Compute Trends Across Three Eras of Machine Learning
• Extracting Training Data from Large Language Models
• Compute and Energy Consumption Trends in Deep Learning Inference
• Controllable Protein Design with Language Models
• ERNIE 3.0: Large-scale Knowledge Enhanced Pre-training for Language Understanding and Generation
• BigSSL: Exploring the Frontier of Large-Scale Semi-Supervised Learning for Automatic Speech Recognition
• MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer
• Scientometric Review of Artificial Intelligence for Operations & Maintenance of Wind Turbines: The Past, Present and Future
• COCO-LM: Correcting and Contrasting Text Sequences for Language Model Pretraining
• Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines
• SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification
• DEPTWEET: A Typology for Social Media Texts to Detect Depression Severities
• PanGu-$\alpha$: Large-scale Autoregressive Pretrained Chinese Language Models with Auto-parallel Computation
• Adversarial Training for Large Neural Language Models
• Multimodal Few-Shot Learning with Frozen Language Models
• Language Models are Open Knowledge Graphs
• ASTRA-sim2.0: Modeling Hierarchical Networks and Disaggregated Systems for Large-model Training at Scale
• (Security) Assertions by Large Language Models
• Understanding the Potential of FPGA-Based Spatial Acceleration for Large Language Model Inference
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• ProtTrans: Towards Cracking the Language of Life's Code Through Self-Supervised Deep Learning and High Performance Computing
• Merak: An Efficient Distributed DNN Training Framework with Automated 3D Parallelism for Giant Foundation Models
• Alexa Teacher Model: Pretraining and Distilling Multi-Billion-Parameter Encoders for Natural Language Understanding Systems
• BASE Layers: Simplifying Training of Large, Sparse Models
• Memory-Efficient Pipeline-Parallel DNN Training
• ZeRO-Offload: Democratizing Billion-Scale Model Training
• Scale Efficiently: Insights from Pre-training and Fine-tuning Transformers
• ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
• Model-Parallel Fourier Neural Operators as Learned Surrogates for Large-Scale Parametric PDEs
• Generalization through Memorization: Nearest Neighbor Language Models
• Cross-lingual Zero- and Few-shot Hate Speech Detection Utilising Frozen Transformer Language Models and AXEL
• Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers
• Compressive Transformers for Long-Range Sequence Modelling
• Machine Reading Comprehension: The Role of Contextualized Language Models and Beyond
• M6: A Chinese Multimodal Pretrainer
• DeLighT: Deep and Light-weight Transformer
• TensorOpt: Exploring the Tradeoffs in Distributed DNN Training with Auto-Parallelism
• BERT Loses Patience: Fast and Robust Inference with Early Exit
• Knowledge-Aware Language Model Pretraining
• Enabling Compute-Communication Overlap in Distributed Deep Learning Training Platforms
• FastMoE: A Fast Mixture-of-Expert Training System
• Achieving Peak Performance for Large Language Models: A Systematic Review
• GSPMD: General and Scalable Parallelization for ML Computation Graphs
• PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management
• Multi-Graph Transformer for Free-Hand Sketch Recognition
• OneFlow: Redesign the Distributed Deep Learning Framework from Scratch
• Scalable and Efficient MoE Training for Multitask Multilingual Models
• Hanayo: Harnessing Wave-like Pipeline Parallelism for Enhanced Large Model Training Efficiency
• ProGen: Language Modeling for Protein Generation
• TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models
• Tesseract: Parallelize the Tensor Parallelism Efficiently
• Controlling Style in Generated Dialogue
• An Efficient Statistical-based Gradient Compression Technique for Distributed Training Systems
• Mixed Dimension Embeddings with Application to Memory-Efficient Recommendation Systems
• Accelerating Training of Transformer-Based Language Models with Progressive Layer Dropping
• Large language models for automated scholarly paper review: A survey
• PLATO-2: Towards Building an Open-Domain Chatbot via Curriculum Learning
• Laplace Redux -- Effortless Bayesian Deep Learning
• How fine can fine-tuning be? Learning efficient language models
• Efficient Quantized Sparse Matrix Operations on Tensor Cores
• Yuan 1.0: Large-Scale Pre-trained Language Model in Zero-Shot and Few-Shot Learning
• Pollux: Co-adaptive Cluster Scheduling for Goodput-Optimized Deep Learning
• Data Movement Is All You Need: A Case Study on Optimizing Transformers
• Parameter-Efficient Transfer Learning with Diff Pruning
• Neural Language Generation: Formulation, Methods, and Evaluation
• Reducing Sentiment Bias in Language Models via Counterfactual Evaluation
• Modeling Protein Using Large-scale Pretrain Language Model
• Towards Efficient Post-training Quantization of Pre-trained Language Models
• Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training
• Maximizing Parallelism in Distributed Training for Huge Neural Networks
• RealFormer: Transformer Likes Residual Attention
• TokenPose: Learning Keypoint Tokens for Human Pose Estimation
• Rethinking gradient sparsification as total error minimization
• Efficient Pipeline Planning for Expedited Distributed DNN Training
• Bandwidth-Effective DRAM Cache for GPUs with Storage-Class Memory
• The Stability-Efficiency Dilemma: Investigating Sequence Length Warmup for Training GPT Models
• M6-10T: A Sharing-Delinking Paradigm for Efficient Multi-Trillion Parameter Pretraining
• EarlyBERT: Efficient BERT Training via Early-bird Lottery Tickets
• Towards Crowdsourced Training of Large Neural Networks using Decentralized Mixture-of-Experts
• Conformer-Kernel with Query Term Independence for Document Retrieval
• I-BERT: Integer-only BERT Quantization
• PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers
• A Survey of Text Representation Methods and Their Genealogy
• Pre-Trained Models: Past, Present and Future
• Explaining Bayesian Neural Networks
• M6-T: Exploring Sparse Expert Models and Beyond
• Stereotype and Skew: Quantifying Gender Bias in Pre-trained and Fine-tuned Language Models
• Style Example-Guided Text Generation using Generative Adversarial Transformers
• NetReduce: RDMA-Compatible In-Network Reduction for Distributed DNN Training Acceleration
• MEGATRON-CNTRL: Controllable Story Generation with External Knowledge Using Large-Scale Language Models
• Pipeline Parallelism for Inference on Heterogeneous Edge Computing
• NLP From Scratch Without Large-Scale Pretraining: A Simple and Efficient Framework
• Cascaded Text Generation with Markov Transformers
• BioMegatron: Larger Biomedical Domain Language Model
• MicroRec: Efficient Recommendation Inference by Hardware and Data Structure Solutions
• An Oracle for Guiding Large-Scale Model/Hybrid Parallel Training of Convolutional Neural Networks
• Rethinking embedding coupling in pre-trained language models
• Optimus: Organizing Sentences via Pre-trained Modeling of a Latent Space
• DRIVE: One-bit Distributed Mean Estimation
• ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training
• Hardware Beyond Backpropagation: a Photonic Co-Processor for Direct Feedback Alignment
• KroneckerBERT: Learning Kronecker Decomposition for Pre-trained Language Models via Knowledge Distillation
• MPipeMoE: Memory Efficient MoE for Pre-trained Models with Adaptive Pipeline Parallelism
• Pipelined Backpropagation at Scale: Training Large Models without Batches
• Length-Adaptive Transformer: Train Once with Length Drop, Use Anytime with Search
• Revisiting BFloat16 Training
• Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training
• BaPipe: Exploration of Balanced Pipeline Parallelism for DNN Training
• Breaking the Memory Wall: A Study of I/O Patterns and GPU Memory Utilization for Hybrid CPU-GPU Offloaded Optimizers
• Improving Machine Reading Comprehension with Single-choice Decision and Transfer Learning
• Activations and Gradients Compression for Model-Parallel Training
• EmbRace: Accelerating Sparse Communication for Distributed Training of NLP Neural Networks
• Counterfactual Variable Control for Robust and Interpretable Question Answering
• LIBRA: Enabling Workload-aware Multi-dimensional Network Topology Optimization for Distributed Training of Large AI Models
• LAMP: Large Deep Nets with Automated Model Parallelism for Image Segmentation
• Optimizing DNN Compilation for Distributed Training with Joint OP and Tensor Fusion
• Tearing Down the Memory Wall
• The Dark Side of the Language: Pre-trained Transformers in the DarkNet
• Contrastive Distillation on Intermediate Representations for Language Model Compression
• TDML -- A Trustworthy Distributed Machine Learning Framework
• Insights into resource utilization of code small language models serving with runtime engines and execution providers
• Conversations Are Not Flat: Modeling the Dynamic Information Flow across Dialogue Utterances
• Pay Attention when Required
• Efficient Transformer-based Large Scale Language Representations using Hardware-friendly Block Structured Pruning
• Solving Machine Learning Problems
• UFO-BLO: Unbiased First-Order Bilevel Optimization
• Direct Feedback Alignment Scales to Modern Deep Learning Tasks and Architectures
• Auto-MAP: A DQN Framework for Exploring Distributed Execution Plans for DNN Workloads
• SpeechMoE: Scaling to Large Acoustic Models with Dynamic Routing Mixture of Experts
• Learning from Videos with Deep Convolutional LSTM Networks
• Distributed Deep Learning in Open Collaborations
• Pelta: Shielding Transformers to Mitigate Evasion Attacks in Federated Learning
• 8-bit Optimizers via Block-wise Quantization
• What Changes Can Large-scale Language Models Bring? Intensive Study on HyperCLOVA: Billions-scale Korean Generative Pretrained Transformers
• Connectivity Matters: Neural Network Pruning Through the Lens of Effective Sparsity
• Distributed Training and Optimization Of Neural Networks
• A Closer Look at Codistillation for Distributed Training
• VirtualFlow: Decoupling Deep Learning Models from the Underlying Hardware
• Testing pre-trained Transformer models for Lithuanian news clustering
• Giving Up Control: Neurons as Reinforcement Learning Agents
• Trends and Advancements in Deep Neural Network Communication
• Undivided Attention: Are Intermediate Layers Necessary for BERT?
• Go Wider Instead of Deeper
• Is Attention All What You Need? -- An Empirical Investigation on Convolution-Based Active Memory and Self-Attention
• Moshpit SGD: Communication-Efficient Decentralized Training on Heterogeneous Unreliable Devices
• ERNIE-Tiny : A Progressive Distillation Framework for Pretrained Transformer Compression
• Layered gradient accumulation and modular pipeline parallelism: fast and efficient training of large language models
• Automatic Graph Partitioning for Very Large-scale Deep Learning
• Beyond Preserved Accuracy: Evaluating Loyalty and Robustness of BERT Compression
• An Empirical Exploration in Quality Filtering of Text Data
• Towards More Effective and Economic Sparsely-Activated Model
• Deeplite Neutrino: An End-to-End Framework for Constrained Deep Learning Model Optimization
• Directed Beam Search: Plug-and-Play Lexically Constrained Language Generation
• Multi-grained Evidence Inference for Multi-choice Reading Comprehension
• I-BERT: Inductive Generalization of Transformer to Arbitrary Context Lengths
• SWIFT: Expedited Failure Recovery for Large-scale DNN Training
• Varuna: Scalable, Low-cost Training of Massive Deep Learning Models
• Knowledgeable Dialogue Reading Comprehension on Key Turns
• Can We Achieve More with Less? Exploring Data Augmentation for Toxic Comment Classification
• Adaptive Semiparametric Language Models
• GPT3Mix: Leveraging Large-scale Language Models for Text Augmentation
• SentenceMIM: A Latent Variable Language Model
• Memory-efficient array redistribution through portable collective communication
• A Survey of Large-Scale Deep Learning Serving System Optimization: Challenges and Opportunities
• DistIR: An Intermediate Representation and Simulator for Efficient Neural Network Distribution
• Mini Minds: Exploring Bebeshka and Zlata Baby Models
• Local Knowledge Powered Conversational Agents
• AutoDrop: Training Deep Learning Models with Automatic Learning Rate Drop
• Defending Democracy: Using Deep Learning to Identify and Prevent Misinformation
• Emergent Properties of Finetuned Language Representation Models
• Privileged Zero-Shot AutoML
• Building Compact and Robust Deep Neural Networks with Toeplitz Matrices
• Scheduling Optimization Techniques for Neural Network Training
• Neural Parameter Allocation Search
• WaLDORf: Wasteless Language-model Distillation On Reading-comprehension
• The Frankfurt Latin Lexicon: From Morphological Expansion and Word Embeddings to SemioGraphs
• A Linear Algebraic Approach to Model Parallelism in Deep Learning
• The curious case of developmental BERTology: On sparsity, transfer learning, generalization and the brain
• Do Transformers Need Deep Long-Range Memory
• Scaling Distributed Deep Learning Workloads beyond the Memory Capacity with KARMA
• HetSeq: Distributed GPU Training on Heterogeneous Infrastructure
• Data-Informed Global Sparseness in Attention Mechanisms for Deep Neural Networks
• META${}^2$: Memory-efficient taxonomic classification and abundance estimation for metagenomics with deep learning
• Dispatcher: A Message-Passing Approach To Language Modelling
• Large Product Key Memory for Pretrained Language Models
• High Throughput Training of Deep Surrogates from Large Ensemble Runs
• Simple and Effective Input Reformulations for Translation
• Benchmarking down-scaled (not so large) pre-trained language models
• Secure Distributed Training at Scale
• Improving Layer-wise Adaptive Rate Methods using Trust Ratio Clipping
• E.T.: Entity-Transformers. Coreference augmented Neural Language Model for richer mention representations via Entity-Transformer blocks
• Task-adaptive Pre-training of Language Models with Word Embedding Regularization
• Hydra: A System for Large Multi-Model Deep Learning
• Error Detection in Large-Scale Natural Language Understanding Systems Using Transformer Models
• Investigating Alternatives to the Root Mean Square for Adaptive Gradient Methods