Megatron-Turing NLG 530B

What it is

Megatron-Turing NLG 530B is a language model with 530 billion parameters, announced jointly by NVIDIA and Microsoft on 11 October 2021. At the time it was three times the size of the largest existing model of its type^[1]. Parameters are the numbers a network adjusts during training; more of them means more capacity to fit patterns, and also more computation to train and to run. The technical paper describes the model as the largest monolithic transformer-based language model at the time^[2]. "Monolithic" here means one dense network rather than a mixture-of-experts design, which routes different inputs to separate sub-networks.

Architecture and training data

The shape of the network is given by a handful of numbers: 105 layers, a hidden dimension of 20,480, 128 attention heads, and a sequence length of 2,048 tokens^[1]. Attention heads are parallel branches that each compute their own attention over the input, and depth (the layer count) sets how many times information is reworked on the way through. The model was trained on 270 billion tokens drawn from a curated set of 339 billion tokens spanning 15 datasets. Those datasets include a subset of The Pile, two Common Crawl snapshots, RealNews and CC-Stories^[1].

Training system

Training ran on Selene, an NVIDIA supercomputer built from 560 DGX A100 servers, each holding eight A100 80GB GPUs^[1]. A model of this size does not fit on a single GPU, so the work has to be split: the model itself is divided across chips, the layers are arranged in stages that pass activations along, and batches of data are spread out as well. To do this the two companies combined NVIDIA's Megatron-LM with Microsoft's DeepSpeed and ran data, pipeline and tensor parallelism at the same time — what is usually called 3D parallelism — reaching 113 to 126 teraFLOP/s per GPU^[1][2].

Results

The model set state-of-the-art zero-, one- and few-shot results on LAMBADA and PiQA^[1]. Zero-, one- and few-shot refer to how many worked examples are placed in the prompt — none, one, or a handful — with the model's weights left untouched. In the few-shot setting it reached 0.872 accuracy on LAMBADA and 0.832 on PiQA^[1]. NVIDIA also reported qualitative analysis showing that the model could infer basic mathematical operations from context and handle syntactic structures that challenge fine-tuned models, in both cases without any fine-tuning^[1].

Paper and release

The technical paper, by Shaden Smith, Mostofa Patwary and 18 other authors, was posted to arXiv on 28 January 2022^[2]. The model weights were not publicly released^[1], so outside researchers could read the methods and the numbers but not run the model themselves.

Where it sits

NVIDIA described MT-NLG as the successor to Turing NLG 17B and Megatron-LM, placing it on a continuing line of large models built by the company^[1]. Both companies acknowledged that the model inherits biases from its training data and stressed responsible AI principles for production use^[1]. The announcement is also a record of an engineering approach: a very large dense model trained by stitching two independently developed software stacks together on a dedicated cluster, at a scale that made the training run itself the headline.