context parallelism
In plain wordsA 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.
A form of 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. that splits each input sequence along its length across devices for every layer, so that activation memory per device falls in proportion to the number of devices and contexts of hundreds of thousands to millions of tokens become trainable. Most of a 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. processes tokens independently, but self-attentionself-attention, n.: Attention in which a passage looks at itself: every word checks every other word in the same text to work out its meaning in context, such as which noun 'it' refers to. needs every query to see the relevant keys and values, so devices must exchange them. Ring attention (Hao Liu, Matei Zaharia and Pieter Abbeel, 2023) passes key–value blocks around a ring of devices while computing attention on the blocks already held, hiding communication behind computation; DeepSpeed Ulysses (2023) instead uses all-to-all collectives to re-partition by attention head during attention; other designs all-gather keys and values. Causal masking creates load imbalance, handled by reordering chunks. Meta's Llama 3 used context parallelism as the fourth dimension of its '4D' layout for long-context training, and it also speeds long prefills at inference.
To train at 1M tokens we added context parallelism of 16 on top of tensor parallelism.
Origin
Grew out of the long-sequence sense of 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. (Shenggui Li and colleagues, 2021). The name 'context parallelism' was adopted in NVIDIA's Megatron-Core around 2023–2024, whose documentation contrasts it with Megatron's own sequence parallelism, and spread through Meta's Llama 3 report (2024) and PyTorch's context-parallel APIs.
Related
- 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.
- long contextlong context, n.: The ability of an AI model to work with very long inputs — whole books, codebases or hours of transcripts — in one go, and the techniques that make this possible.
- self-attentionself-attention, n.: Attention in which a passage looks at itself: every word checks every other word in the same text to work out its meaning in context, such as which noun 'it' refers to.
- 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.
- collective communicationcollective communication, n.: The standard patterns by which a whole group of chips exchanges data at once — for example, everyone adding up their numbers and all getting the total — which large training jobs use constantly.
- context extensioncontext extension, n.: Techniques for stretching how much text a trained AI model can handle at once, usually by adjusting how it represents word positions and then training it briefly on longer examples.
Sources
- Liu, H., Zaharia, M. & Abbeel, P. (2023). Ring Attention with Blockwise Transformers for Near-Infinite Context. arXiv:2310.01889. arxiv.org/abs/2310.01889
- Jacobs, S. A., Tanaka, M., Zhang, C., Zhang, M., Song, S. L., Rajbhandari, S. & He, Y. (2023). DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models. arXiv:2309.14509. arxiv.org/abs/2309.14509
Mentioned in
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., 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., 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., Megatron-LMMegatron-LM, prop. n.: NVIDIA'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., 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.