Megatron-LM: GPU cluster training and parallelism explained

Jared Casper explains how NVIDIA and Microsoft approached large language model training with Megatron-LM. Uploaded on October 27, 2023, the talk has no supplied recording date or software version. Its results and implementation details describe the work presented at that time.

The central example is the 530-billion-parameter Megatron-Turing NLG model. Casper reports that its training state required 8.4 TB of device memory. The supplied description places the training on a DGX SuperPOD with more than 3,000 A100 GPUs and an InfiniBand interconnect, a cluster-scale workload rather than a desktop local LLM setup.

Casper explains why replicating a model through data parallelism becomes insufficient when it cannot fit on one GPU. Tensor parallelism divides individual layers across devices; pipeline parallelism assigns groups of layers to different devices. Microbatches reduce idle time in the pipeline. In the experiment he discusses, keeping tensor parallelism within an eight-GPU node gave the best utilization. He also describes an interleaved schedule that reduced pipeline idle time, particularly with smaller batches.

The talk presents Megatron as an open source library built with PyTorch and describes combining it with DeepSpeed for the 530B training run. Practical lessons cover checkpoint throughput, choosing batch and parallelism sizes that permit different GPU allocations on a shared cluster, and validating ideas on smaller models before committing to months of training.