To help address the growing demand for ever-longer sequence lengths in
transformer models, Liu et al. recently proposed Ring Attention, an exact
attention algorithm capable of overcoming per-device memory bottle- necks by
distributing self-attention across multiple devices. In this paper, we study
the performance characteristics of Ring Attention in the important special case
of causal transformer models, and identify a key workload imbal- ance due to
triangular structure of causal attention computations. We propose a simple
extension to Ring Attention, which we call Striped Attention to fix this
imbalance. Instead of devices having contiguous subsequences, each device has a
subset of tokens distributed uniformly throughout the sequence, which we
demonstrate leads to more even workloads. In experiments running Striped
Attention on A100 GPUs and TPUv4s, we are able to achieve up to 1.45x
end-to-end throughput improvements over the original Ring Attention algorithm
on causal transformer training at a sequence length of 256k. Furthermore, on 16
TPUv4 chips, we were able to achieve 1.65x speedups at sequence lengths of
786k. We release the code for our experiments as open source

通过在转换模型中引入 Ring Attention 和 Striped Attention 算法，能够解决设备内存瓶颈和工作负载不均衡问题，实现了在 256k 和 786k 序列长度下分别提高 1.45 倍和 1.65 倍的端到端吞吐量。

条纹注意力：用于因果性变形器的更快环形注意力

Striped Attention: Faster Ring Attention for Causal Transformers

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 creating challenges for
tasks involving extended sequences or long-term dependencies. We present a
distinct approach, Ring Attention, which leverages blockwise computation of
self-attention to distribute long sequences across multiple devices while
concurrently overlapping the communication of key-value blocks with the
computation of blockwise attention. By processing longer input sequences while
maintaining memory efficiency, Ring Attention enables training and inference of
sequences that are device count times longer than those of prior
memory-efficient Transformers, effectively eliminating the memory constraints
imposed by individual devices. Extensive experiments on language modeling tasks
demonstrate the effectiveness of Ring Attention in allowing large sequence
input size and improving performance.

传统的 Transformer 架构由于内存需求限制了其处理长序列的能力，而我们提出的 Ring Attention 方法通过分块计算自注意力和重叠通信的方式，提高了内存利用效率，允许处理更长的输入序列，从而消除了设备内存的限制，同时在语言建模任务上取得了显著的性能提升。