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

The memory problem in training large transformers

Recent work has shown that training large transformer models advances the state of the art in natural language processing, but very large models can be difficult to train because of memory constraints ^[1]. During training a model's parameters, their gradients and the optimizer's state all have to be held somewhere, and a single accelerator has a fixed amount of memory. Once a model outgrows that budget, part of the work has to be spread over several devices.

The paper was written by Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper and Bryan Catanzaro of NVIDIA, and was submitted to arXiv on 17 September 2019 as 1909.08053 ^[1].

Splitting a layer across GPUs

The authors present a simple, efficient intra-layer model parallel approach that makes it possible to train transformer models with billions of parameters ^[1]. Model parallelism, in general, means dividing one model among several devices rather than replicating it on each of them: different pieces of the computation live on different chips, and the devices exchange the intermediate results they need from one another. In an intra-layer scheme the split happens inside a layer, so a single large matrix multiplication is itself shared out across devices.

What the authors emphasise is how little machinery the method requires. The approach does not need a new compiler or library changes, it is orthogonal and complementary to pipeline model parallelism, and it can be implemented by inserting a few communication operations in native PyTorch ^[1]. In other words, the change lives in the model code rather than in the toolchain, and it can be combined with other ways of splitting a model rather than replacing them.

What it achieved at scale

The authors illustrate the method by converging transformer models of up to 8.3 billion parameters using 512 GPUs ^[1]. They report that the system sustains 15.1 PetaFLOPs across the entire application, with 76 percent scaling efficiency compared with a strong single-GPU baseline that sustains 39 TeraFLOPs, or 30 percent of peak FLOPs ^[1]. Scaling efficiency here is a measure of how much of the single-device throughput survives when the same work is spread over many devices; the comparison is made against that strong single-GPU baseline rather than against the theoretical peak of the hardware.

Two large models and their benchmarks

To show that large language models can further advance the state of the art, the authors train an 8.3 billion parameter model similar to GPT-2 and a 3.9 billion parameter model similar to BERT ^[1]. They also report a structural finding: careful attention to the placement of layer normalization in BERT-like models is critical to achieving increased performance as model size grows ^[1]. Layer normalization is the step that rescales activations inside a transformer block; where it sits relative to the other sub-layers turns out to matter more as the model gets bigger.

The numbers are reported against previous state-of-the-art results. With the GPT-2-like model the paper reports a WikiText103 perplexity of 10.8 against a previous state of the art of 15.8, and LAMBADA accuracy of 66.5 percent against 63.2 percent; the BERT-like model reaches 90.9 percent on RACE against 89.4 percent ^[1]. Perplexity measures how surprised a language model is by held-out text, so a lower value is better; the other two figures are accuracies, where higher is better.

Why it matters

The paper makes two separate points. The first is engineering: models with billions of parameters can be trained with an intra-layer split that needs no new compiler and no library changes, only a few communication operations inserted into native PyTorch, and that coexists with pipeline model parallelism ^[1]. The second is empirical: scaling up, done with attention to architectural details such as where layer normalization sits in BERT-like models, produced better results than the previous state of the art on the benchmarks the authors measured, across both a GPT-2-like and a BERT-like design ^[1]. Together they describe a practical route from a model that fits on one device to one that does not.