PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, Bryan Catanzaro

arXiv:1909.08053v4cs.CL

TL;DR

Very large transformer models are difficult to train under processor memory limits. The paper develops a simple intra-layer model-parallel approach in native PyTorch and combines it with scalable training and architectural analysis. It trains models up to 8.3 billion parameters and reports improved language-modeling and downstream-task results, including state-of-the-art benchmarks.

  • Problem

    Very large transformer models exceed modern processors’ memory limits, making scalable training difficult.

  • Method

    The paper implements efficient intra-layer model parallelism in existing PyTorch transformers using targeted modifications and a few communication operations.

  • Results

    The approach trains models up to 8.3 billion parameters and yields improved performance across GPT-2 and BERT evaluations, with state-of-the-art results on WikiText103, LAMBADA, and RACE.

  • Takeaways & Limitations

    Careful layer-normalization placement is critical for increasing BERT-like model accuracy as model size grows.

  • Takeaways & Limitations

    Models exceeding 16 billion parameters require hybrid intra-layer, inter-layer, and inter-node model parallelism because their memory exceeds that available within 16 GPUs of a DGX-2H box.

Abstract

from arXiv · show

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

1. Introduction

Large transformer models improve NLP performance but exceed single-processor memory limits. The paper introduces an efficient PyTorch intra-layer model-parallel approach and demonstrates scalable training and improved downstream results.

  • Motivation: Memory constraints make very large language models difficult to train, motivating additional model-parallel techniques.Existing approaches partition models to distribute weights and optimizer state, but may require model rewriting or custom compiler support.
  • Approach: The proposed intra-layer model parallelism uses a few targeted modifications to existing PyTorch transformer implementations and is orthogonal to pipeline parallelism.The approach requires no custom C++ code or compiler.
  • Scalability: 76% scaling efficiency is demonstrated using 512 GPUs for an empirical model- and data-parallel training analysis.The reported scaling uses a strong single-GPU baseline.
  • Model Scaling: Careful layer-normalization placement is critical for increasing accuracy as BERT-like models grow.The paper identifies normalization placement as an architectural factor affecting large-model performance.
  • Model Scaling: 8.3-billion-parameter GPT-2 and 3.9-billion-parameter BERT models show improved accuracy as model size increases.The study evaluates both left-to-right GPT-2 and bidirectional BERT models on downstream tasks.
  • Results: 10.8 perplexity on WikiText103, 66.5% accuracy on LAMBADA, and 90.9% accuracy on RACE are reported as state-of-the-art results.These are the headline benchmark results highlighted by the paper.

2. Background and Challenges

Transformer language models build on pretrained representations and are increasingly used for NLP because of their accuracy and compute efficiency. Scaling training across accelerators introduces memory, optimization, and parallelization challenges that motivate distributed tensor computation.

  • Neural Language Model Pretraining: Pretrained language models transfer contextual language representations and can be fine-tuned end to end on downstream tasks.This development has made pretrained language models central to NLP research.
  • Transformer Language Models: Transformer language models use encoder-only or decoder-only architectures depending on the application, including BERT and GPT-2.The paper studies both architectures.
  • Transformer Language Models: GPT-2 and BERT use GeLU nonlinearities and apply layer normalization to the inputs of attention and feed-forward layers.This differs from the original Transformer, which uses ReLU and output layer normalization.
  • Parallelism: Data parallelism splits minibatches across workers, whereas model parallelism distributes model memory and computation across workers.Weak scaling can increase data throughput, but large batches can complicate optimization.
  • Parallelism: Activation checkpointing reduces memory requirements by recomputing activations during the backward pass instead of storing them during the forward pass.It is presented as a complementary technique for scaling neural-network training.
  • Parallelism: Distributed tensor computation partitions tensor operations across devices to accelerate computation or increase model size.The paper exploits transformer attention-head parallelism while avoiding a new framework or compiler.

3. Model Parallel Transformers

The paper parallelizes transformer layers within their attention and MLP blocks, using partitioned matrix operations and a few strategically placed all-reduce operations. The design reduces synchronization and communication while remaining simple to implement in PyTorch and compatible with pipeline parallelism.

  • Transformer block parallelism: Transformer layers are parallelized separately in their self-attention and two-layer MLP blocks.The approach exploits transformer structure rather than introducing a new compiler or library.
  • MLP parallelism: Column-splitting the first MLP GEMM allows GeLU to run independently on each partition, removing a synchronization point.The second GEMM is split along rows, and its output is reduced before dropout.
  • Communication pattern: The MLP block requires one forward all-reduce and one backward all-reduce across GPUs.The conjugate f and g operators implement identity/all-reduce behavior in opposite passes.
  • Self-attention parallelism: Column-parallel partitioning lets each GPU compute attention heads locally without immediate communication.Key, query, and value GEMMs are partitioned across GPUs, while the output projection is parallelized along rows.
  • Embedding communication: Fusing the parallel output embedding with cross-entropy communicates scalar losses of size b × s instead of logits of size b × s × v.This avoids an all-gather over the large vocabulary dimension.
  • Communication pattern: A transformer layer uses four total communication operations across its forward and backward passes.The method duplicates some local computations, including layer normalization, dropout, and residual connections, to keep GPUs compute bound.
  • Implementation: The implementation adds only a few synchronization primitives or all-reduce operations, requires no compiler changes, and complements pipeline model parallelism.The approach is implemented in native PyTorch.

4. Setup

The experiments train GPT-2 and BERT-style models on a deduplicated aggregate corpus using mixed precision, activation checkpointing, and large-scale optimization settings. The setup combines multiple language-modeling datasets with task-specific training schedules and memory-management techniques.

  • Model configurations: The study evaluates left-to-right GPT-2 and bidirectional BERT transformer models.The models represent generative and masked-language-modeling approaches, respectively.
  • Training data: The aggregate training corpus combines Wikipedia, CC-Stories, RealNews, and OpenWebText after removing WikiText103 test articles.The authors filter documents shorter than 128 tokens and deduplicate similar content using locality-sensitive hashing.
  • Training data: 174 GB of deduplicated text remains after filtering documents and removing content with Jaccard similarity greater than 0.7.The filtering and deduplication target repeated or unsuitable training documents.
  • Training optimization: Mixed-precision training with dynamic loss scaling uses V100 Tensor Cores, while activation checkpointing is applied after every transformer layer.The setup also uses Adam with weight decay, global gradient-norm clipping of 1.0, and dropout of 0.1.
  • GPT-2 training: GPT-2 training uses sequences of 1024 subword units, batch size 512, and 300k iterations.The learning rate warms up for 3k iterations before cosine decay.
  • BERT training: BERT training uses batch size 1024, a 1.0e-4 learning rate, 10,000 warmup iterations, and linear decay over 2 million iterations.The configuration uses the original BERT dictionary and replaces next-sentence prediction with sentence-order prediction.

5. Experiments

The experiments evaluate weak scaling for model and model+data parallelism, then test how increasing GPT-2 and BERT model sizes affects language-modeling and downstream-task performance. They also examine architectural changes needed to scale BERT beyond BERT-Large.

  • Model and data parallelism: 39 TeraFLOPs is sustained by the 1.2B-parameter single-GPU baseline, reaching 30% of theoretical peak FLOPS.This configuration is used as the baseline for scaling comparisons.
  • Model and data parallelism: 74% scaling relative to the strong single-GPU baseline is achieved for the 8.3B model on 512 GPUs with model+data parallelism.The 8.3B model with 8-way model parallelism alone reaches 77% of linear scaling; added gradient communication lowers the 512-GPU result slightly.
  • Language modeling results using GPT-2: 10.81 adjusted perplexity and 66.51% accuracy are achieved by the 8.3B GPT-2 model on WikiText103 and LAMBADA, respectively.Increasing model size is associated with lower WikiText103 perplexity and higher LAMBADA cloze accuracy.
  • Bi-directional transformer results using BERT: Rearranging layer normalization and residual connections eliminates observed instabilities and enables scaling BERT-style models beyond BERT-Large.The rearranged architecture also produces lower training loss than the original architecture in the reported comparison.
  • Bi-directional transformer results using BERT: 1.58, 1.30, and 1.16 validation perplexities are obtained by the 336M, 1.3B, and 3.9B BERT models, respectively, showing a monotonic decrease with model size.The models use the rearranged architecture described for scaling BERT-style transformers.
  • Bi-directional transformer results using BERT: Downstream-task performance improves in all reported cases as model size increases, and the 3.9B model achieves state-of-the-art results on RACE.The evaluation covers MNLI, QQP, SQuAD 1.1, SQuAD 2.0, and RACE.

6. Conclusion and Future Work

The paper demonstrates scalable training of transformer models beyond single-GPU limits and identifies future directions for extending model scale, efficiency, architectures, and evaluations.

  • Conclusion: Models up to 8.3 billion parameters were trained on 512 NVIDIA V100 GPUs using 8-way model parallelism, sustaining 15.1 PetaFLOPs across the application.The approach required only a few modifications to existing PyTorch transformer implementations.
  • Conclusion: Careful placement of layer normalization in BERT-like models was critical to increasing accuracy as model size grew.The study also reported new state-of-the-art results on WikiText103, LAMBADA, and RACE.
  • Future Work: Further scaling beyond 16 billion parameters will require hybrid intra-layer and inter-layer model parallelism with inter-node model parallelism.The paper notes that such models demand more memory than is available within 16 GPUs of a DGX-2H box.
  • Future Work: Future investigations include more efficient optimizers, other model families, harder and more diverse downstream tasks, and knowledge distillation.Examples include XLNet, T5, generative question answering, summarization, and conversation.

B.1. Hybrid Model and Data Parallelism

Hybrid model and data parallelism distributes model instances across GPU groups while replicating corresponding parameters across data-parallel groups.

  • GPU Grouping: Two or more GPUs within a server form a model-parallel group containing one model instance distributed across those GPUs.The section describes GPUs 1 to 8 as an example grouping.
  • Section Scope: The section introduces hybrid model and data parallelism and refers to a table of BERT finetuning hyperparameters.The supplied table passage only identifies its role and does not provide the hyperparameter values.
  • GPU Grouping: GPUs occupying the same position in different model-parallel groups form data-parallel groups holding identical model parameters.These groups can span GPUs within one server or, more typically, multiple servers.

B.2. Model Parallel Random Number Generation

Random number generation is handled separately inside and outside model-parallel regions so dropout remains synchronized where required while retaining worker-specific randomness.

  • Random Number Generation: Dropout occurs both outside model-parallel regions and inside the self-attention block, requiring distinct random-number handling.The paper seeds generators identically at training start to synchronize residual-connection dropout across model-parallel workers.
  • Random Number Generation: A separate dropout random-number generator is uniquely seeded for each model-parallel worker to produce randomness across the entire operation.This worker-specific generator applies to dropout within model-parallel regions.
  • GPU Grouping: Figure 8 depicts GPU grouping for hybrid model and data parallelism with 8-way model parallelism and 64-way data parallelism.The figure provides the grouping context for the parallel workers discussed in this section.

C. Text Samples

This section presents generated text samples alongside contextual excerpts covering flowers, language models, model limitations, and a machine-learning conference.

  • Generated Samples: The section introduces text samples generated by Megatron-LM from context prompts and notes that some samples are cut short.The supplied passages include a flower-related generation after a botanical context prompt.
  • Flower Sample: The supplied excerpts describe flower structures and their reproductive and pollination-related functions.They mention petals, sepals, the calyx, stamens, anthers, ovaries, and pollinators.
  • Other Contexts: Other excerpts discuss GPT-2, language-model limitations, multilingual phrasing differences, ICML 2019, and natural language processing.The language-model discussion mentions data limitations, ambiguity, overfitting, and differing rules across languages.

D. Further Scaling Analysis

The further scaling analysis examines how attention heads affect scaling and reports strong-scaling results for a 1.2 billion parameter model.

  • The analysis studies the effect of attention-head count on scaling results.
  • The experiments include strong-scaling measurements for a 1.2 billion parameter model.
  • The analysis considers both model-parallel scaling and training acceleration for smaller models.

D.1. Attention Heads and Scaling

For an 8.3 billion parameter configuration using 8-way model parallelism, increasing attention heads from 16 to 32 slightly reduces scaling efficiency because self-attention computations become less favorable.

  • D.1. Attention Heads and Scaling: 16 to 32 attention heads are evaluated for the 8.3 billion parameter configuration with 8-way model parallelism.
  • D.1. Attention Heads and Scaling: Increasing attention heads makes some self-attention GEMMs smaller.
  • D.1. Attention Heads and Scaling: Increasing attention heads slightly decreases scaling efficiency as self-attention softmax size also increases.

D.2. Strong Scaling

The strong-scaling analysis evaluates model parallelism across fixed-size models and several language-model benchmarks. It reports faster training for a 1.2 billion parameter model while describing diminishing returns as per-GPU computation falls and overheads rise.

  • D.2. Strong Scaling: Model parallelism is designed primarily to train models larger than a single GPU can accommodate.
  • D.2. Strong Scaling: 64% faster training is obtained with two GPUs for a fixed 1.2 billion parameter model and batch size of 8 samples per iteration.
  • D.2. Strong Scaling: Additional GPUs produce diminishing returns because per-GPU computation decreases while memory-bandwidth and communication overheads increase.
  • D.2. Strong Scaling: The evaluation methodology covers WikiText103 perplexity and LAMBADA cloze-style prediction accuracy.
  • D.2. Strong Scaling: WikiText103 perplexity is computed from average cross entropy, with normalization based on the original tokenization and overlapping evaluation for fixed transformer contexts.
  • D.2. Strong Scaling: LAMBADA evaluation masks the final word in a 4–5 sentence context and requires all corresponding subword predictions to be correct.
Loading 1909.08053v4…