Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
Deepak Narayanan‡⋆, Mohammad Shoeybi†, Jared Casper†, Patrick LeGresley†, Mostofa Patwary†, Vijay Korthikanti†, Dmitri Vainbrand†, Prethvi Kashinkunti†, Julie Bernauer†, Bryan Catanzaro†, Amar Phanishayee∗, Matei Zaharia‡ †NVIDIA ‡Stanford University ∗Microsoft Research
Abstract
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).
中文速览
训练超大规模语言模型面临两大拦路虎：单张GPU根本装不下千亿参数的模型，就算勉强塞进去训练时间也长得离谱。为此，研究者提出将张量并行（tensor parallelism）、流水线并行（pipeline parallelism）和数据并行（data parallelism）三种策略巧妙组合，形成"PTD-P"方案，并设计了一种新颖的交错流水线调度方式（interleaved pipeline schedule），能在内存占用相近的前提下将训练吞吐量提升10%以上。最终，该方法在3072块A100 GPU上以502 petaFLOP/s的聚合算力成功训练了一个万亿参数的GPT模型，单卡利用率达到理论峰值的52%，整个训练周期估计约需3个月，比同类方案（如ZeRO-3）快约70%。这项工作证明了在数千张GPU上高效训练超大语言模型是完全可行的，为后续千亿乃至万亿参数模型的实用化训练提供了清晰的工程路线图。
原文 arXiv:2104.04473;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2104.04473v5