Megatron-LM
In plain wordsNVIDIA's open-source software, first described in a 2019 research paper, for training very large language models across thousands of graphics chips by splitting each model in several ways at once.
A research framework and code base from NVIDIA for training large transformertransformer, n.: The design behind today's chatbots and much of modern AI. It handles text by letting every word look at every other word and decide which of them matter for understanding it. language models with combined parallelism, and the title of the 2019 paper that introduced its tensor parallelismtensor parallelism, n.: Splitting the individual giant tables of numbers inside each layer of a model across several chips, so that every chip does a slice of every calculation at the same time. scheme: splitting the feed-forward and attention blocks across GPUs so that each transformer layer needs only a few all-reduceall-reduce, n.: A group operation in which every chip contributes a set of numbers, the sets are added together, and every chip ends up with the same total. It is how parallel training keeps model copies in sync. operations — a scheme simple enough to write directly in PyTorch. Later versions added pipeline parallelismpipeline parallelism, n.: Splitting a model's layers into consecutive stages on different chips, like stations on an assembly line, and streaming small batches of data through so that every station stays busy. with an interleaved schedule, sequence parallelismsequence parallelism, n.: Splitting the long string of words a model is working on into chunks handled by different chips, which saves memory and makes it possible to train on very long documents., selective activation recomputation, contextcontext parallelism, n.: A way of training or running a model on extremely long inputs — whole books or codebases — by splitting the text across many chips and having them swap the information each part needs to 'see' the others. and expert parallelismexpert parallelism, n.: A way of running a 'mixture of experts' model — one made of many specialist sub-networks — by giving each chip its own set of specialists and sending each piece of text to whichever chip holds the specialists it needs., distributed optimizers and FP8 training; its reusable core was split out as the Megatron-Core library. Megatron-LM and its forks have trained many open and commercial models, including the 530-billion-parameter Megatron-Turing NLG built with Microsoft (2021), and its papers supplied much of the field's vocabulary of parallelism.
We forked Megatron-LM for the pretraining run and used Hugging Face tools only for evaluation.
Origin
Introduced by Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper and Bryan Catanzaro at NVIDIA (arXiv, September 2019), who trained an 8.3-billion-parameter model on 512 GPUs at 15.1 petaFLOP/s; extended by Deepak Narayanan and colleagues (2021) and Vijay Korthikanti and colleagues (2022).
Related
- tensor parallelismtensor parallelism, n.: Splitting the individual giant tables of numbers inside each layer of a model across several chips, so that every chip does a slice of every calculation at the same time.
- pipeline parallelismpipeline parallelism, n.: Splitting a model's layers into consecutive stages on different chips, like stations on an assembly line, and streaming small batches of data through so that every station stays busy.
- 3D parallelism3D parallelism, n.: Combining three ways of dividing up training work — by data, by layer stage and by slices within layers — so that thousands of chips can train one giant model efficiently.
- DeepSpeedDeepSpeed, prop. n.: Free software from Microsoft for training and running AI models too big for ordinary tools, spread over many chips; it is best known for a memory-saving technique called ZeRO.
- sequence parallelismsequence parallelism, n.: Splitting the long string of words a model is working on into chunks handled by different chips, which saves memory and makes it possible to train on very long documents.
- distributed trainingdistributed training, n.: Training one AI model on many computer chips at once, splitting up the data or the model itself among them and keeping them in sync, so that jobs that would take a single chip thousands of years finish in weeks.
Sources
- Shoeybi, M. et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053. arxiv.org/abs/1909.08053
- Narayanan, D. et al. (2021). Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. arXiv:2104.04473. arxiv.org/abs/2104.04473
Mentioned in
model parallelismmodel parallelism, n.: Splitting a single AI model that is too big for one chip into pieces spread over several chips, which must constantly pass intermediate results to each other as they work., DeepSpeedDeepSpeed, prop. n.: Free software from Microsoft for training and running AI models too big for ordinary tools, spread over many chips; it is best known for a memory-saving technique called ZeRO., NCCLNCCL, prop. n.: NVIDIA's software library that lets its graphics chips exchange data with each other quickly — inside one computer and across networks — and that nearly all large-scale AI training on NVIDIA hardware relies on., pipeline parallelismpipeline parallelism, n.: Splitting a model's layers into consecutive stages on different chips, like stations on an assembly line, and streaming small batches of data through so that every station stays busy., sequence parallelismsequence parallelism, n.: Splitting the long string of words a model is working on into chunks handled by different chips, which saves memory and makes it possible to train on very long documents., tensor parallelismtensor parallelism, n.: Splitting the individual giant tables of numbers inside each layer of a model across several chips, so that every chip does a slice of every calculation at the same time.