PaperlayerView exact revision on arXiv

Source-linked AI summary

Ring Attention with Blockwise Transformers for Near-Infinite Context

Hao Liu, Matei Zaharia, Pieter Abbeel

arXiv:2310.01889v4cs.CL

TL;DR

Transformers’ memory demands limit their use on long sequences. Ring Attention distributes blockwise attention and feedforward across devices while overlapping key-value communication with computation. Experiments report context lengths beyond 100 million tokens and improved reinforcement-learning performance.

  • Problem

    Transformer self-attention and layer outputs impose memory costs that constrain long-context training and inference.

  • Method

    Ring Attention distributes blockwise self-attention and feedforward across devices while overlapping key-value-block communication with computation.

  • Results

    Ring Attention enables sequences exceeding 100 million tokens without attention approximations and outperforms blockwise parallel transformers across six reinforcement-learning tasks.

  • Takeaways & Limitations

    Context length scales linearly with device count while maintaining performance and avoiding the memory bottleneck imposed by individual devices.

Abstract

from arXiv · show

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.

1 Introduction

Transformers are powerful but their quadratic self-attention memory cost limits context length. Ring Attention distributes blockwise computation across devices, overlapping communication and computation to scale context with device count.

  • Motivation: Transformers support state-of-the-art AI systems, but self-attention has memory cost quadratic in input sequence length.Long contexts are important for books, images, videos, codebases, and scientific data.
  • Motivation: Blockwise attention and feedforward reduce memory without approximating attention, but storing every layer’s output remains a significant challenge.The stored outputs arise from self-attention interactions among all sequence elements.
  • Motivation: 100 million tokens require over 1000GB of memory for a modest model with hidden size 1024, exceeding contemporary GPU and TPU HBM capacity.The stated comparison uses batch size 1 and devices typically having less than 100GB of HBM.
  • Approach: Ring Attention distributes blockwise attention’s outer loop across hosts and overlaps key-value communication with blockwise attention and feedforward computation.The approach relies on blockwise computation whose results are invariant to operation ordering when combined correctly.
  • Results: More than 500 times longer sequences than prior memory-efficient state-of-the-art methods can be trained, including sequences exceeding 100 million tokens without attention approximations.The introduction also states that training and inference lengths scale with the number of devices.
  • Contributions: The contributions are a memory-efficient architecture with context length scaling linearly with device count and experiments demonstrating its effectiveness.The claimed scaling maintains performance while removing the individual-device memory bottleneck.

2 Large Context Memory Constraint

Memory-efficient Transformers reduce attention storage through blockwise computation, but long-context scaling remains constrained by storing each layer’s output. Self-attention’s all-to-all interactions make avoiding this storage computationally impractical without recomputation.

  • Blockwise attention reduces attention memory overhead to 2bsh bytes per layer.Here, b is batch size, s is sequence length, and h is hidden size.
  • BPT still requires storing each layer’s output to support subsequent self-attention.The stored outputs are needed because self-attention involves interactions among all sequence elements.
  • Without stored layer outputs, subsequent self-attention becomes computationally impractical and requires recomputation for each sequence.

3 Ring Attention with Blockwise Parallel Transformers

Ring Attention distributes sequence blocks across hosts and circulates key-value blocks around a ring while overlapping communication with blockwise attention computation. The design requires six blocks of memory per host and selects block and sequence sizes from the compute-to-bandwidth ratio.

  • Ring-Based Blockwise Attention: Ring Attention assigns each host one query block and circulates key-value blocks among neighboring hosts in a ring.Each host sends key-value blocks onward while receiving blocks from the preceding host.
  • Ring-Based Blockwise Attention: Permutation invariance lets each query block attend to key-value blocks in any order when block statistics are combined correctly for rescaling.
  • Arithmetic Intensity Between Hosts: c ≥ F/B is required to hide key-value transfer behind computation, where c is block size, F is host compute, and B is inter-host bandwidth.The condition follows from 4dc^2/F ≥ 4cd/B.
  • Algorithm and Implementation: Communication is overlapped with blockwise attention computation as key-value blocks are sent and received between neighboring hosts.The implementation uses memory-efficient local computation together with collective permutation of key-value blocks.
  • Memory Requirement: Each host stores six blocks, requiring 6bch bytes of memory, and the minimum sequence length is s = 6c.The six blocks include query, current key-value, received key-value, and attention-feedforward output storage.
  • Algorithm and Implementation: The method is compatible with existing memory-efficient Transformer code by combining local blockwise computation with overlapped key-value communication.

4 Setting

The evaluation benchmarks Ring Attention across LLaMA model sizes, Transformer memory baselines, and GPU and TPU configurations. Experiments measure maximum sequence length and model FLOPs utilization under full gradient checkpointing and full precision.

  • Evaluation measures maximum sequence length and model FLOPs utilization.
  • Model Configuration: The study evaluates LLaMA models with 3B, 7B, 13B, and 30B parameters.
  • Baselines: Baselines include vanilla Transformers, memory-efficient attention, efficient CUDA attention, and memory-efficient attention plus feedforward.
  • Training Configuration: Experiments span a single 8-GPU DGX A100, distributed 32 A100 GPUs, and TPUv3, TPUv4, and TPUv5e systems.
  • Training Configuration: All methods use full gradient checkpointing for attention and feedforward, and results use full precision rather than mixed precision.

5 Results

Experiments evaluate Ring Attention’s context-length scalability, MFU, and reinforcement-learning performance against memory-efficient Transformer baselines. Across these settings, Ring Attention supports substantially longer contexts while maintaining performance and improving RL returns.

  • Evaluation setup: Ring Attention is evaluated against vanilla Transformers, memory-efficient attention, and blockwise parallel Transformers across model sizes and accelerator configurations.The evaluation measures maximum supported sequence length, model flops utilization, and throughput.
  • Maximum context size: Over 1 million tokens on 32 A100 GPUs for a 7B model represents a 32× improvement over the previous best context size.On TPUv4-512, the method enables a 256× increase and training sequences exceeding 30 million tokens.
  • Maximum context size: Ring Attention scales context size linearly with device count: a context size s on n GPUs becomes ns under the method.The reported improvements include 8× on 8 A100 GPUs and 256× on TPUv3-512.
  • Model flops utilization: Ring Attention is expected to have lower MFU than baselines because longer contexts increase self-attention FLOPs, whose utilization is lower than feedforward utilization.The method nevertheless maintains MFU while supporting longer contexts.
  • Model flops utilization: Ring Attention maintains MFU while enabling significantly longer context lengths, with large models trained on contexts over 4M tokens with negligible overheads.The MFU evaluation covers model sizes from 7B to 65B.
  • In-context RL performance: Across six ExoRL tasks, AT with Ring Attention achieves a total average return of 113.66 versus 111.13 for AT with BPT.The Ring Attention model consistently outperforms or matches original AT with BPT when conditioning on more trajectories.

6 Related Work

Related work reduces Transformer memory through exact blockwise computation and sequence parallelism, while Ring Attention combines blockwise computation with overlapped ring communication. This design targets larger contexts without communication or computation overheads.

  • Memory-efficient Transformers: Prior methods compute exact self-attention block by block, reducing memory overhead without approximating attention.These methods include memory-efficient attention, efficient CUDA implementations, and blockwise parallel Transformers.
  • Ring Attention: Ring Attention distributes blockwise attention and feedforward across devices while overlapping key-value communication with query-key-value computation and feedforward.The communication occurs circularly among hosts.
  • Ring communication: Earlier ring-topology sequence parallelism reduces communication cost relative to standard sequence parallelism, but communication-computation overlap remains difficult.The cited work is described as infeasible for large-context training and inference because of communication overheads.
  • Ring communication: Ring Attention is presented as the first application of ring communication to self-attention that fits efficiently into Transformer training and inference without added overheads.The paper distinguishes this use from prior ring communication applications in other parallel-computing scenarios.

7 Conclusion

The paper concludes that Ring Attention reduces Transformer memory requirements by distributing long sequences across devices, allowing context length to scale linearly with device count. Experiments support training sequences exceeding 100 million tokens without attention approximations, with future applications in long-form modalities and reasoning.

  • Conclusion: Ring Attention reduces Transformer memory requirements and removes the memory bottleneck imposed by individual devices.The conclusion frames this as a memory-efficient approach for state-of-the-art Transformer models.
  • Conclusion: Context length scales linearly with the number of devices while maintaining performance.The approach supports training sequences up to device count times longer than prior memory-efficient Transformers.
  • Conclusion: Sequences exceeding 100 million tokens can be trained without approximations to attention.The conclusion reports effectiveness on language-modeling and reinforcement-learning experiments.
  • Future prospects: The paper identifies potential applications in large video-audio-language models, extended feedback and trial-and-error learning, codebases, scientific sequences, and reasoning from link-gathering data.These are presented as future prospects enabled by near-infinite context.
  • Implementation: The Jax implementation exchanges key-value blocks among hosts using the collective operation jax.lax.ppermute.Forward and backward passes are defined with defvjp.
  • Practical deployment: For large-scale TPU or high-bandwidth GPU-cluster training, the authors recommend combining FSDP for model sharding with Ring Attention for large context.Tensor parallelism can be added when the total batch size is too large.

B.1 Evaluation of context length

The context-length evaluation uses fully sharded tensor parallelism to measure the maximum sequence length achievable in end-to-end training. Ring Attention is used with FSDP to extend sequence length while keeping the total batch size in tokens fixed.

  • Evaluation setup: Maximum supported context length is evaluated with fully sharded tensor parallelism, a configuration widely used in prior end-to-end training.The evaluation sets the same batch size across comparisons.
  • Evaluation setup: Ring Attention extends sequence length while the total batch size in tokens remains the same.The setup therefore compares larger contexts under a fixed token-budget condition.

B.2 Evaluation of MFU

The evaluation uses large token batches on GPUs and TPUs, with separate finetuning of LLaMA-13B at up to 512K tokens on 32 A100 GPUs.

  • 2 million tokens per batch were used on GPUs, versus 4 million tokens per batch on TPUs.
  • Training used FSDP with Jax SPMD and no-saveable gradient-checkpointing policies for attention and the feedforward network.
  • LLaMA-13B finetuning was limited to a 512K-token context on 32 A100 80GB cloud GPUs.The context limit reflected cloud compute budget constraints.

C Inference requirement

Ring Attention performs blockwise attention while rotating key-value blocks among hosts, supporting inference with context lengths that scale with device count without added overhead.

  • Ring Attention applies a blockwise forward pass that scans key-value blocks and returns normalized attention outputs.The implementation accumulates numerators, denominators, and maximum scores before dividing to produce the output.
  • The implementation defines both forward and backward Ring Attention passes, including blockwise gradient computation and key-value rotation.
  • Key-value blocks are exchanged between previous and next hosts using the collective operation lax.ppermute.

D Training FLOPs Scaling of Context Size

Training FLOPs increase with longer contexts, but the per-dataset cost grows subquadratically because the total number of tokens remains fixed and the ratio decreases for larger models.

  • The Figure 5 x-axis reports context length relative to a 4K context, including notation such as 32x(128K).
  • The per-dataset FLOPs ratio is expressed as (6h + s2)/(6h + s1), where h is hidden dimension and s1, s2 are old and new context lengths.
  • 162.6x higher per-dataset FLOPs results when scaling the 170B model from 4K to 10M tokens, despite a 3072x longer context.
  • Approximately 20-40x more FLOPs are required for small models scaled to 1M-token contexts, with higher ratios at 10M and 100M tokens.
Loading 2310.01889v4…