Paper: RingAttention with Blockwise Transformers for Near-Infinite Context (ICLR-2024)

Reviewer comment (Experiments):
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 works (some of which are also not referenced), the novelty of the work might be limited.

Author response:
**#1: Relation to DeepSpeed-inference and FlexGen**

There are some key differences: 

1. CPU memory and the CPU-GPU interconnect are much slower than the GPU-GPU interconnect Ring Attention utilizes. Thus, using CPU offloading would not work for large context training and inference. 

2. Ring Attention allows training and inference with a context size that scales linearly with device count, while prior CPU offloading is only applicable to inference. Ring Attention does not add communication overheads while prior CPU offloading addes large communication costs. 

3. DeepSpeed-inference and FlexGen require storing the entire full sequence on each device, while Ring Attention distributes a long sequence to a ring of devices, and never needs to store full sequence on any device. This allows Ring Attention to process device count times longer sequences. 

4. Due to limitations in physical space and CPU memory, offloading to the CPU is insufficient for managing long sequence lengths. The small CPU-accelerator bandwidth necessitates more stringent requirements for hiding communication in CPU offloading, reducing its flexibility and applicability. In contrast, Ring Attention, which is widely applicable on GPU and TPU, requires only a fast interconnect between nearby devices, and keeps all data fast memory only while CPU offloading saves data to slower CPU memory. This minimal requirement enables Ring Attention to be scalable to any number of devices. Unlike CPU offloading that necessitates storing the full sequence on each device, Ring Attention distributes the sequence across multiple devices.


**#2: 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 parallelism methods, like data, pipeline, tensor parallelism, and FSDP, necessitate storing entire sequences on each device, rendering them incapable of processing large context sizes. Pipeline parallelism also leads to device idle time due to non-overlapped communication. Sequence parallelism distributes the sequence across devices but requires extensive communication that cannot be overlapped, making it infeasible for large context sizes.  

We have also included a comparison with Deep Speed Ulysses, which was published on Arxiv three days prior to the ICLR submission deadline. Deep Speed Ulysses combines sequence parallelism and tensor parallelism for their optimized NVSwitch all-to-all topology. It can reduce the communication cost of sequence parallelism. Importantly, it requires gathering and storing full sequence on each device, while Ring Attention distributes long sequences to a ring of devices and never needs to keep full sequence on any device. Ring Attention outperforms it substantially both in maximum context size and training efficiency. Please refer to our response to reviewer haGA for the full results.

Ring Attention is the first work that distributes transformers computation on long sequences across devices and enables fully overlapped communication with computation. Ring Attention is orthogonal and fully compatible with existing large-scale training parallelism methods. 

We demonstrate this by applying Ring Attention to FSDP end-to-end large-scale training of 7B-65B LLM on 32-1024 TPUv4, covering various compute budgets. Ring Attention facilitates the effortless training of very large context sizes without incurring overheads or making approximations. For example, on 1024 TPUv4, Ring Attention enables training with context sizes over millions in tokens, which is 512 times larger than prior state-of-the-art memory efficient / flash attention, without adding communication nor computation overheads.

Reviewer comment (Evaluation):
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, comparison with Claude and GPT‑3.5 in Figure 3 is 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).

Author response:
**#4: Long context impact on model performance**

The model trained with Ring Attention maintains competitive accuracy levels comparable to Claude and GPT4 and can process sequences 5-16 times longer. We include results of applying Ring Attention to extend the LLaMA context length to 512K tokens and evaluate its performance using LongChat tasks, as shown in Figure 3. Following the reviewer’s request, we also studied model performance on additional tasks, including training agentic transformers with more trajectories for improved in-context RL. The results in Table 5 demonstrate that a large context window significantly enhances the performance of learning decision-making policies from large-scale experiences. A large context allows agentic transformers to encode many more trajectories simultaneously, making hindsight relabeling more informative and enabling transformers to learn policies more effectively.

Regarding the “lost in the middle” observation, it is an interesting observation that some LLMs may not fully leverage the full context information. This observation is orthogonal to our paper, as our work focuses on making it possible to train long-context transformers without making approximations. Ring Attention could enable academic and industry labs to train long-context transformers to further investigate this phenomenon. Anthropic’s blog (https://www.anthropic.com/index/prompting-long-context#b3) suggests a different trend compared to the “lost in the middle” observation. This discrepancy might relate to task-specific characteristics, where additional texts tend to contain relevant information primarily at the beginning or end due to human biases, indicating a potential data issue. These are intriguing research areas orthogonal to Ring Attention, and Ring Attention can facilitate future research on the impact of long contexts on model performance.

Reviewer comment (Experiments):
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.

Author response:
Thank you for the suggestion. We have conducted additional experiments on more tasks. We applied Ring Attention to train agentic transformers on hundreds of trajectories to achieve better in-context RL. The results in Table 5 show that a large context window significantly improves performance of learning decision making policy from large scale experience. Large context allows agentic transformers to encode many more trajectories at the same time, so the hindsight relabeling is more informative and makes transformers learn a policy more effectively.

Reviewer comment (Evaluation):
It would be better to include the results on network bandwidth utilizations in Table 3.

Author response:
Thanks for suggesting to report network utilization rate in addition to maximum context sizes in end-to-end large-scale training, which can provide more information about the training. 
We have included the measurement of network utilization rates in Table 5. We measure the bandwidth utilization using \texttt{jax.profiler} and, based on the average measurements from 3 runs. The results show that Ring Attention has significantly higher bandwidth utilization rates. This is attributed to more communication from larger context sizes.

Reviewer comment (Experiments):
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?

Author response:
You are right that in causal attention mode, our method will make some devices underutilized. This is a great, insightful suggestion to Ring Attention. By skipping upper triangular blocks computation and balancing the computation load between devices when using causal attention, the compute cost can be reduced and potentially leading to a twofold increase in speed. We are definitely interested in this and planning to research this direction in future work.

Reviewer comment (Experiments):
I am wondering if the same idea could be adopted to distribute linear projections in Attn and FFN computations as well.

Author response:
We focus on self-attention in this work, since it contributes the largest activation size in today’s state-of-the-art AI. We believe the idea of Ring Attention can be applied to any all-to-all ops and even mixture of experts or split up a FFN computation or some aggregation.

Reviewer comment (Theory):
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?

Author response:
**#2: Saving compute in casual mode.**

This is a great, insightful suggestion for Ring Attention. By skipping the computation of upper triangular blocks and balancing the computation load between devices when using causal attention, the compute cost can be reduced, potentially leading to an increase in speed. We are definitely interested in this and plan to research this direction in future work.

Reviewer comment (Experiments):
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.

Author response:
**#1: Comparison with Deep Speed Ulysses.**

We appreciate the pointer to Deep Speed Ulysses which appeared on Arxiv 3 days before the ICLR deadline. 

The idea of Deep Speed Ulysses involves combining sequence parallelism and tensor parallelism to leverage their optimized all-to-all NVSwitch-based topology. First, it splits along the sequence dimension with sequence parallelism. Then, before computing attention, it aggregates query, key, and value using all-to-all to ensure that each host has the complete sequence; each host processes one attention head for tensor parallelism of attention scores. Finally, it uses all-to-all again to collect results along attention heads and reshards along the sequence dimension. This method helps in reducing the communication cost of sequence parallelism.

Importantly, Deep Speed Ulysses requires gathering and storing the entire sequence on each device, which becomes an issue for large context windows. In contrast, Ring Attention distributes sequences across devices and communicates key-value pairs in a ring that no device ever holds the entire sequence. Therefore, Ring Attention can scale context size linearly with device count, whereas Deep Speed Ulysses cannot.

Deep Speed Ulysses is great for leveraging all-to-all topology but all-to-all is difficult to scale to a large number of devices. For instance, for H100, the maximum number of GPUs supported with nvlink all-to-all interconnect is limited by network switch. In contrast, TPU torus design can maintain a fast interconnect arbitrarily, such as on thousands of TPUs with optical circuit switching. Since Ring Attention requires a minimal ring topology, it works great on both GPU and TPU, and scales particularly well to large TPU pods. For instance, Ring Attention allows a >10M context size for LLaMa LLM with TPU 1024, which is 512 times longer than the prior state-of-the-art in memory efficiency/flash attention.

We conducted a comparison in end-to-end training of LLaMa (7B, 13B, 34B) using 512x A100 80GB. We compared the maximum context size achieved by Deep Speed Ulysses and Ring Attention, and measured their model flops utilization (MFU). The table below shows the results. Ring Attention not only achieves 64 times longer context size thanks to distributing the sequence across devices but also attains much higher MFU due to overlapping communication in a ring of devices. It’s worth noting that self-attention has lower MFU than feedforward network, and larger context can lower MFU. Even so, Ring Attention outperforms the baseline, showing its effectiveness.


| | Context size (x1e3) | | MFU(%) | |

|---------|---------------------|-------|-------|----|

| | Ulysses | RingAttention | Ulysses | RingAttention |

| 7B | 256 | 16384 | 38 | 42 |

| 13B | 128 | 8192 | 40 | 44 |

| 34B | 64 | 4096 | 41 | 45 |