シーケンスの並列性とリング アテンション

A transformer processes a sequence of tokens, and its attention step lets every token look at every other token. That makes memory grow quickly with context length. The Ring Attention paper gives a concrete example: processing 100 million tokens with batch size 1 needs over 1,000 GB of memory even for a modest model with a hidden size of 1,024, while GPUs and TPUs typically have less than 100 GB of high-bandwidth memory. Sequence parallelism is the family of techniques that split a single sequence along its length, the token dimension, across several accelerators, so each one stores only a slice. It complements the older forms of parallelism: data parallelism splits the batch, tensor parallelism splits weight matrices, and pipeline parallelism splits layers.

Most transformer operations are easy to split by token, because they treat each token independently: layer normalization, the feed-forward network, dropout and the query, key and value projections. Attention is the exception. Each token's query must be compared with the keys of every other token in the sequence, and its output is a weighted sum of their values. If device 1 holds tokens 1 to 1,000 and device 2 holds tokens 1,001 to 2,000, device 2's queries still need device 1's keys and values. Every sequence-parallel method is really an answer to one question: how to get keys and values, or queries, to the right device without stalling the computation.

Ring Attention's answer is to arrange the devices in a ring. Each device keeps its own block of queries for the whole computation. Key and value blocks travel around the ring: at each step a device computes attention between its queries and the key-value block it currently holds, while simultaneously sending that block to the next device and receiving a new one from the previous device. After as many steps as there are devices, every query block has seen every key-value block. Because attention is computed blockwise with a running softmax, the same technique used by memory-efficient attention kernels, the partial results can be combined exactly, so the output is identical to ordinary attention, not an approximation. If each block's computation takes longer than its transfer, the communication is fully hidden.

The idea has a short history. Shenggui Li and colleagues at the National University of Singapore proposed sequence parallelism with Ring Self-Attention in May 2021; compared with tensor parallelism on 64 NVIDIA P100 GPUs, they reported 13.7 times larger maximum batch size and 3.0 times longer maximum sequence length. In May 2022, NVIDIA's Vijay Korthikanti and colleagues used the name sequence parallelism for a narrower technique inside Megatron-LM: splitting the LayerNorm and dropout regions along the sequence and swapping all-reduce operations for all-gather and reduce-scatter. Combined with selective recomputation, it cut activation memory by 5 times. In October 2023, Hao Liu, Matei Zaharia and Pieter Abbeel at UC Berkeley published Ring Attention with Blockwise Transformers, which fully overlapped communication with blockwise attention and feed-forward computation. In February 2024 the same group released Large World Model, a family of 7B models handling text and video beyond 1 million tokens.

The main alternative is DeepSpeed-Ulysses from Microsoft, published in September 2023. It also starts with the sequence split across GPUs, but just before attention it performs an all-to-all exchange so that each GPU receives the full sequence for a subset of attention heads, computes those heads, and then reverses the exchange. Its authors show that communication volume stays constant when sequence length and device count grow together, and report up to 2.5 times higher throughput with 4 times longer sequences than the prior state of the art. Because work is divided by heads, the number of devices it can use for one sequence is bounded by the number of attention heads, while Ring Attention divides by blocks of tokens. Megatron Core's context parallelism follows the Ring Attention pattern and adds cuDNN flash attention kernels and causal-mask load balancing.

The headline results are large. On 32 A100 GPUs, the Ring Attention paper reports over 1 million tokens of context for a 7B model, a 32 times improvement over the previous best memory-efficient transformer, and on TPUv4-512 a 256 times increase, allowing training sequences of over 30 million tokens. In general, Ring Attention allows sequences up to device-count times longer than memory-efficient transformers on a single device. Large World Model used it to train on TPUv4-1024, which its authors describe as roughly equivalent to 450 A100 GPUs, scaling RoPE positional embeddings as context grew.

There are real limitations. Hiding communication needs enough work per block: the Ring Attention paper derives that the block size must be at least the device's FLOPS divided by its interconnect bandwidth, which is about 1,000 tokens for A100s on NVLink or TPUs, but about 24,500 tokens for A100s linked by InfiniBand at 12.5 GB per second. Causal language models create load imbalance, because with contiguous blocks the device holding the first tokens has almost nothing to attend to while the last device does the most work; Striped Attention fixes this by giving each device tokens spread throughout the sequence, reporting up to 1.45 times speedups at 256K tokens on A100s and TPUs and 1.65 times at 786K tokens on 16 TPUv4 chips. Sequence parallelism also does not reduce the total compute of attention, which still grows with the square of sequence length, so a million-token step remains expensive.

In practice, using these methods requires a framework that supports them, such as Megatron Core with Transformer Engine, DeepSpeed, or a JAX implementation like the one published with Ring Attention, plus fast device-to-device links. Total GPU count becomes the product of all parallel dimensions; Megatron Core's documentation states it as tensor times context times pipeline times data parallel size. Training a long-context model also involves more than the parallelism scheme: the Large World Model authors describe data curation, progressive context extension and positional-embedding scaling as necessary parts of the process.