Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism Mohammad Shoeybi 1 2 Mostofa Patwary 1 2 Raul Puri 1 2 Patrick LeGresley 2 Jared Casper 2 Bryan Catanzaro 2 Abstract 1. Introduction Recent work in language modeling demonstrates Natural Language Processing (NLP) is advancing quickly in that training large transformer models advances part due to an increase in available compute and dataset size. the state of the art in Natural Language Processing The abundance of compute and data enables training increas- applications. However, very large models can be ingly larger language models via unsupervised pretraining quite difficult to train due to memory constraints. (Devlin et al., 2018; Radford et al., 2019). Empirical evi- In this work, we present our techniques for train- dence indicates that larger language models are dramatically ing very large transformer models and implement more useful for NLP tasks such as article completion, ques- a simple, efficient intra-layer model parallel ap- tion answering, and natural language inference (Lan et al., proach that enables training transformer models 2019; Raffel et al., 2019). By finetuning these pretrained with billions of parameters. Our approach does language models on downstream natural language tasks, not require a new compiler or library changes, is one can achieve state of the art results as shown in recent orthogonal and complimentary to pipeline model work (Devlin et al., 2018; Peters et al., 2018; Howard & parallelism, and can be fully implemented with Ruder, 2018; Radford et al., 2018; 2017; Ramachandran the insertion of a few communication operations et al., 2016; Liu et al., 2019b; Dai et al., 2019; Yang et al., in native PyTorch. We illustrate this approach 2019; Liu et al., 2019a; Lan et al., 2019). by converging transformer based models up to As these models become larger, they exceed the memory 8.3 billion parameters using 512 GPUs. We sus- limit of modern processors, and require additional memory tain 15.1 PetaFLOPs across the entire applica- management techniques such as activation checkpointing tion with 76% scaling efficiency when compared (Chen et al., 2016). Widely used optimization algorithms to a strong single GPU baseline that sustains 39 such as ADAM require additional memory per parameter to TeraFLOPs, which is 30% of peak FLOPs. To store momentum and other optimizer state, which reduces demonstrate that large language models can fur- the size of models that can be effectively trained. Several ther advance the state of the art (SOTA), we train approaches to model parallelism overcome this limit by an 8.3 billion parameter transformer language partitioning the model such that the weights and their asso- model similar to GPT-2 and a 3.9 billion parame- ciated optimizer state do not need to reside concurrently on ter model similar to BERT. We show that careful the processor. For example, GPipe (Huang et al., 2018) and attention to the placement of layer normalization Mesh-Tensorflow (Shazeer et al., 2018) provide frameworks arXiv:1909.08053v4 [cs.CL] 13 Mar 2020 in BERT-like models is critical to achieving in- for model parallelism of different kinds. However, they creased performance as the model size grows. Us- require rewriting the model, and rely on custom compilers ing the GPT-2 model we achieve SOTA results and frameworks that are still under development. on the WikiText103 (10.8 compared to SOTA per- plexity of 15.8) and LAMBADA (66.5% com- In this work, we implement a simple and efficient model pared to SOTA accuracy of 63.2%) datasets. Our parallel approach using intra-layer model-parallelism. We BERT model achieves SOTA results on the RACE exploit the inherent structure in transformer based language dataset (90.9% compared to SOTA accuracy of models to make a simple model-parallel implementation that 89.4%). trains efficiently in PyTorch, with no custom C++ code or compiler required. This approach is orthogonal to pipeline- based model parallelism as advocated by approaches such 1Equal contribution 2NVIDIA. Correspondence to: Mohammad as GPipe (Huang et al., 2018). Shoeybi <[email protected]>. To demonstrate the scalability of our approach, we establish Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism • We show that careful attention to the placement of layer normalization in BERT-like models is critical to achieving increased accuracies as the model grows. • We demonstrate that scaling the model size results in improved accuracies for both GPT-2 (studied up to 8.3 billion parameters) and BERT (studied up to 3.9B parameters) models. • We showcase that our models achieve state of the art results on test sets: perplexity on WikiText103 (10.8 ppl), accuracy on LAMBADA (66.5%), and accuracy on RACE (90.9%). Figure 1. Model (blue) and model+data (green) parallel FLOPS as a function of number of GPUs. Model parallel (blue): up to • We open source our code along with the training 8-way model parallel weak scaling with approximately 1 billion and evaluation pipelines at https://github:com/ parameters per GPU (e.g. 2 billion for 2 GPUs and 4 billion for NVIDIA/Megatron-LM 4 GPUs). Model+data parallel (green): similar configuration as model parallel combined with 64-way data parallel. 2. Background and Challenges 2.1. Neural Language Model Pretraining a baseline by training a model of 1.2 billion parameters on a single NVIDIA V100 32GB GPU, that sustains 39 Pretrained language models have become an indispensable TeraFLOPs. This is 30% of the theoretical peak FLOPS part of NLP researchers’ toolkits. Leveraging large corpus for a single GPU as configured in a DGX-2H server, and pretraining to learn robust neural representations of lan- is thus a strong baseline. Scaling the model to 8.3 billion guage is an active area of research that has spanned the parameters on 512 GPUs with 8-way model parallelism, past decade. Early examples of pretraining and transferring we achieve up to 15.1 PetaFLOPs per second sustained neural representations of language demonstrated that pre- over the entire application. This is 76% scaling efficiency trained word embedding tables improve downstream task compared to the single GPU case. Figure1 shows more results compared to word embedding tables learned from detailed scaling results. scratch (Mikolov et al., 2013; Pennington et al., 2014; Turian et al., 2010). Later work advanced research in this area by To analyze the effect of model size scaling on accuracy, learning and transferring neural models that capture contex- we train both left-to-right GPT-2 (Radford et al., 2019) lan- tual representations of words (Melamud et al., 2016; Mc- guage models as well as BERT (Devlin et al., 2018) bidi- Cann et al., 2017; Peters et al., 2018; Radford et al., 2017; rectional transformers and evaluate them on several down- 2019). Recent parallel work (Ramachandran et al., 2016; stream tasks. We show that the existing BERT architecture Howard & Ruder, 2018; Radford et al., 2018; Devlin et al., results in model degradation as the size increases. We over- 2018; Liu et al., 2019b; Dai et al., 2019; Yang et al., 2019; come this challenge by rearranging the layer normalization Liu et al., 2019a; Lan et al., 2019) further builds upon these and residual connection in the transformer layers and show ideas by not just transferring the language model to extract that with this change, results for the downstream tasks on contextual word representations, but by also finetuning the development sets improve monotonically as the model size language model in an end to end fashion on downstream increases. In addition, we show that our models achieve tasks. Through these works, the state of the art has advanced test set state of the art (SOTA) results on WikiText103, from transferring just word embedding tables to transferring cloze-style prediction accuracy on LAMBADA, and reading entire multi-billion parameter language models. This pro- comprehension RACE datasets. gression of methods has necessitated the need for hardware, In summary, our contributions are as follows: systems techniques, and frameworks that are able to oper- ate efficiently at scale and satisfy increasing computational needs. Our work aims to provide the tools necessary to take • We implement a simple and efficient model parallel another step forward in this trend. approach by making only a few targeted modifications to an existing PyTorch transformer implementation. 2.2. Transformer Language Models and Multi-Head Attention • We perform an in-depth empirical analysis of our model and data parallel technique and demonstrate Current work in NLP trends towards using transformer mod- up to 76% scaling efficiency using 512 GPUs. els (Vaswani et al., 2017) due to their superior accuracy Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism gate these effects and drive down the training time of large neural networks. To scale out training even further, parallel work (Chen et al., 2016) has combined data parallelism with activation checkpointing: recomputing activations in the backward pass without storing them in the forward pass to reduce memory requirements. However, these techniques have one fundamental limitation in the problem size they can tackle: the model must fit entirely on one worker. With language models of increasing size and complexity like BERT and GPT-2, neural networks have approached the memory capacity of modern hardware accelerators. One solution to this problem is to employ parameter sharing to reduce the memory footprint of the model (Lan et al., 2019), but this limits the overall capacity Figure 2. Transformer Architecture. Purple blocks correspond to of the model. Our approach is to utilize model parallelism fully connected layers. Each blue block represents a single trans- to split the model across multiple accelerators. This not former layer that is replicated N times.