PaperlayerView exact revision on arXiv

Source-linked AI summary

Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model

Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, Elton Zhang, Rewon Child, Reza Yazdani Aminabadi, Julie Bernauer, Xia Song, Mohammad Shoeybi, Yuxiong He, Michael Houston, Saurabh Tiwary, Bryan Catanzaro

arXiv:2201.11990v3cs.CL

TL;DR

Large language models increasingly require specialized infrastructure, algorithms, and data practices because their parameter sizes exceed ordinary training resources. This paper presents MT-NLG, a 530-billion-parameter autoregressive transformer, and its 3D-parallel training system, curated corpus, and training procedures. MT-NLG achieved superior zero-, one-, and few-shot results on several NLP benchmarks, while analyses found substantial bias and limits to in-context generalization.

  • Problem

    Training rapidly growing language models requires advances in infrastructure, parallelism, software, algorithms, and data curation because their memory and compute demands exceed available single-GPU resources.

  • Method

    The paper trains MT-NLG, a 530-billion-parameter autoregressive transformer, using DeepSpeed–Megatron 3D parallelism, curated training data, and documented stability-oriented configurations.

  • Results

    MT-NLG achieved superior zero-, one-, and few-shot performance on several NLP benchmarks and exhibited in-context learning whose performance depends on shot-example distributions.

  • Takeaways & Limitations

    In-context learning can reflect grammatical and syntactic structure, but matching shot and evaluation distributions is necessary for optimal performance and does not ensure out-of-distribution generalization.

  • Takeaways & Limitations

    The model was trained without anti-bias countermeasures and should not be deployed as-is; benchmark performance can also exploit spurious dataset correlations.

Abstract

from arXiv · show

Pretrained general-purpose language models can achieve state-of-the-art accuracies in various natural language processing domains by adapting to downstream tasks via zero-shot, few-shot and fine-tuning techniques. Because of their success, the size of these models has increased rapidly, requiring high-performance hardware, software, and algorithmic techniques to enable training such large models. As the result of a joint effort between Microsoft and NVIDIA, we present details on the training of the largest monolithic transformer based language model, Megatron-Turing NLG 530B (MT-NLG), with 530 billion parameters. In this paper, we first focus on the infrastructure as well as the 3D parallelism methodology used to train this model using DeepSpeed and Megatron. Next, we detail the training process, the design of our training corpus, and our data curation techniques, which we believe is a key ingredient to the success of the model. Finally, we discuss various evaluation results, as well as other interesting observations and new properties exhibited by MT-NLG. We demonstrate that MT-NLG achieves superior zero-, one-, and few-shot learning accuracies on several NLP benchmarks and establishes new state-of-the-art results. We believe that our contributions will help further the development of large-scale training infrastructures, large-scale language models, and natural language generations.

1 Introduction

Rapidly growing language models improve downstream performance but create memory, compute, and training-time challenges. The paper presents MT-NLG, a 530-billion-parameter model, and the infrastructure, data, and training innovations used to build it.

  • Scaling model size has improved zero-shot and few-shot performance, strengthening the case for large-scale language modeling.Foundation models combine transformers, self-supervised learning, few-shot conditioning, and fine-tuning for downstream adaptation.
  • Training very large models exceeds single-GPU memory and can require impractically long runtimes without coordinated algorithm, software, and hardware optimization.Efficient parallelism must scale across memory and compute resources on thousands of GPUs.
  • MT-NLG contains 530 billion parameters, making it, to the authors’ knowledge, the largest monolithic language model trained at that time.Its parameter count is three times larger than GPT-3, while sparse models may have larger total parameter counts.
  • DeepSpeed and Megatron enabled an efficient 3D parallel system combining data, pipeline, and tensor-slicing parallelism.The design combines tensor-slicing and pipeline parallelism in the regimes where each is most effective.
  • The paper documents training infrastructure, datasets, evaluation, social-bias analysis, in-context learning, and qualitative generation results.These topics are organized across Sections 2–7.

2 Large Model Training Infrastructure

Training MT-NLG requires addressing memory and compute constraints that no single parallelism strategy solves alone. The system therefore combines complementary data, tensor, and pipeline parallelism to improve scalability and efficiency.

  • Memory and Compute Efficiency: A 530-billion-parameter model requires over 10 terabytes of aggregate memory for weights, gradients, and optimizer states.With Adam and mixed precision, training uses 20 bytes per parameter.
  • Memory and Compute Efficiency: Activation memory can also be substantial, scaling with batch size, sequence length, and model dimensions.The stated configuration requires approximately 16.9 terabytes of activation memory.
  • Memory and Compute Efficiency: Gradient accumulation reduces peak activation memory by processing micro-batches sequentially and freeing their activations before weight updates.The full training batch is split into micro-batches whose gradients are accumulated.
  • Tradeoffs of Data, Tensor, and Pipeline Parallelism: Data parallelism offers compute efficiency but replicates model state and incurs increasing gradient-aggregation costs as model size grows.ZeRO improves memory efficiency by partitioning replicated data among data-parallel workers.
  • Tradeoffs of Data, Tensor, and Pipeline Parallelism: Tensor parallelism partitions individual layers and reduces model memory, but requires high-bandwidth communication and can reduce compute efficiency.Megatron uses tensor parallelism to partition transformer blocks across workers.
  • Tradeoffs of Data, Tensor, and Pipeline Parallelism: Pipeline parallelism divides layers into stages and reduces model memory, but introduces activation storage, load-balancing, and pipeline-bubble constraints.Its efficiency depends on the number of micro-batches relative to pipeline stages.
  • Tradeoffs of Data, Tensor, and Pipeline Parallelism: 3D parallelism combines data, tensor, and pipeline parallelism because each individual technique has complementary strengths and limitations.The combination addresses compute and memory efficiency simultaneously.

3 Training Dataset and Model Configuration

The paper constructs a large curated training corpus and specifies a 530-billion-parameter transformer configuration with stability-oriented training choices. Data filtering, deduplication, and controlled blending support the training process.

  • Training Dataset and Preprocessing: The training corpus combines selected Pile datasets with two Common Crawl snapshots, RealNews, and CC-Stories.Table 1 lists the 15 datasets used for MT-NLG training.
  • Training Dataset and Preprocessing: Common Crawl processing includes language detection, HTML text extraction, classifier scoring, and score-based document filtering.The pipeline uses the 2020-50 and 2021-04 snapshots to acquire approximately 150 billion tokens.
  • Training Dataset and Preprocessing: A fastText quality classifier achieved 90.3% accuracy on held-out documents before filtering extracted Common Crawl text.Positive examples came from OpenWebText2, Wikipedia, and Books3, while negatives came from extracted Common Crawl text.
  • Training Dataset and Preprocessing: The pipeline applies Unicode correction, non-English and short-document removal, downstream-task data removal, fuzzy deduplication, and dataset-priority selection.Fuzzy deduplication uses hashing, min-hashes, locality-sensitive hashing, and approximate Jaccard similarity.
  • Training Dataset and Preprocessing: Datasets are blended into heterogeneous batches while tracking under- and oversampling to maintain the chosen mixing-weight distribution.Batch composition is adjusted at each step rather than relying on an even split.
  • Model and Training Process: MT-NLG uses a 105-layer autoregressive transformer with hidden dimension 20480, 128 attention heads, sequence length 2048, and global batch size 1920.The model uses 8-way tensor and 35-way pipeline parallelism and trains on 270 billion tokens from a 339-billion-token dataset.
  • Model and Training Process: Training stability depends on learning rate, weight initialization, and Adam settings, including a lower β2 and initialization standard deviation of approximately 1/(3 ∗H).The paper reports that higher learning rates and higher-variance initialization can destabilize or prevent convergence.

4 Results and Achievements

MT-NLG is evaluated across diverse NLP benchmarks using zero-, one-, and few-shot prompting, with results showing strong performance and task-dependent effects of additional examples. It establishes new state-of-the-art results on LAMBADA and substantially improves over prior models on several tasks.

  • Training Evaluation: The validation cross-entropy loss decreases from 3.15 after 1 billion training tokens to 1.85 after 270 billion tokens.At the end of the initial 12-billion-token batch-size phase, the loss is 2.31.
  • Completion Prediction: MT-NLG establishes new state-of-the-art accuracy on LAMBADA in zero-shot, one-shot, and few-shot settings.The evaluation uses narrative passages and tests prediction of the final word from broader discourse context.
  • Reading Comprehension: MT-NLG’s zero-shot RACE-h performance surpasses GPT-3’s few-shot performance by +1.14%.RACE-h does not benefit substantially from including more examples in the prompt.
  • Reading Comprehension: BoolQ improves significantly from zero-shot to few-shot, while RACE-h benefits comparatively less from additional examples.The paper attributes this difference to BoolQ’s structured single-word answer format versus RACE-h’s more natural-sentence answers.
  • Commonsense Reasoning: MT-NLG improves over GPT-3 by +2.81% in zero-shot Winogrande accuracy, while few-shot performance gains are stronger across commonsense tasks.One example provides only marginal improvement, whereas additional examples substantially improve performance.
  • Natural Language Inference: On HANS, MT-NLG moves from near-chance zero-shot performance to a large few-shot improvement, whereas GPT-2 remains near chance.On ANLI, one example improves performance, but additional examples do not provide further gains.
  • Word Sense Disambiguation: Word-in-Context performance is slightly below chance at zero-shot but surpasses chance with few-shot examples.The task is evaluated by comparing likelihoods for “yes” and “no” continuations.

5 Exploring Social Biases

MT-NLG was analyzed for gender, ethnicity, and religious bias after being trained without anti-bias countermeasures. The analyses found pervasive and sometimes offensive biases, supporting the authors’ warning against deploying the model as-is.

  • 5.1 Introducing the Challenge of Social Bias in NLP Models: The model was trained without anti-bias countermeasures, and the authors state that it should not be deployed without them.They position MT-NLG as a resource for research on future countermeasures and deployments that incorporate them.
  • 5.2 Gender and Occupation Analysis: 78% of 323 occupations received higher model probabilities for male than female identifiers, indicating a general masculine bias.The occupation prompts used gender identifiers such as male, man, female, and woman.
  • 5.2 Gender and Occupation Analysis: The average occupation bias score was −0.77, indicating that the model leaned toward male identifiers across more occupations.A score of 0 indicates no biased association; negative scores indicate a skew toward male identifiers.
  • 5.3 Adjective Co-Occurrence Analysis: Co-occurrence analysis captures surface-level adjective frequencies but does not account for sentiment or context.The authors also caution that the gender study used only three templates and was less robust than their other analyses.
  • 5.3 Adjective Co-Occurrence Analysis: Gender co-occurrence analysis found stereotypical associations, including more appearance-related adjectives for female identifiers and more diverse adjectives for male identifiers.Almost 80 of the top 100 gender-associated adjectives were identical across genders, while distinct-word analysis highlighted stereotypical differences.
  • 5.3 Adjective Co-Occurrence Analysis: Ethnicity analyses found more positive adjectives associated with White ethnicity and offensive stereotypes associated with other ethnicities, while Black prompts showed substantially more negative sentiment.The authors describe these results as unacceptable for deployment and call for anti-bias countermeasures.

6 Natural Language Understanding and In-Context Learning

MT-NLG is evaluated on systematic language understanding and in-context learning, revealing gains from scale, pretraining, and carefully composed prompts, alongside persistent heuristic and formatting sensitivities.

  • Natural Language Understanding: HANS tests whether models apply entailment rules rather than superficial heuristics such as vocabulary overlap or shared subsequences.
  • Model Size and Training: A 530-billion-parameter MT-NLG can largely escape superficial heuristics, whereas 1.5-billion-parameter GPT-2 remains near 50% accuracy on HANS regardless of shot count.
  • Model Size and Training: Additional autoregressive pretraining and more prompt examples both affect HANS performance, with peak accuracy for most checkpoints at 2-shot prompting.
  • Distribution of Shots: Sanitized prompts use examples from subcases different from the evaluated example, isolating guidance from same-subcase distribution overlap.
  • Shot Composition: Shot order, label composition, and sample distribution substantially affect performance, with one reported configuration reaching 78.6% accuracy.
  • Evaluation Summary: MT-NLG can apply many grammatical rules, but models initially rely on lexical heuristics, class biases, and task formatting, requiring special effort to elicit correct responses.

7 Qualitative Examples for MT-NLG Generation Capabilities

Qualitative evaluations probe MT-NLG’s generation across riddles, Jeopardy questions, code, arithmetic, and free-form writing, showing useful but uneven capabilities without task-specific fine-tuning.

  • Overview: MT-NLG is qualitatively capable of solving riddles, answering Jeopardy questions, and generating code off-the-shelf.
  • Riddle Answer Generation: Riddle prompts elicit line-by-line interpretations and answers that usually make sense, though performance is not always perfect.
  • Jeopardy Questions: The Jeopardy evaluation tests whether memorized knowledge can support answers in a guessing-game format requiring reasoning over hints.
  • Code Generation: MT-NLG consistently generates syntactically correct code and can implement simple tasks correctly from commented function signatures.
  • Inferring Arithmetic Operations: Arithmetic experiments examine whether answers reflect understanding of mathematical expressions rather than merely reproducing encountered examples.
  • Free-form Generative Writing Assistance: For free-form writing, the authors prompted MT-NLG with the paper’s introduction and sampled candidate sentences iteratively to help author the abstract.

8 Related Works

Related work traces scaling from early billion-parameter models to dense transformer systems, contrasts monolithic models with mixture-of-experts approaches, and notes multitask fine-tuning for zero-shot learning.

  • Model Scaling: Language-model scaling progressed from billion-parameter LSTMs through BERT and GPT-2 as large-scale pretraining with transformers became dominant.
  • Mixture-of-Experts: Mixture-of-experts models reach trillion-parameter scales by selectively activating subsets of parameters, whereas MT-NLG focuses on monolithic dense transformers.
  • Multitask Fine-Tuning: T0 and FLAN improve zero-shot learning through large-scale multitask fine-tuning, with larger pretrained models extracting greater benefit from these methods.

9 Conclusions

The paper presents MT-NLG, a 530-billion-parameter autoregressive transformer with strong in-context learning, and documents the infrastructure, training practices, evaluations, biases, and limitations surrounding it.

  • Conclusions: MT-NLG is a 530-billion-parameter left-to-right autoregressive transformer with strong in-context learning capabilities.
  • Conclusions: The model achieved superior zero-, one-, and few-shot performance on several NLP benchmarks, establishing new state-of-the-art results.
  • Conclusions: The paper presents 3D-parallelism strategies, hardware infrastructures, training configurations, and datasets for efficient large-scale training and future research.
  • Conclusions: Its analysis also examines social biases and factors affecting in-context learning, highlighting limitations of current large language models.

A Analysis Results on HANS Dataset

HANS evaluates whether NLI models rely on three superficial syntactic heuristics using balanced entailment and non-entailment examples generated from structured templates.

  • Dataset design: HANS tests lexical overlap, subsequence, and constituent heuristics that can encourage incorrect entailment judgments.Each heuristic represents a distinct superficial relationship between premise and hypothesis.
  • Dataset design: 30 subcases arise from five entailment-supporting and five contradicting templates for each of the three heuristics.Each subcase targets a specific grammatical or syntactic structure.
  • Dataset construction: 30,000 balanced test examples and 30,000 training examples are generated, with training examples used for few-shot prompts.The test set contains equal numbers of entailment and non-entailment examples.

A.2 Performance per subcase

Per-subcase analysis shows that MT-NLG both relies on superficial heuristics and captures several linguistic rules, while some structurally ambiguous cases remain difficult.

  • Evaluation setup: 32-shot MT-NLG accuracy is reported for every HANS subcase after shifting prediction means to counter class biases.The evaluation uses a model pretrained on 270 billion tokens.
  • Heuristic behavior: Heuristic reliance remains visible: non-entailment is nearly perfect for lexical overlap but harder when the hypothesis appears verbatim as a subsequence.The model struggles more when superficial similarity strongly favors entailment.
  • Syntactic reasoning: Despite autoregressive-only training, MT-NLG learns rules involving passive voice, subject-object order, relative clauses, and verb transitivity.It can use syntactic structure to escape misleading textual similarity during inference.
  • Difficult cases: Garden-path sentences and omitted-relative-pronoun constructions are among the most problematic cases, often also confusing to humans.These examples involve temporary ambiguity or difficult past-participle structures.

B Qualitative Examples of MT-NLG Text Generation

The paper presents qualitative MT-NLG generations covering open-ended reasoning, riddles, reverse trivia, code, mathematical operations, and association tests.

  • Qualitative examples: Qualitative examples examine open-ended, creative, or logical generation beyond benchmark evaluations.The section presents these outputs as qualitative examples.
  • Riddles: Table 13 shows MT-NLG solving riddles and often generating interpretations after providing an answer.The examples focus on the model’s explanatory behavior after answering.
  • Riddles: Table 14 shows different interpretations of the same riddle depending on the answer generated.The alternative-answer examples connect the model’s explanation to its selected answer.
  • Trivia: Table 15 presents reverse-trivia answers with few-shot context prepended to each question.The examples use Jeopardy!-style reverse trivia questions.
  • Code and mathematics: Tables 16 and 17 provide samples of code generation and mathematical-operation inference.These examples extend the qualitative evaluation to programming and arithmetic-like tasks.
  • Association analysis: Tables 18 and 19 list occupation lexicons used to test associations between gender and profession.The tables document the occupation vocabulary for the association analysis.
Loading 2201.11990v3…