RingAttention with Blockwise Transformers for Near-Infinite Context
Authors: Hao Liu
Matei Zaharia
Pieter Abbeel
Venue: ICLR 2024 poster (2024)
Keywords: Language Model
Large Context
Transformers
Long Context Model
Memory Efficiency
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We present an efficient method of computing the standard transformer architecture, enabling effective processing of long contextual information.

Abstract: 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, Blockwise RingAttention, 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.

Published as a conference paper at ICLR 2024
RINGATTENTION WITH BLOCKWISE TRANSFORMERS FOR NEAR-INFINITE CONTEXT
Hao Liu, Matei Zaharia, Pieter Abbeel
UC Berkeley
ABSTRACT
Transformers have emerged as the architecture of choice for many statc-of- the-art AI models, showcasing exceptional performance across a wide range of AI applications. However, the memory demands imposed by Transform- ers limit their ability to handle long sequences, thereby posing challenges in utilizing videos, actions, and other long-form sequences and modali- ties in complex environments. We present a novel approach, Blockwise RingAttention, 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 device count times longer than those achievable by prior memory-effiiien Transformers, without resorting to approximations incurring additional communication and computation overheads. Exten- sive experiments on language modeling and reinforcemnet learning tasks demonstrate the effectiveness of our approach in allowing millions of tokens context size and improving performance.
INTRODUCTION
Transformers (Vaswani et al., 2017) have become the backbone of many state-of-the-art AT systems that have demonstrated impressive performance across a wide range of AT problems. Transformers achieve this success through their architecture design that uses self- attention and position-wise feedforward mechanisms. However, scaling up the context length of Transformers is challenge (OpenAI, 2023), since the inherited architecture design of Transformers, ie. the self-attention has memory cost quadratic in the input sequence length, which makes it challenging to scale to longer input sequences. Large context Transformers are essential for (ackling a diverse array of AI challenges, ranging from processing books and high-resolution images to analyzing long videos and complex codebases. They excel al extracting information from the interconnected web and hyperlinked content, and are crucial for handling complex scientific experiment data. There have been emerging use cases of language models with significantly expanded context than before: GPT-3.5 (Schulman et al., 2022) with context length 16K, GPT-4 (OpenAI, 2023) with context length 32k, MosaicML'ss MPT (MosaicML, 2023) with context length 65k, and Anthropic's Clande (Anthropic, 2023) with context length 100k.
Driven by the significance, there has been surging research interests in reducing mem- ory cost. One line of research leverages the observation that the softmax matrix self- attention can be computed without materializing the full matrix (Milakov and Gimelshein, 2018) which has led to the development of blockwise computation of self-attention and feedforward (Rabe and Staats, 2021; Dao et al., 2022; Liu and Abbecl, 2023b) with- out making approximations. Despite the reduced memory, a significant challenge still arises from storing the output of each layer. This necessity arises from self-attention's inherent nature, involving interactions among all elements (n to interactions). The subsequent layer's self-altention relies on accessing all of the prior layer's outputs. Fail-
Code available. Email: hao liudcs. berkeley.edee

---

Published as a conference paper at ICLR 2024
ing to do so would increase computational costs cubically, every output must be recomputed for each sequence element, rendering it impractical for longer sequences. These components facilitate the efficient cap- ture of long-range dependencies between in- put tokens, and enable scalability through highly parallel computations. To put the memory demand in perspective, even when dealing with batch size of 1, processing 100 million tokens requires over 1000GB of
Contaxt Size
LOOK
memory for modest. model with hidden size of 1024. This much greater than the capacity of contemporary GPUs and TPUs, which typically have less than 100GB of high- bandwidth memory (HBM).
10K
38
78
305
Model Size
Figure 1 Maximum context length under end-to-end large-scale training on TPUv4-1024. Baselines are vanilla transformers (Vaswani et al., 2017). memory efficient transform- ers (Rabe and Staats, 2021), and memory ef- ficient, attention and feedforward (blockwise parallel transformers) (Liu and Abbeel, 2023b). Our proposed approach Blockwise RingAtten- tion allows training up to device count times longer sequence than baselines and enables the training of sequences that exceed millions in length without making approximations nor adding any overheads to communication and computation.
To tackle this challenge, we make a key observation: by performing self-attention and feedforward network computations in blockwise fashion (Liu and Abbeel, 2023b), we can distribute sequence dimensions across multiple devices, allowing concurrent com- putation and communication. This insight stems from the fact that when we compute the attention on a block-by-block basis, the results are invariant the ordering of these blockwise computations. Our method dis- tributes the outer loop of computing block- wise attention among hosts, with each device managing its respective imput block. For the inner loop, every device computes blockwise attention and feedforward operations specific to its designated input block. Host devices form a conceptual ring, where during the inner loop, each device sends copy of its kcy-value blocks being used for blockwise computation to the next device in the ring, while simultaneously receiving key-value blocks from the previous one. As long as block computations take longer than block transfers, overlapping these processes results in no added overhead compared to standard transformers. Our work utilizes blockwise parallel transformers (Liu and Abbeel, 2023b) costs, enabling xero-overhead scaling of context size across tens of millions of tokens during both training and inference, and allowing for the use of an arbitrarily large context size. Since our approach overlaps the communication of key-value blocks between hosts in ring through blockwise computation of transformers, we name it RingAttention with Blockwise Transformers.
substantially reduce memory
We evaluate the effectiveness of our approach on language modeling benchmarks. Our experiments show that RingAttention can reduce the memory requirements of Transformers, enabling us to train more than 500 times longer sequence than prior memory efficient state-of-the-arts and enables the training of sequences that exceed 100 million in length without making approximations to attention. Importantly, RingAttention eliminates the memory constraints imposed by individual devices, empowering the training and inference of sequences with lengths that scale in proportion to the number of devices, essentially achicving ncar-infinite context size.
Our contributions are twofold: (a) proposing a memory efficient transformers architecture that allows the context length scale linearly with the number of devices while maintaining performance, eliminating the memory bottleneck imposed by individual devices, and (b) demonstrating the effectiveness of our approach through extensive experiments.
2

---

Published as a conference paper at ICLR 2024
2 LARGE CONTEXT MEMORY CONSTRAINT
Given input sequences Q, K, E Rxxx where s is the sequence length and the head dimension. We compute the matrix of outputs as:
Attention(Q2 K, V) softmax( QKT
Vd )V,
where soflmax is applied row-wise. Each self-altention sub-layer is accompanied with leedforward network, which is applied lo each position separately and identically. This consists of two linear transformations with ReLU activation in between.
FFN(x) max(0, xW1 + b1) b2.
Blockwise Parallel Transformers. Prior state-of-the-arts have led to substantial reduc- tions in memory utilization, achieved through innovative techniques that enable attention computation without full materialization by computing attention in a block by block man- ner (Rabe and Staats, 2021; Dao et al., 2022; Liu and Abbeel, 2023b). These advancements lowered the memory overhead of attention to 2bsh bytes per layer, where represents the batch size, 8 denoles the sequence length, and hh stands for the hidden size of the model. To further reduce memory usage, blockwise parallel transformer (BPT) (Liu and Abbeel, 2023b) introduced a strategy where the feedforward network associated with each self-attention sub-layer is computed in a block-wise fashion. This approach effectively limits the maximum activation size of feedforward network from 8bsh 26sh. For a more detailed analysis of memory efficiency, please refer to the discussion provided therein. In summary, the stato-of-the-art transformer layer's memory cost of activation is 2bsh.
Large Output of Each Layer. While BPT significantly reduces memory demand in Transformers, it still presents a major challenge for scaling up context length because it requires storing the output of each layer. This storage is crucial due to the inherent nature of self-attention, which involves interactions among all elements (n to n interactions). Without these stored outputs, the subsequent layer's self-attention becomes computationally impractical, necessitating recomputation for each sequence element. To put it simply, for
processing 100 million tokens with batch size of 1 requires over 1000GB of memory a modest model with a hidden size of 1024. In contrast, modern GPUs and TPUs typically provide less than 100GB of high-bandwidth memory (HBM), and the prospects for significant HBM expansion are hindered by physical limitations and high manufacturing costs.
3 RINGATTENTION WITH BLOCKWISE TRANSFORMERS
Our primary objective is to eliminates the memory constraints imposed by individual devices by efficiently distribute long sequences across multiple hosts without adding overhead. To achieve this goal, we propose an enhancement to the blockwise parallel transformers (BPT) framework (Lin and Abbeel, 2023b). When distributing hosts, each host is responsible for running one element the outer loop of blockwise attention corresponding to its designated block, well as the feedforward network specific to that block. These operations do not necessitate communication with other hosts. However, a challenge arises in the inner loop, which involves key-value block interactions that require fetching blocks from other hosts. Since each host possesses only one key-value block, the naive approach of fetching blocks from other hosts results two significant issues. Firstly, it roduces computation ion delay as the system waits to receive the necessary key-value blocks. Secondly, the accumulation of key-value blocks leads increased memory usage, which defeats the purpose of reducing memory cost.
input sequence across different
Ring-Based Blockwise Transformer. To tackle the aforementioned challenges, we leverage the permutation invariance property of the inner loop's key-value block operations. This property stems from the fact that the self-attention between a query block and a group of key-value blocks can be computed in any order, as long the statistics of each block are combined correctly for rescaling. We leverage this property by conceptualizing all hosts as forming a ring structure: host-1, host-2, host-N. As we compute blockwise attention and
3

---

Published as a conference paper at ICLR 2024
(a)
Blockwise FeedForward
Blockwise FeodForward
Key Value Block
Key Value Block
Key Value Block
Blockwise Attention
Blockwise Attention
query block
query block
Device
Device 2
(b)
Blockwisee FeedForward
Blockwise FeadForward
Blockwise FeedForward
Blockwise FeedForward
Query1
Blockwiie Attention
Blockwise Attention
Blockwise Attention
Blockwise Attention
Query Outer Loop
Quary2
Key1
Key2
Key3
Key4
Query3
Valle1
Vallee
Query4
compute, next device
proviaus devios
Key architecture as the original Transformer Figure 2 Top (a): We use the same model but reorganize the compute. In the diagram, we explain this by showing that in a ring of hosts, each host holds one query block, and key-value blocks traverse through ring of hosts for attention and feedforward computations in a block-by-block fashion. As we compute attention, cach host sends kcy-valuc blocks to the next host while reccives key-value blocks from the precoding host. The communication is overlapped with the computation of blockwise attention and feedforward. Bottom (b): We compute the original Transformer block-by-block, Each host is responsible for one iteration of the query's outer loop, while the key-value blocks rotate among the hosts. As visualized, device starts with the first query block on the left; then we iterate over the key-value blocks sequence positioned horizontally, The query block, combined with the key-value blocks, are used to compute self-attention (yellow box), whose output is pass to foedforward network (cyan box).
Loop
feedforward, each host efficiently coordinates by concurrently sending key- value blocks being used for attention computation to the next host while receiving key-value blocks from the preceding host, effectively overlapping transferring blocks with blockwise computation. Concretely, for any host-i, during the computation of attention between its query block and key-value block, it concurrently sends key-value blocks to the next host-(i+ 1) while receiving key-value blocks from the preceding host-(i 1). If the computation time exceeds the time required for transferring key-value blocks, this results in no additional communication cost. This overlapping mechanism applies to both forward and backward passes of our approach since the same operations and techniques can be used. Prior work has also proposed leveraging ring topology to compute full attention (Li al., 2023b), aiming to reduce communication
4

---

Published as a conference paper at ICLR 2024
Table 1 Comparison of maximum activation sizes among different Transformer architectures. Here, b is batch size, h is hidden dimension, is number of head, is sequence length, is block size, the block size (c) is indopendent of the input sequence length (s). The comparison is between vanilla transformer Vaswani et al. (2017), memory efficient attention (Rabe and Staats, 2021; Dao et al., 2022), blockwise parallel transformers (Liu and Abbeel, 2023b), and our proposed approach RingAltention Numbers are shown in byles per layer, assuming 66loat16 precision.
Layer Type
Self- Self-Autention
FeedForward
Total
Vanilla Transformers
2lms²
8bsh
2bhus
Memory Efficient Attention Blockwise Parallel Transformers
2bsh + 4bch
8bsh
8bsh
2bsh
2bsh
2bsh
Blockwise RingAttention
6bch
26ch
6bch
costs. Our work differs by utilizing blockwise parallel transformers substantially reduce memory costs. As we show in the next section, this enables zero- veroooverhead scaling of context size during both training and inference and allows arbitrarily large context size.
Arithmetic Intensity Between Hosts. In order to determine the minimal required block size to overlap transforring with computation, assume that each host has F FLOPS and that the bandwidth between hosts is denoted as B. It's worth noting that interactions only with the immediately previous and next hosts in a circular configuration, thus our analysis applies to both GPU all-to-all topology and TPU torus topology. Let's consider the variables: block size denoted as с and hidden size as d. When computing blockwise self-attention, and keys, and an additional 2dc² FLOPs for multiplying these attention scores by values. In total, the computation demands amount to 4dc² FLOPs. We exclude the projection of queries, keys, and values, as well as blockwise feedforward operations, since they only add compute complexity without any communication costs between hosts. This simplification leads to more stringent condition and does not compromise the validity of our approach. On the communication front, both key and value blocks require total of 2cd bytes. Thus, the combined communication demand is 4cd bytes. To achieve an overlap between communication and computation, the following condition must hold: 4de2/ / 2 4cd/B, B. This implies that the block size, denoted as c, should be greater than or equal to F F/B. Effectively, this means that the block size needs to be larger than the ratio of FLOPs over bandwidth.
approach involves
require 2dc² FLOPs for calculating attention scores using queries
Memory Requirement. A host needs to store multiple blocks, including one block size to store the current query block, two block sizes for the current key and value blocks, and two block sizes for receiving key and value blocks. Furthermore, storing the output of blockwise attention and feedforward necessitates one block size, as the output retains the shape of the query block. Therefore, total of six blocks are required, which translates to 6bch bytes of memory. It's worth noting that the blockwise feedforward network has maximum activation size of 26ch (Liu and Abbeel, 2023b). Consequently, the total maximum activation size remains at 6bch bytes. Table provides a detailed comparison of the memory costs between our method and other approaches. Notably, our method exhibits the advantage of linear memory scaling with respect to the block size c, and is independent of the input sequence length+
Our analysis shows that the model needs to have a sequence length of s = 6c, which is six times the minimal block size. Requirements for popular computing servers are shown in Table 2. The required minimal sequence length (rightmost column) for each host varies between 6K and 10K, and the minimal block size (second-to-rightmmoet column) for each host is around 1K for TPUs and GPUs with high bandwidth interconnect. For GPUs connected via InfiniBand, which offers lower bandwidth, the requirements are more strict. These requirements are easy to meet using blockwise transformers (Liu and Abbeel, 2023b) and standard parallelism such as data and tensor parallelism, which we will show in experiment Section 5.
5

---

Published as a conference paper at ICLR 2024
Table Minimal sequence length needed on each device. Interconnect Bandwidth is the unidirectional bandwidth between hosts, i.c., NVLink InfiniBand bandwidth between GPUs and ICI bandwidth between TPUs. The minimal block size required and minimal sequence length 6c.
FLOPS FLOPS/Bandwide Banddidthh,
Interconnect Bandwidth
Minimal Blocksize
Minimal Sequence Len
Spec Per Host
FLOPS
HBM
(TF)
(GB)
(GB/s)
(xle3)
(x le3)
A100 NVLink
312
80
300
1.0
6.2
A100 InfiniBand
312
80
12.5
24.5
149.5
TPU v3
123
16
112
1.1
6.6
TPU v4
275
32
268
1.0
6.2
TPU v5e
196
16
186
1.1
6.3
Algorithm Large Context Transformors using RingAttention with Blockwise Transformers.
Required: Input sequence æ. Number of hosts N₁. Initialize
Split input sequence into N₂ blocks that each host has one input block. Compule query, key, and value for its input block on each host.
for Each transformer layer do for count = 1 to N₁ 1 do for For each host concurrently. do
Compute memory efficient attention incrementally using local query, key, value blocks.
Send key and value blocks to next host and receive key and value blocks from previous host.
end for end for for For each host concurrently, do
Compute memory efficient feedforward using local attention output. end for end for
Algorithm and Implementation. Algorithm provides the pscudocode of the algo- rithm. RingAttention is compatible with existing code for memory efficient transformers: RingAttention just needs to call whatever available memory efficient computation locally on each host, and overlap the communication of key-value blocks between hosts with blockwise computation. We use collective operation jax. lax. ppermute to send and receive key value blocks between nearby hosts. A Jax implementation is provided in Appendix A.
SETTING
We evaluate the impact of using RingAttention in improving Transformer models by bench- marking maximum sequence length and model flops utilization.
Model Configuration. Our study is built upon the LLaMA architecture, we consider 3B, 7B, 13B, and 30B model sizes in our experiments.
Baselines. We evaluate our method by comparing it with vanilla transformers (Vaswani et al, 2017) which computes self-attention by materializing the attention matrix and computes the feedforward network normally, transformers with memory efficient attention (Rabe and Staats, 2021) and its efficient CUDA implementation (Dao et al., 2022), and transformers with both memory efficiont attention and feodforward (Liu and Abbeol, 2023b).
Training Configuration. For all methods, apply full gradient checkpointing (Chen et al., 2016) to both attention and feedforward, following prior works (Rabe and Staats, 2021; Liu and Abbeel, 2023b). The experiments are both GPUs and TPUs. For GPUs,

---

Published as a conference paper at ICLR 2024
consider both single DGX A100 server with 8 GPUs and distributed 32 A100 GPUs. We also experiment with TPUs, from older generations TPUv3 to newer generations of TPUv4 and TPUv5e. We note that all of our results are obtained using full precision instead of mixed precision.
5 RESULTS
In our experiments, our primary objective is to comprehensively evaluate the performance of RingAttention across multiple key metrics, including maximum supported sequence length within accelerator memory, model flops utilization, and throughput. We compare RingAtten- tion's performance with several baseline models, including the vanilla transformers (Vaswani et al., 2017), transformers with memory efficient attention (Rabe and Staats, 2021), and transformers with both memory efficient attention and feedforward (Liu and Abbeel, 2023b), across different model sizes and accelerator configurations.
5.1 EVALUATING MAX CONTEXT SIZE
We evaluate maximum supported context length using fully sharded tensor parallelsim (FSDP) (Facebook, 2023) which is widely used in prior end-to-end training (Touvron et al., 2023; Geng and Liu, 2023). We note that no tensor parallelism is considered in our evaluations since our approach independent of tensor parallelism. Practitioners can combine our method with tensor parallelism, which we will show in Section 5.2. Using FSDP allows us to set the same batch size in tokens for basclines and our approach, ensuring a fair comparison. Concretely, on devices, FSDP is used to shard the model for basclines, which gives sequence length of The total batch size tokens is nl. We utilize FSDP along with RingAttention to extend the sequence length to al and m sequences. This means that the lotal batch size in lokens remains the same, but RingAttention enables a siguificantly larger context size. Table 3 summarizes the results of our experiments.
za
Our RingAttention model consistently surpasses baselines, delivering superior scalability across diverse hardware setups. For example, with 32 A100 GPUs, achieve over 1 million tokens in context size for 7B model, a 32 times improvement over previous best. Furthermore, when utilizing larger accelerators like TPUv4-512, RingAttention enables a 256 times increase in context size, allows training sequences of over 30 million tokens. Furthermore, our Ring Attention model scales linearly with the number of devices, as demonstrated by the 8x improvement over previous best on 8 A100 and the 256x improvement on TPUv3-512. If model can be trained with context size s on n GPUs using the blockwise attention and feedforward, with our RingAttention approach, it becomes possible to train a model with context size of ns.
5.2 EVALUATING MODEL FLOPS UTILIZATION
We evaluate the model flops utilization (MFU) of Ritgnttention in standard training settings using fully sharded data parallelism(FSDP) (Facebook, 2023) and tensor parallelism following LLaMA and OpenLLaMA (Touvron et al., 2023; Geng and Liu, 2023) with Jax SPMD. The batch size in tokens are 2M on 8/32x A100 and 4M on TPUv4-256. Our goal is investigating the impact of model size and context length on MFU, critical performance metries while highlighting the benefits of our approach. Table 4 presents the results of our experiments on MFU for different model sizes and context lengths. We present the achieved MFU using state-of-the-art. memory elficient transformers BPT (Liu and Abbeel, 2023b), compare it to our anticipated MFU based on these results, and demonstrate the actual MFU obtained with our approach (RingAttention). For fair comparison, both BPT and our approach are based on the same BPT implementation on both GPUs and TPUs.
KingAttentiog trains much longer context sizes for self-attention, resulting in higher self- attention FLOPs compared to bascline models. Since self-attention has a lower MFU than feedforward, Aitgnitenntonn is expected to have lower MFU than the baseline models. Our method offers a clear advantage in terms of maintaining MFU while enabling training with significantly longer context lengths. As shown in Table 4, when comparing our approach
7

---

Published as a conference paper at ICLR 2024
Table 3 The maximum context length supported in end-to-end training using fully sharded data parallelism and various transformers architectures. We show different model sizes and accelerators. Basclines are vanilla transformer (Vaswani al., 2017), transformer with memory efficient attention (Rabe and Staats, 2021), and transformer with memory efficient attention and feedforward (Liu and Abbeel, 2023b). The context size reported in tokens (1e3). Our approach Blockwise RingAutention substantially pulperforms baselines and enables training sequences that are up to device count times longer than prior state-of-te tate-off-hh--sss.....
Max context size supported (xle3)
Memory Efficient Attn
Memory Efficient RingAttention (Ours)
Ours SOTA
Vanilla
Attn and FFN
8x A100
NVLink
3B
32
64
512
8x
7B
2
16
32
256
8x
13B
2
128
32x A100
InfiniBand
7B
4
128
4096
32x
13B
4
32
64
2048
32x
TPUv3-512
7B
1
4
8
2048
256x
13B
1024
128x
TPUv4-1022
3B
8
16
32
16384
512x
7B
4
16
8192
512x
13B
4
8
16
4096
256x
30B
2
4
8
2048
256x
TPUv5e-256
3B
32
4096
128x
7B
2
8
16
2048
128x
to prior state-of-the-arts. it is evident that we can train very large context models without compromising MFU or throughput.
5.3 IMPACT ON IN CONTEXT RL PERFORMANCE
We present results of applying RingAttention for learning trial-and-error RL experience using Transformers. We report our results in Table 5, where we evaluate our proposed model on the ExoRL benchmark across six different tasks. On ExoRL, we report the cumulative return, per ExoRL (Yarats et al., 2022). We compare BC, DT (Chen et al., 2021), AT (Liu and Abbeel, 2023a), and AT with memory efficient attention (Rabe and Staats, 2021) (AT+ME), AT with blockwise parallel transformers (Liu and Abbeel, 2023b) (AT+BPT), and AT with our RingAttention and AT paper. AT RingAttention numbers run by ourselves. Since the ExoRL data is highly diverse, having been collected using unsupervised RL (Laskin el. al., 2021), it has been found that TD learning performs best, while behavior cloning struggles (Yarats et al., 2022). AT (Liu and Abbeel, 2023a) shows that conditioning Transformer on multiple trajectorics with relabeled target return can achieve competitive results with TD learning. For more details, please refer to their papers. We are interested in applying Ring Attentttotion to improve the performance than 32 trajectories in prior works. It is worth noting that each trajectory has 1000 x 4 length where 1000 is sequence length while is aituu--stato-actio---rtttt 128 trajectories with modest 350M size model infeasible for prior state-of-the-ar blockwise parallel transformers. Results in Table show that, by scaling up the sequence length (number of trajectories), AT + Ring Attention consistently outperforms oringal AT with
Tnn-RingAttentee
The numbers of BC, DT, AT are from the ExoRL
AT by conditioning on a larger number of trajectories rather
making training

---

Published as a conference paper at ICLR 2024
Table 4
Model flops utilization (MFU) with different training configurations: model sizes, compute, and context lengths. RingAttention enables training large models (7B-65B) large input context sizes (over 4M) with negligible overheads.
Memory Efficient Attention &
Blockwise RingAttention (expected)
Blockwise RingAttention
35
30
MFU (%)
25
5
0
7B on A100-8
13B on A100-8
138 on A100-32
30B on TPUv4-1024
65B on TPUv4-1024
32K 256K
16K 128K
64K 2048K
16K 2048K
8K vs 1024K
Training config
Model size
7B
13B
13B
30B
65B
Compute
8x
A100
8x A100
32x
A
TPUv4-1021
TPUv1-1021
Memory efficient
Context size (x1e3)
16
8
attention & FFN
Blockwise
Context size
256
128
2048
2048
1024
Ring Attentttn
(xle3)
Vicuna-138(16K)
OpenAl GPT3.5(16K)
Anthropic Claude2(100k)
RingAttention 138/512K)
0.8
Accuracy (%)
0.6
0.4
0.2
8K
12K 16K
64K
100K
200K
500K
Context Len
Figure Comparison of different models on the long-range line retricval task.
BPT across all six tasks, achieving a total average return of 113.66 compared the AT with BPT model's total average return of 111.13. The results show that the advanlage of RingAttention for training and inference with long sequences.
Table 5 Application of RingAttention improving Transformer in RL. BC and DT use vanilla attention. AT ME denotes using memory efficient attention, AT BPT denotes using blockwise parallel transformer. AT RA denotes using RingAttentton.
ExoRL
BC-10%
DT
AT + ME AT + BPT AT + BPT AT + RA
Task
N Trajs
32
N Trajs
32
N Trajs
128
Trajs
128
Walker Stand
52.91
34.54
00m
95.45
o0m
98,23
Walker Run
34.81
49.82
00m
105.88
00m
110.45
Walker Walk
13.53
34.94
00m
78.56
00m
78.95
Chectah Run
34.66
67.53
00m
178.75
00m
181.34
Jaco Reach
23.95
18.64
87.56
89.51
Cartpole Swingup
56.82
67.56
120.56
123.45
Total Average
36.11
45.51
oom
111.13
00m
113.66
9

---

Published as a conference paper at ICLR 2024
5.4 IMPACT ON LLM PERFORMANCE
We evaluate RingAttention by applying our method to finetune LLaMA model to longer context. In this experiment, while our approach enables training with millions of context tokens, we conducted finetuning on the LLaMA-13B model, limiting the context length to 512K tokens due to constraints on our cloud compute budget. This finetuning was carried out on 32 A100 GPUs, using the ShareGPT dataset, following methodologies as outlined in prior works (Chiang et al., 2023; Geng et al., 2023). We then evaluated our finetuned model on the line retrieval test (Li et al., 2023a). In this test, the model needs to preciscly retrieve a number from a long document, the task can effectively capture the abilities of text generation, retrieval, and information association at long context, reflected by the retrieving accuracy. Figure 3 presents the accuracy results for different. models across varying context lengths (measured in tokens). Notably, our model, Ring B 13B-512K, stands out as maintains high accuracy levels even with long contexts. GPT335--trtte 16K, (reuna-16B-Lan 16K, and Claude 2-100K demonstrate competitive accuracy within short context lengths. However, they cannot handle extended context lengths.
6 RELATED WORK
Transformers have garnered significant attention in the field of AI and have become the backbone for numerous stale-of-the-art models. Several works have explored memory-efficient techniques to address the memory limilations of Transformers and enable their application to wider range of problems. Computing exact. self-attention in blockwise manner using the tiling technique (Milakov and Gimelshein, 2018) has led to the development of memory efficient attention mechanisms (Rabe and Staats, 2021) and its efficient CUDA implementation (Dao et al., 2022), and blockwise parallel transformer (Liu and Abbeel, 2023b) that proposes computing both feodforward and self-attention block-by-block, resulting in significant reduction in memory requirements. In line with these advancements, our work falls into the category of memory efficient computation for Transformers. Other works have investigated the approximation of attention mechanisms, yet these efforts have often yielded sub-optimal results or encountered challenges during scaling up. For an in- depth review of these techniques, we recommend referring to the surveys (Narang et al., 2021; Tay et al., 2022). Another avenue of research explores various parallelism methods, including data parallelism (Dean et al., 2012), tensor parallelism (Shoeybi et al., 2019), pipeline parallelism (Narayanan et al., 2019; Huang et al., 2019; Narayanan et al., 2021), sequence parallelism (Li et al., 2023b; Korthikanti et al., 2022; Jacobs et al., 2023), and FSDP (Facebook, 2023; Rajbhandari et al., 2020). The activations of self-attention take a substantial amount of memory for large context models. Tensor parallelism can only reduce parts of activations memory and sequence parallelism introduces a significant communication overhead that cannot be fully overlapped with compulation ion. Prior work has studied sharding along sequence and attention heads, and gathering sequences via an optimized all-to-all topology, achieving reduced communication (Jacobs et al., 2023). However, this method is restricted by the number of attention heads and requires gathering the full sequence each device. In comparison, our approach fully overlaps communication with blockwise computation, enhancing its scalability. Prior work study sequence parallelism for computing self-attention using & ring topology (Li et al., 2023b), but is not optimized for blockwise parallel transformers and is incompatible with memory-efficiet atlention, which are crucial to large context training. Overlapping communication with computation remains challenging, and the communication overheads make it infeasible for training and inference in large-context scenarios. Our work leverages on blockwise parallel transformers to distribute attention and foedforward across devices and concurrently overlaps the communication of key- vayue blocks in circular of hosts with the computation of query-key-value blocks and feedforward, reducing memory cost substantially and allowing device count times larger context size with zero overheads. Overlapping communication with computation has been studied in high performance computing literature (Danalis et al., 2005; Wang et al., 2022; Danalis et al., 2009, inter alia). While ring communication has found applications in other parallel computing scenarios (Bischof, 2008; Hursey and Graham, 2011; Gibiansky, 2017; Sergeev and Del Balso, 2018), our work stands out as the first work to show that it can be applied to
10

---

Published as a conference paper at ICLR 2024
self-attention as used in Transformers and to make it fit efficiently into Transformer training and inference without adding significant overhead by overlapping blockwise computation and communication.
CONCLUSION
In conclusion, we propose a memory efficient approach to reduce the memory requirements of Transformers, the backbone of state-of-the-art Al models. Our approach allows the context length to scale linearly with the number of devices while maintaining performance, eliminating the memory bottleneck imposed by individual devices. Through extensive experiments on language modeling and reinforcement learning, we demonstrate its effectiveness, enabling training sequences that are up to device count times longer than those of prior memory-effccien Transformers, exceeding context length of 100 million without. making approximations to attention. In terms of future prospects, the possibility of ncar-infinite context introduces vast array of exciting opportunities, such large aao--aaddoo-anuaaee models, learning from extended feedback and -erroos, understanding and generating codebase, adapting Al models to understand scientific data such gene sequences, and developing strong reasoning from link gathering data.
ACKNOWLEDGMENTS
This project is supported in part by Office of Naval Research grant N00014-21-1-2760. We express our gratitude to the BAIR and RLL communities for their insightful discussions and feedback. We are also thankful David Patterson for addressing our questions about TPUs and giving insightful feedback on early versions of this work. Our appreciation goes out to Yash Katariya and Sharad Vikram from the Jax developers' team for assisting with our Jax related questions. We also thank Tri Dao for the valuable feedback on this work. We thank Google TPU Research Cloud for granting us access to TPUs.
REFERENCES
Anthropic. Introducing claude, 2023. URL http:://ww..attbrooicccccmm//m introducing-cl.n claude.
Christian Bischof. Parallel computing: Architectures, algorithms, and applications, volume 15. IOS Press, 2008
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:1877-188 1901, 2020.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34: 15097, 2021.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint ar 1rx::16004.00 74, 2016.
Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E Gonzalez, et al. Vicuna: An open-source chatbot impressing gpt-4 with 90%* * chatgpt quality. See https:/ /vw.//na..sssee
2023.
Anthony Danalis, Ki- Yony Kim, Lori Pollock, and Martin Swany. Transformations to parallel codes for commuication-compttaat overlap. In SC'05: Proceedings of the 2005 A CM/IEEF conference Supercompuling, pages 58 58. IEFF, 2005.
11

---

Published as a conference paper at ICLR 2024
Anthony Danalis, Lori Pollock, Martin Swany, and John Cavazos. Mpi-aware compiler optimizations for improving communication-comttaatii overlap. In Proceedings of the 23rd international conference Supercomputing, pages 516-325, 2009.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neurol Information Processing Systems, 35:16344-163599 2022.
Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc'aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. Large scale dis- tributed deep networks. Advances in neural information processing systems, 25, 2012.
Facebook. Fully Sharded Data Parallel: faster Al training with fewer GPUs engineer- ing.fb.com. bttps://engineerrt
fb. comm/2021007115/ppe http://////ppeoouueeeeeeee
2023.
Xinyang Geng and Hao Liu. Openllama: An open reproduction of llama, may 2023. URL https://github. com/ hom/opentm-researell /opnn Ilama, 2023.
Xinyang Geng, Arnav Gudibande, Hao Liu, Eric Wallace, Pieter Abbeel, Sergey Levine, and Dawn Song. Koala: A dialogue model for academic research. Blog post, April, 1, 2023.
Andrew Gibiansky. Bringing hpce techniques to deep learning. Baidu Research, Tech. Rep., 2017.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoe V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in newral information processing systems, 32, 2019.
Joshua Hursey and Richard T. Graham. Building a fault tolerant. mpi application: A ring communication example. In 2011 IEEE International Symposium on Parallel and Distributed Processing Workshops and Phd Forum, pages 1549 1556. IEEE, 2011.
Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Leon Song, Samyam Rajbhandari, and Yuxiong He. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models, ar preprint ar Xxxv:2309.14509 2023.
Vijay Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. arXiv preprint arXiv:22050051: 2022.
Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang, Lerrel Pinto, and Picter Abbeel. Urlb: Unsupervised reinforcement learning benchmark. 2021.
XrXi preprint
Dacheng Li, Rulin Shao, Anze Xie, Ying Sheng, Lianmin Zheng, Joseph E. Gonzalez, lon Stoica, Xuczhe Ma, and Hao Zhang. How long can open-source 1lms truly promise on context length?, June 2023a. URL https://Imsy:: orgo/og////oo 20110//222300-225 19-longh
Shenggui Li, Fuzhao Xue, Chaitauya Baranwal, Yongbin Li, and Yang You. Sequence parallelism: Long sequence training from system perspective. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki, editors, Proceedings of the 61st Annual Meeting of the Association for Compulational Lánguistics (Volume 1: Long Papers), pages 2391- 2404, Toronto, Canada, July 2023b. Association for Computational Linguisties. doi: 18653/ 1/22223.ac/ acl-alln..n... 134. URL https: //cpatto//cl org/2023. acl acl-long. 134. 134.
Hao Liu and Pieter Abbeel. Emergent agentic transformer from chain of hindsight experience. International Conference on Machine Learning, 2023a.
Hao Liu and Pieter Abbeel. Blockwise parallel transformer for large context models. Advances in newral information processing systems, 2023b.
12

---

Published as a conference paper at ICLR 2024
Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint arXiv: 02867,, 2018.
MosaicML. Introducing mpt-7b: A new standard for open-source, commercially usable llms, 2023. URL https://wwa mosai cml. cowmoo///////. commbblg/mptttt/e
Sharan Narang, Hyung Won Chung, Yi Tay, William Fedus, Thibault Fevry, Michael Matena, Karishma Malkan, Noah Fiedel, Noam Shazeer, Zhenzhong Lan, et al. Do transformer modifications transfer across implementations and applications? arXiw preprint arXiv:2102. 1972, 2021.
Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R. Devanur, Gregory R Ganger, Phillip B Gibbons, and Matei Zaharia. Pipedream: Generalized pipeline parallelism for dnn training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles, pages 1 15, 2019.
Doepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. Memory- efficient pipeline-parallel dnn training. In International Conference on Machine Learning, pages 7937-7947. PMLR, 2021.
OpenAL. Gpt-4 technical report, 2023.
Markus N Rabe and Charles Staats. Sclf-attention docs not need o(n2) memory. arXi preprint orXiv:2112.0500 2021.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-16. IEFF, 2020.
J. Schulman, B. Zoph, C. Kim, J. Hilton, J. Menick, J. Weng, J. F. C. Uribe, L. Fedus, L. Metz, M. Pokorny, R. G. Lopes, S. Zhao, A. Vijayvergiya, E. Sigler, A. Perelman, C. Voss, M. Heaton, J. Parish, D. Cummings, R. Nayak, V. Balcom, D. Schnurr, T. Kaftan, C. Hallacy, N. Turley, N. Deutsch, and V. Goel. Chatgpt: Optimizing language models for dialogue. OpenAl I Blog, 2022. URL https:// mtmpenaii.cmm/iiie cww/blog////////occcooooooooooo. chatgpt.
Alexander Sergeev and Mike Del Balso. Horovod: fast and easy distributed deep learning in tensorflow. Xiv preprint arxiv: AKK2:1802005110 2018.
Mohammad Shocybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Calanzaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint a 1909. 009.08053, 2019.
Yi Tay, Mostafa Dehghani, Samira Abnar, Hyung Won Chung, William Fedus, Jinfeng Rao, Sharan Narang, Vinh Q Tran, Dani Yogatama, and Donald Metzler. Scaling laws vs model architectures: How does inductive bias influence scaling? arXie preprint ar arxix:2207. 10551, 2022.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXio preprint ar Xrv:22300 13971, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkorcit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Shibo Wang, Jinliang Wei, Amit Sabne, Andy Davis, Berkin Ilbeyi, Blake Hechtman, Dehao Chen, Karthik Srinivasa Murthy, Marcello Maggioni, Qiao Zhang, et al. Overlap communication with dependent computation via decomposition in large deep learning models. In Proceedings of the 28th ACM International Conference Architectural Support for Programming Languages and Operating Systems, Volume 1, pages 93-106, 2022.
Denis Yarats, David Brandfonbrener, Hao Liu, Michael Laskin, Pieter Abbeel, Alessandro Lazaric, and Lerrel Pinto. Don't change the algorithm, change the data: Exploratory data. for offline reinforcement learning. ar preprint or Xiv:2201. 13425, 2022.
13

---

Published as a conference paper at ICLR 2024
A CODE
The implementation of RingAAteention in Jax is provided in Figure 4. We use defvjp function to define both the forward and backward passes, and collective operation jax. jax. wox.1xxxppeemmen ppermute to facilitate the exchange of key-value blocks among a ring of hosts. The provided code snippet highlights essential components of RingAttention. We provide the complete code on github.
For large scale end-to-end training TPU or on GPU cluster with high bandwidth inter connection, we recommend using FSDP to shard large models and using RingAttention to achieve large context. If total batch size is too large, add tensor parallelism to reduce the global batch size. The degree of parallelism can be adjusted using the mesh dim paramcter within the codebase. To illustrate, consider a sctup with 512 devices, such as 512x A100. If the model size is 30B, you shard it across 8 devices and allocate the remaining 32 devices for RingAttention. This setup allows the context size to be expanded 32 times more than if you didn't use RingAllenticg need for FSDP. This means you can utilize all 512 devices exclusively to expand the context using RingAttention by 512 times. Building upon the result that our approach allows for 256K context size when using 8x A100 GPUs, suggests that by employing 512 A100 GPUs, the potential context size be expanded to 16 million.
Conversely, for models sized 7B 3B, there is no
B EXPERIMENT DETAILS
B.1 EVALUATION OF CONTEXT LENGTH
In the experimental results presented in Section 5.1, we used fully sharded tensor parallelism (FSDP) to partition the model across GPUs or TPU devices. Our evaluation focused determining the maximum achievable sequence length in commonly used FSDP training scenarios. For TPUs, we utilized its default training configuration, which involved performing malmul operations in boattil 16 format with weight accumulation in float32. On the other hand, for GPUs, we adopted the default selup, where all operations were performed in float32.
B.2 EVALUATION OF MFU
In the evaluation presented in Scction 5.2. The batch size in tokens million per batch on GPU and 4 million per batch on TPU. The training was conducted using FSDP Facebook (2023) with Jax SPMD. For gradient checkpointing (Chen et al., 2016), we used noth- ing saveable as checkpointing policies for attention and feedforward network (FFN). For more details, please refer to Jax documentation.
B.3 EVALUATION ON LINE RETRIEVAL
In the evaluation presented in Section 5.4, we finetuned the LLaMA-13B model (Touvron et al., 2023), limiting context length to 512K tokens due to constraints on our cloud compute budget, the training was conducted on 32x A100 80GB Cloud GPUs. We use user-shared conversations gathered from ShareGPT .com with its public APIs for finetuning, following methodologies outlined in prior works (Chiang et al., 2023; Geng et al., 2023). ShareGPT is a website where users can share their ChatGPT conversations. To ensure data quality, we convert the HTML back to markdown and filter out some inappropriate or low-quality samples, which results in 125K conversations after data cleaning.
C INFERENCE REQUIREMENT
We provide the minimal sequence length required to overlap communication with computation during training in Table 2. Our Blockwise RingAttention enables effortless training of context size that scales linearly with the number devices. While we focus on introducing training as it is more memory demanding than autoregressive inference where the number of query
14

---

Published as a conference paper at ICLR 2024
def _rring_attention_fddle k, bias, axis_name, flloat32_ogiits, blockvise_kwargs) if float332loggits:
3
9, k q. q.astype(jnn (jo..flaat3)) float32), astyype wttp.ffaaa....... batch, g_len, num_heade dim_per_ine _ead q.shape batch, kv xv_len, num sum_heads, dim_per_heet _eadd = k. shape
numerator jup. zeros( g_len, num_heads, (er_per_hedd) astype dtype)
demominator jmp zeros num_heads, q_len)) astype dtype) axis_size lax. psumm1, axis_name)
block_size q_len # assumes this function is pre-sharded inside shard query stry__chuk_ssie size blockvise (atrrisse_swsss_ (queryy chunk_sszee" bey_chunk_size blockwise_ wetpwissekwargslloo ("key_ def scan_kv_blockkcarr)) prev
""arrrs("kky_ccuu_sssze""] chunk_size"] idx):
12
13
prev_max_score numerator, denominator, k, v carry
attn_ attn_bias_alice = lax lax. ltt.__siieeeeiiaaaarreeeeeeeeeee
(lax.axis_ www._iiieidd_d_eeeeeeeeee (aaxi__amme) idx) % axis_size kv_len, kv sv_len, axis=-1)
b_block_ixx lax. aris_indees (axis_name)
17
k_block_idz (lax axis axis_index((_ee (axis (amis_nnmme idx) axis_size
18
t_chunk_idx_stene = q_block_idx (block_size //
k_ t_chunk_idx_stare = k_block_idx
fblock_size // key_ feey_chunkssiee
numerator, denominator, max_score _blockwise__tte tttpii/s_aaeetton_fudeeee
V,
(numerator, denominator, prev_ prev_max_score, p.chunk_ixx_start k_chunk bauukkidx_staree btlasattt__baas_slaa
22
##blockwise_kwarrg
k, v = map(lambda x: lax pax.ppperuute (x, axis_name, perm= (i 1) % axis_size) for in raxissaxis_sizel) return (max_ (max_score, numerator, denominator, v), None prev_max_seee (max_score, numerator, denominator, 2), lax. scan (scanssan_Kk brccan_kybblock init=(prev_ satev_ma__se_ maee___e__scoe, numerator, denominator, k, v), xs= output numerator rearrange (denominator,, 'b h q return output aetpue.attype(( dt.dt)))) (output, q, k, v, attn ittn_bias, denominator, max_score)
(k, v))
jup.full((bttch,, num_heads, q_len), -jnp. inf) .af)).ssttp((g astype(q.dtypar
27
28
init=(prev_mxx_score,
axis_size))
q h') h.),..,, None]
32
def sing.atttantion_br((riiii _ame,, float32_logits, blockwise_kwarge res, g): output, 9, k, sttn_bias, denominator, max_score = res batch, kv_len, num_heads, dim_per_heat k. shape axis_size lax (sump,, axis_name)
36
dq aap.zerrs__1nn zeros_like dtype-jnp
jnee-joaffoo.3. float32)
37
dk = 2eoozeroa___11
dv = jnp.zeros. beros_like(t dtype=jnp float32) jtype==jap.float33)
query setry_chunk_aste size blockwise blockwise_kwarra aeguur_ccuukkiiiee. size"]
key bey_chhun_size blockwise (wwyrgs"key_chuuti
41
block_size shapee[1] assumes this function is pre-sharded inside shard_map idx):
def ccan_kv_block(cane dq, dk, dv, k, carry
attn_bias_ slice lax. www.aicnnsllliiicm slicicccs__i____ee in_ dim dim bbas,
45
(lax. aax__xxi_index index (axis_ (axis_name) idx) % axis_size kv_len, kv wv_len, axis=-1)
q. blockc__ise idx = lax. axis index (axisexxxxx__axxxxx _ammee)
k_ x_block_idz = (lax cecis_iidexxaxii__eeee
idx) % axis_size
c_chunk_ix__ttt
=q_ q_block idx
(block_size
query_chunk...
49
k_chunk_idx_start k b_block_idr
(block_size // key_chunn____ (chunk_size)
50
d www.ooovlseeeeeee...ee dq, dk, dv L_chunk_idx_starte http://ss_atttetton_bbieeeee bwd(c, k, v, g, (dq, dk, dv, output, denominator, max_score),
k_chunk__ _da_sttrt,, bias-attn_a sestn_baa_ssh_s bias_slice, *eblockvise_kwarge
k, dk, dv map (lambda lax. ppermute axis name, perm [(1,
1) % axis_sizes for return (dq, dk, dv, k, v), None
range(axis_aize)]] (k, v, dk, dv))
55
(dq, dk, dv, k, v) = lax. acan (scan block, init= dk, dv, k, v), xs=jnp arange(0, axis_size))
dq, dk, dv dq. astype (q dtype), dk astype (k. dtype), dv as (ytypeeyydtypp dtype) return dq. dk, dv, None
59
@partial custom_ vip, nondiff. Saciiff_rrgumss(C [4, 5, 6])
def ring attentionno 3, v, bias, axis_name Float32_logits, fecdkwise_wwargee y. ainggatttetioon_uuude k, v, attn_bias, axis name_ float32_logits, blockwise_kvarge return y
ring_atie hing.atrenicoo.eeerrpprr.r...ate wwt.:/ip..iigg.attettioooo.
.ring_attention_baa))
Figure 4 Key parts of the implementation in Jax. We use collective operation lax. ppermute to send and receive key value blocks between previous and next hosts. The full code is implemented in Jax and Pallas for best performance.

---

Published as a conference paper at ICLR 2024
token is one, RingAttention is applicable to inference too. For example, serving LLaMa 7B on 32x TPUv5e, the conventional approach to distribute the model along the attention heads dimension, with each device computing of 1, this can serve up to a 256K context length due to key-value cache activation size. RingAttention allow 32 times larger context by circulating the key-value cache between ring of devices. To overlap the communication with computation, it needs d2/F >= 2 where B/F >-2. With a bandwidth of 186 GB/s and flops of 196 TFLOPs, and assuming an unrcasonably high MFU of 40% for this large context, then B/F 2.4, mcaning that RingAttention allows 32 times larger context for inference without adding overheads.
attention head. Assuming batch size
D TRAINING FLOPS SCALING OF CONTEXT SIZE
Given that our proposed approach unlocks the possibility of training with context size exceeding 100 million tokens and allows for linear scaling of the context size based on the mumber of devices, it is essential to understand how the training FLOPs per dataset scale with the context size. While larger context size results in higher number of FLOPs, the increased ratio does not scale quadratically because the number of tokens remains fixed. We present these results in Figure which showcases various model sizes and context lengths, representing different computational budgets. The figure shows the ratio of FLOPs for larger context lengths compared the same model with shorter 4K context size. We calculated the
5000
1TB
1.1
1.3
1.6
56.8
596.8
1000
1758
1.1
1.2
1.4
1.8
2.6
4.3
14.4
162.6
1725.6
Model Size
65B
1.1
1,2
1.5
2.2
3.4
5.8
20.6
237.2
2521.5
2.3
268.0
2850.3
138
1.1
1.4
1.8
2.8
4.6
8.4
31.0
362.3
3855.9
7B
2.0
3.1
5.4
10.0
37.4
439.7
4682.0
2x
4x
8x
16x
32x
64x
256x (1M)
3072x
32768x
(16K)
(32K)
(64K)
(128K)
(256K)
(10M)
(100M)
Context Len
Figure The per dataset trainig FLOPs cost. ratio relative to a 4k context size, considering different model dimensions. On the x-axis, you'll find the context length, where, for example, 32x(102 (2x8k) denotes context length of 128k, 32x the size of the same model's 4k context length.
per sequence FLOPs using (24bsh² + 4hs2h)n where model hidden dimension, b batch size, s is total sequence length, and is number of layers. The datasct FLOPs ratio is then given by ((24bsgh² (hasss/))/((abb))
(bbs 81) (6h+ (/h))/2)//-/))// (6h++))))) where 82 and and old context lengths. Model sizes and their hidden dinensions as follows: LLaMA-7B (4096), LLaMA-13B (5140). LLaMA-33B (7168), LLaMA-65B (8192), GPT3- 175B (12288), and 1TB (36864). These model configurations from LLaMA (Touvron et al., 2023) and GPT-3 (Brown et al., 2020) papers, except the 1TB model size and dimension defined by
16

---

Published as a conference paper at ICLR 2024
As depicted in Figure 5, scaling up small models to 1M context size results in approximately 20-40 times more FLOPs, and even more for 10M and 100M token context sizes. However, as the model sizes increase, the cost ratio decreases. For instance, scaling up the 170B model from 4K to 10M incurs 162.6x higher per dataset FLOPs, despite the context size being 3072 times longer.
17


Reviews:
[{"review_id": "zZGbCRRlar", "reviewer": "Reviewer_haGA", "summary": "This paper proposed an efficient distributed computation of full attention mechanism of long sequences across the sequence dimension. By leveraging the blockwise computation of full attention, the proposed Ring Attention distributes long sequences across multiple devices while overlapping the communication of key-value blocks with the computation of blockwise attention. Importantly, Ring Attention enables us to store the output of each layer across multiple devices, significantly reduces the memory demand in Transformer for long sequences.", "questions": "NA.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The challenge the paper aims to address is well-motivated and the proposed method is well-introduced. The experimental results are strong.", "weaknesses": "There are concerns/questions about the proposed Ring Attention:\n\n1. If I understand correctly, the Ring Attention is based on the assumption that the blockwise computation of self-attention on different devices is well-balanced. However, in the causal mode of self-attention, which is widely used in auto-regressive LLMs, this assumption does NOT hold. It is because under the causal mode, only the lower triangular blocks are necessarily to compute, yielding around half of FLOPs of the non-causal mode with full attention. But in Ring Attention, even under causal mode, we still need to compute all the blocks in the full attention matrix. When the sequence is super long, this may require significantly more FLOPs? \n\n2. For the results in Table 3, Ring Attention are not compared with other sequence parallel mechanism, such as Deep Speed Ulysses which leverages all-to-all communication to speed up sequence parallelism.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OFPtXBT579", "reviewer": "Reviewer_dMuy", "summary": "The paper proposes ring attention, which is a way of implementing softmax dense attention.  The attention operation is  \ndistributed over multiple GPUs, and computed in parallel along the sequence dimension.  Attention is also computed in a blockwise fashion, without ever instantiating the full attention matrix, which reduces memory requirements. Together, these two techniques allow dense attention to be scaled to very long sequences.", "questions": "The authors claim that ring attention over very long sequences is suitable for both training and inference.  However, the cost of attention in both cases is still O(N^2).  For training, this is less of an issue; all elements of the (long) sequence are processed in parallel, so the sequence length is essentially just a batch dimension, and the number of tokens per training step is not excessive for an LLM.  \n\nHowever, for inference, tokens are processed autoregressively, so very long sequence lengths will lead to very slow generation times, and thus very high latency when serving the model.  The authors do not discuss or even mention this problem.  How does the latency of ring attention over, e.g., 8M tokens during inference compare to a more conventional short-context model?", "limitations": "", "rating": 3, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "The paper is well-written, and it cites the appropriate literature.  The technique described in this paper does, in fact, allow transformers to be trained with very long sequences.", "weaknesses": "Unfortunately, although the paper is technically sound, it does not really contribute anything new.  The ideas presented\nin this paper will already be obvious to most practitioners.  As a result, I do not feel that it is appropriate for publication at a conference. \n\nThe fact that transformers process all elements of a sequence in parallel, (at least during training) was the whole inspiration for the original 2017 paper that introduced the transformer architecture.  Thus, the idea of distributing the computation in parallel along sequence length is most definitely nothing new.  This paper merely distributes the computation over multiple devices.\n\nA naive implementation of attention instantiates the entire attention matrix, which is too large to fit in memory for long sequences.  The blockwise computation of dense attention avoids such instantiation, and it does involve some subtle implementation details wrt. to the handling of softmax.  However, it is also not a contribution of this paper; the authors cite the appropriate prior work. \n\nThe remaining piece of the puzzle is to overlap the matrix multiplications for each block with the communication overhead of transferring blocks between different hosts.  However, this is also nothing new; tiled matrix multiplication, distributed over multiple GPUs, is implemented by every major machine learning library today (e.g. Jax), and existing implementations overlap matmuls and communication.\n\nThus, the authors merely observe that it is possible to implement long-range attention using a simple combination of well-understood and previously published techniques.  This fact should be obvious to most practitioners, and IMO, does not rise to the level of a conference paper.  I would encourage the authors to publish this work in a workshop or on arXiv.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4bWNdZC4OX", "reviewer": "Reviewer_rUcK", "summary": "The paper proposed a distributed computing strategy for distributing self-attention computation to across multiple devices without introducing latency on host communication. The authors observed that when distributing blockwise computation of self-attention in multiple devices, the order of which block is computed first does not change the output. So, they proposed a ring style communication. Each device will take care of self-attention output of a subset of queries, and it only receives a subset of keys and values from the next device at each time, which reduces the communication cost. Also, by overlapping the computation of queries’ attention to the current subset of keys and values and the communication of the next subset of keys and values, this host communication does not introduce additional latency overhead.", "questions": "1. I am wondering if the same idea could be adopted to distribute linear projections in Attn and FFN computations as well", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea is simple and effective. \n\n2. The proposed distributed self-attention computation allows bypassing the hardware limitation of single device. \n\n3. The ring style communication allows the communication requirement stays constant when scaling to more devices. \n\n4. The overlapping of computation and communication hides the communication overhead.", "weaknesses": "1. There is only one toy experiment comparing the performance of RingAttention with other efficient attention baselines. The proposed method enables full self-attention for longer sequences, so it would be interesting to see the performance potential of full self-attention on the real world datasets. \n\n2. It would be better to include the results on network bandwidth utilizations in Table 3.\n\n3. In causal attention (used in most LLMs), the compute cost for different queries are different. Queries at the begin of sequence require less compute. In the proposed method, since each device takes care of a subset of queries, the compute load on different devices are different. Will this method make some devices underutilized?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EPkKuj04LW", "reviewer": "Reviewer_Z7SR", "summary": "Using large context lengths poses memory challenges since it is needed to maintain the key value vectors of earlier tokens. This paper proposes to distribute these vectors across multiple nodes allowing linearly scaling the context length with number of nodes. To address the challenge that computing the attention for a token requires access to all previous key value vectors, the authors propose a scheme to move the vectors across the node. Namely, they suggest forming a ring of nodes and passing the vector to the next node in the ring continuously until all vectors pass through all nodes. Moreover, the authors point out that the communication can be overlapped with the computation, thus making the overhead negligible. The paper shows the method can be practically applied through various experiments, namely fine-tuning LLaMA with 500k token context length.", "questions": "(Please also see the weakness section.)\n\n1. How is ring attention applied while doing auto regressive decoding (in the stage of decoding tokens one by one)?\n\n2. How does the training time scale with number of nodes? For example, what is the rate of slowdown observed when training a model with 2k context length on two nodes in comparison with k context length on one node (expected 4x) and in comparison with using normal attention to train with 2k context length on a single node?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The method effectively allows scaling the context length with the number of nodes. The proposal to overlap communication and computation could allow hiding the communication cost. It can also be implemented in a fairly straightforward manner. The authors show Ring Attention's practicality through experiments. It is very nice to see experiments done both on GPUs and TPUs. The paper is written clearly and can be followed reasonably easily.", "weaknesses": "While the paper contains experiments showcasing that ring attention can be applied in practice, a comparison with other methods is missing. For example, if applicable, is it better to apply pipeline parallelism or ring attention? Other similar methods also include DeepSpeed Inference [1] which suggests to offload the kv cache to cpu and similarly overlap communication and computation and FlexGen [2]. Additionally, in light of these existing work (some of which are also not referenced), the novelty of the work might be limited.\n\nFurthermore, an ablation study is missing. For example, how important is it to overlap communication and computation?\n\nI would also like to point out that while the method allows scaling the context length by increasing the number of nodes, it does not address challenges of handling long contexts by Transformers (e.g. [3]). As such, I find comparison with Claude and GPT3.5 in Figure 3 a bit misleading. In particular, we do not expect this method to improve the performance of the model on small context lengths (though a small increase is observed possibly due to more fine-tuning). \n\n[1] Aminabadi, Reza Yazdani, Samyam Rajbhandari, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Olatunji Ruwase et al. \"DeepSpeed-inference: enabling efficient inference of transformer models at unprecedented scale.\" In SC22: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-15. IEEE, 2022.\n\n[2] Sheng, Ying, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Re, Ion Stoica, and Ce Zhang. \"FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU.\" (2023).\n\n[3] Liu, Nelson F., Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. \"Lost in the middle: How language models use long contexts.\" arXiv preprint arXiv:2307.03172 (2023).\n\nSome minor suggestions:\n* Table 1 is not table but a figure. \n* The last paragraph of section 2 is an almost replica of the introduction (this is not a problem on its own. I am merely pointing this out in case the authors would like to make better use of the space)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)