Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM

By D. Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, M. Patwary, Vijay Anand Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro and 2 more. 2021

1,287 citations recorded by the source corpus.

Large language models have led to state-of-the-art accuracies across several tasks. However, training these models efficiently is challenging because: a) GPU memory capacity is limited, making it impossible to fit large models on even a multi-GPU server, and b) the number of compute operations required can result in unrealistically long training times. Consequently, new methods of model parallelism such as tensor and pipeline parallelism have been proposed. Unfortunately, naive usage of these methods leads to scaling issues at thousands of GPUs. In this paper, we show how tensor, pipeline, and data parallelism can be composed to scale to thousands of GPUs. We propose a novel interleaved pipelining schedule that can improve throughput by 10+% with memory footprint comparable to existing approaches. Our approach allows us to perform training iterations on a model with 1 trillion parameters at 502 petaFLOP/s on 3072 GPUs (per-GPU throughput of 52% of theoretical peak).

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Large Language Models, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Large Language Models

In Efficiency & Compression

Citations in the corpus

PaLM: Scaling Language Modeling with Pathways (2022, methodology): “Many other works aim to increase of the scale of models, while limiting communication overheads (Rajbhandari et al., 2020; Lepikhin et al., 2020; Li et al., 2020; Rasley et al., 2020; Rajbhandari et al., 2021; Ren et al., 2021; Narayanan et al., 2021a).”
Direct Preference Optimization: Your Language Model is Secretly a Reward Model (2023, background): “Self-supervised language models of increasing scale learn to complete some tasks zero-shot [31] or with few-shot prompts [6, 25, 11].”
Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model (2022, methodology): “Through a collaboration between NVIDIA Megatron-LM [63, 43] and Microsoft DeepSpeed [57, 65], we created an efﬁcient and scalable 3D parallel system capable of combining data, pipeline, and tensor-slicing based parallelism.”
Pathways: Asynchronous Distributed Dataflow for ML (2022, methodology): “System designers have adopted inge-nious techniques to execute pipelined (Narayanan et al., 2021; Rasley et al., 2020; Narayanan et al., 2019; Huang et al., 2019) and homogeneous MoE (Lepikhin et al., 2020; Fedus et al., 2021) models on MPI-style systems, but as we argue in detail later, the MPI…”
Reducing Activation Recomputation in Large Transformer Models (2022, methodology): “Compared to model parallelism, these techniques, which are based on data parallelism, are less eﬃcient and do not scale well to a large numbers of GPUs [13] and are thus a better ﬁt for ﬁnetuning models in resource-constrained environments.”
PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023, methodology): “For example, Megatron [21] demonstrated highly efficient deep transformer training on large clusters using 3D (data, tensor and pipeline) parallelism.”
MegaBlocks: Efficient Sparse Training with Mixture-of-Experts (2022, methodology): “We use gradient accumulation for all models and train with a batch size of 512 sequences and the largest micro batch size that does not run out of memory (Narayanan et al., 2021a).”
What Language Model to Train if You Have One Million GPU Hours? (2022, methodology): “We refer to hardware FLOP the estimates reported by our codebase, using the formula from Narayanan et al. (2021).”
DeepSpeed- Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale (2022, methodology): “Under the hood, it leverages the single GPU kernels to maximize per GPU memory bandwidth utilization, while using NCCL all-reduce collectives to perform the necessary across GPU communication as described in [19].”
A Survey of Large Language Models (2023, positive): “• Megatron-LM [124] is a PyTorch-based deep learning library developed by NVIDIA for training large-scale language models.”
This paper is included in the Proceedings of the 16th USENIX Symposium on Operating Systems Design and Implementation. (neutral): “…stage has different sets of operators and may locate on meshes with different shapes, in contrast to many SPMD pipeline-parallel training systems [40, 57], Alpa adopts an MPMD-style runtime to orchestrate the inter-op parallel execution – Alpa generates distinct static execution instructions for…”
Visually-Augmented Language Modeling (2022, background): “To further strengthen such contextual self-supervision, recent seminal works, e.g. GPT-3 (Brown et al., 2020) and Megatron-LM (Narayanan et al., 2021), focus on increasing the model size and the scale of pre-training corpus.”
Holistic Evaluation of Language Models (2023, background): “36 These changes have also translated from research to deployment: LMs are directly exposed as commercial APIs or are integratedintoubiquitousproductsaspartofanemergingcommercial ecosystem.”
HybridFlow: A Flexible and Efficient RLHF Framework (2024, neutral): “Numerous systems [23, 36, 48, 54, 71, 75, 89] optimized large model training through model parallelisms such as tensor parallelism and pipeline parallelism to partition models across devices.”
Kimi K2: Open Agentic Intelligence (2025, neutral): “Our strategy leverages a combination of 16-way Pipeline Parallelism (PP) with virtual stages [29, 54, 39, 58, 48, 22], 16-way Expert Parallelism (EP) [40], and ZeRO-1 Data Parallelism [61].”
Scaling Data-Constrained Language Models (2023, background): “These ﬁndings called for a renewed focus on the scaling of pre-training data rather than scaling model size via complex parallelization strategies [98, 91, 9, 78].”
SmoothLLM: Defending Large Language Models Against Jailbreaking Attacks (2023, neutral): “3 Estimates suggest that training GPT-3 took in excess of 800,000 GPU-hours (see, e.g., the estimates in §5.1 in [28]) and cost nearly $4 million, with training GPT-4 is thought to have cost nearly ten times that amount [29].”