Ring Attention with Blockwise Transformers for Near-Infinite Context

By Hao Liu, Matei A. Zaharia, Pieter Abbeel. International Conference on Learning Representations, 2023

539 citations recorded by the source corpus.

Transformers have emerged as the architecture of choice for many state-of-the-art AI models, showcasing exceptional performance across a wide range of AI applications. However, the memory demands imposed by Transformers limit their ability to handle long sequences, thereby posing challenges in utilizing videos, actions, and other long-form sequences and modalities in complex environments. We present a novel approach, Ring Attention with Blockwise Transformers (Ring Attention), which leverages blockwise computation of self-attention and feedforward to distribute long sequences across multiple devices while fully overlapping the communication of key-value blocks with the computation of blockwise attention. Our approach enables training and inference of sequences that are up to device count times longer than those achievable by prior memory-efficient Transformers, without resorting to approximations or incurring additional communication and computation overheads. Extensive experiments on language modeling and reinforcement learning tasks demonstrate the effectiveness of our approach in allowing millions of tokens context size and improving performance.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition

View Paper: https://www.semanticscholar.org/paper/02ad9f3fefe33cb9ca546591bec65dbdf7766c80

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

Citations in the corpus

RULER: What's the Real Context Size of Your Long-Context Language Models? (2024, background): “Flash attention (Dao et al., 2022; Dao, 2023) and Ring attention (Liu et al., 2023) significantly reduce the memory footprint required for processing long context.”
LongVILA: Scaling Long-Context Visual Language Models for Long Videos (2024, methodology): “Sequence Length 32K 64K 96K 128K 8K 16K 32K 64K 96K 128K Sequence Length Sequence Length Specifically, Ring-style systems Li et al. (2021; 2023a); Liu et al. (2023a) use Point-to-Point (P2P) communication primitives to collectively compute the attention module, while DeepSpeed-Ulysses Jacobs et al.…”
Long Context Transfer from Language to Vision (2024, methodology): “Efforts have also been made on data curation [22, 79, 5] and system optimization [40, 45, 27] during long context training.”
FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (2024, unknown): “We repeat the benchmarks 10 times and take the average timing.”
HunyuanVideo: A Systematic Framework For Large Video Generative Models (2024, unknown): “Specifically, Ring Attention [30] is used to achieve efficient training of long sequences through multiple GPUs, breaking through the GPU memory limitation of a single GPU.”
Video-MME: The First-Ever Comprehensive Evaluation Benchmark of Multi-modal LLMs in Video Analysis (2024, unknown): “For instance, exploring techniques like ring attention [37], as investigated by large world models [36], and training-free context extension methods could be beneficial [2].”
Parallelizing Linear Transformers with the Delta Rule over Sequence Length (2024, unknown): “…have made it possible to scale attention to longer sequences through hardware-aware restructuring of the intermediate computations [19, 17, 57, 14], these methods still require storing the key/value vectors of previous elements, and this “KV cache” (whose size grows linearly) can be…”
MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention (2024, unknown): “Similarly, we build a sparse index i b and use it to compute the sparse attention weights and attention output.”