RingAttention with Blockwise Transformers for Near-Infinite Context

Dear Reviewer rUcK,

Thank you so much for the positive review and for highlighting that our work is simple and effective. We appreciate your feedback and suggestions, which we have incorporated into the revised version to improve it.

#1: 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.

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.

#2: It would be better to include the results on network bandwidth utilizations in Table 3.

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 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.

#3: Balancing compute load on different devices for causal attention (used in most LLMs)

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.

#4: Whether the same idea could be adopted to distribute other computations as well.

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.

Please let us know if our response resolves your concerns. We look forward to hearing from you. Thank you so much.

Dear reviewer Z7SR,

Thank you so much for your review and for highlighting the effectiveness of our method, the comprehensiveness of our experiments, and the clear, easily-followable writing of our paper. We appreciate your feedback and suggestions, which we have incorporated into the revised version to improve it.

#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.

Dear reviewer dMuy,

Thank you so much for the review. We believe that there are misunderstandings. We appreciate your feedback and suggestions, which we have incorporated into the revised version to improve it.

#1: 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.

We respectfully disagree with the reviewer’s criticism regarding novelty. To the best of our knowledge, Ring Attention is the first work that proposes distributing blockwise transformers in a ring topology and overlaps communication with computation. It is also the first to enable training transformers on extremely long sequences without incurring extra computation or communication costs. Ring Attention allows the context size to scale linearly with the device count, addressing the longstanding challenge of freeing transformers from memory and context limitations. For example, Ring Attention enables training LLaMa with a context size of more than millions of tokens on 1,024 TPUv4, which is 512 times larger in context size compared to the previous state-of-the-art blockwise transformers and memory-efficient/flash attention. In contrast, today’s state-of-the-art commercial models have much more limited context sizes, such as GPT-4-turbo has only 128K.

We believe that an effective, novel, and intuitive idea should be encouraged rather than discouraged.

We would like to point out that attention is not tiled matrix multiplication since it involves softmax, that is why the ML community has published memory efficient attention which computes attention cumulatively to allow long sequence and flash attention which optimizes the speed while allowing long sequence. However, these works require storing full sequences on each device, making it impossible for processing long sequences such as books, videos, and trajectories, among others.

Past published work in the ML community has suggested less scalable solutions, such as sequence parallelism, tensor parallelism, and fully sharded data parallelism. These methods are less scalable because they require storing full sequences on each device (as in tensor parallelism and fully sharded data parallelism) or incur large communication costs that cannot be overlapped (as in sequence parallelism). Ring Attention, however, is an innovative approach that distributes long sequences across a ring of devices and never needs to store the full sequence on any single device. It is the first work that enables linear scaling up of context size with the number of devices for extremely long sequences without incurring extra computation or communication costs.

A second important contribution of our paper is the first set of experiments with transformers trained to this massive context length, showing good performance on several ML tasks. We demonstrate that Ring Attention enables us to train Vicuna chatbot to 512K context size with just 32 GPUs, while prior state-of-the-art memory-efficient/flash attention-based methods can only handle 16K context size. This is a 32-fold increase in context window size compared to the previous long context Vicuna. The model's performance is competitive with GPT-4 and Claude-2 in long-context line retrieval tasks, and it can process sequences over five times longer than those models. We also applied Ring Attention to scale agentic transformers across hundreds of hindsight-relabeled trajectories. This approach led to substantially improved in-context RL performance, achieving state-of-the-art results in in-context RL tasks.

Apart from proposing the idea of Ring Attention, this work provides an open-source implementation, which is plug-and-play and easy to use.

In sum, Ring Attention is the first work proposing the distribution of blockwise transformers in a ring topology and overlaps communication with computation. This approach addresses the long-standing memory challenge of transformers without making approximations, allowing the context size to scale linearly with device count without additional computation or communication costs.

#2: 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).

Since we do not make approximations to attention, the compute cost is indeed still O(N^2). Our aim is to make it possible to fit long sequences into transformers, i.e., addressing the memory challenge. Combining Ring Attention with approaches that reduce compute costs, such as sparse attention, is a promising direction, but our work makes even precise attention efficiently scalable to sequence lengths in the millions.

Dear Reviewer haGA,

Thank you so much for the positive review and for highlighting that our paper is well-motivated, the proposed method is well-introduced, and for acknowledging our strong experimental results. We appreciate your feedback and suggestions, which we have incorporated into the revised version to improve it.

#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 |

#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.

Please let us know if our response resolves your concerns. We look forward to hearing from you. Thank you so much.