  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.
