作者: William Brandon, Aniruddha Nrusimha, Kevin Qian, Zachary Ankner, Tian Jin, Zhiye Song, Jonathan Ragan-Kelley

A1 主要贡献

核心问题与研究目标：
为了满足Transformer模型【22, Attention is all you need, 2017, NIPS】对日益增长的序列长度的需求，Liu等人近期提出了Ring Attention【13, Ring attention with blockwise transformers for near-infinite context, 2023】，这是一种精确的注意力算法，通过将自注意力计算分布到多个设备上，克服了单设备内存瓶颈。本文专注于因果Transformer模型这一重要特例，研究了Ring Attention的性能特点，并发现因果注意力计算的三角结构导致了关键的工作负载不平衡问题。在因果自注意力中，每个查询只与序列中较早出现的键进行交互，这原则上能将总操作数减少约一半。然而，现有的Ring Attention算法无法有效利用这一结构来提升吞吐量。

关键发现：Ring Attention的负载不平衡：
本文观察到，Ring Attention在处理因果注意力时存在一个限制：除了算法的第一次迭代外，在后续所有迭代中，某些设备的工作负载是完全必要的（未被掩码），而其他设备的工作负载则是完全不必要的（被掩码）。由于Ring Attention的延迟由每轮迭代中所有参与设备的最大延迟决定，因此无论单个设备如何优化，每轮迭代的延迟都等同于计算一个完全未被掩码的工作负载所需的时间。这导致Ring Attention的运行速度与没有注意力掩码的工作负载相当，尽管原则上它只需要计算一半的操作。

解决方案：Striped Attention：
为了解决这一问题，本文提出了Ring Attention的一个变体，称为Striped Attention。该方法通过对输入序列进行置换，几乎完全消除了Ring Attention中存在的负载不平衡。具体而言，Striped Attention对序列进行置换，使得每个设备总是操作一个非连续的、在原始序列中近似均匀分布的词元子集。这种分区方式确保了在每个设备上的查询/键交互中，大约有一半会被因果掩码抑制。如下图1所示，Striped Attention将词元（token）交错分配给不同设备，而Ring Attention则是分配连续的词元块。

图 1. Ring Attention和Striped Attention的Q、K、V序列初始分区成块的方式。由于K和V序列在Ring Attention算法中一起传输，它们被描绘为单个序列。注意，对于Ring Attention和Striped Attention，输入序列中的词元在运行模型的第一层之前就被分配到设备上，并在整个前向和后向传播过程中保持相同的布局。因此，在使用Ring Attention和Striped Attention时，Q、K、V在每个注意力层的开始时会自动以期望的布局进行分区，无需额外的每层通信来准备它们。

核心贡献：
本文的贡献总结如下：
1. 识别了工作负载不平衡问题：指出了Ring Attention在利用因果注意力结构以节省计算工作方面存在的负载不平衡。
2. 提出了Striped Attention：提出对Ring Attention进行修改，通过解决这种负载不平衡来提高因果序列建模任务的吞吐量。
3. 实验验证：通过实验测量了使用Striped Attention进行因果Transformer训练的吞吐量，发现在长序列训练中，该实现相比原始Ring Attention算法可实现高达1.65倍的端到端加速。

实验结果概述：
在使用JAX【1, JAX: composable transformations of Python+NumPy programs, 2018】实现的Striped Attention进行实验后，发现在8个A100 80GB GPU的服务器上训练十亿级参数的因果语言模型（序列长度达数十万）时，端到端速度提升高达1.45倍。在TPU上，速度提升更为显著，在16个TPUv4芯片上处理超过五十万词元的问题时，实现了1.65倍的加速。本文开源了实验代码，希望Striped Attention能成为一项基础性技术，推动研究人员在超长上下文因果Transformer模型的新兴领域中探索新应用。

A3 背景知识

2.1 因果自注意力

定义：本文的核心目标是为Transformer架构中的核心操作——因果自注意力，定义一个高效的分布式实现。因果自注意力接收输入矩阵 $Q, K, V \in \mathbb{R}^{n_{seq} \times d_{head}}$，其中 $n_{seq}$ 是输入序列长度，$d_{head}$ 是一个超参数。忽略无关的缩放因子，其计算公式为：

其中Softmax是按行计算的，而 $C$ 是一个因果掩码矩阵，其对角线以上的所有元素均为 $-\infty$，其他位置为零：

在 $QK^\top$ 矩阵的某个条目上加上 $-\infty$ 的效果是使其在Softmax计算中不被考虑，并确保该条目的Softmax输出为零。因此，因果掩码可以被解释为确保在序列中给定位置的注意力输出行仅受K和V矩阵中相同或更早位置的行的影响。

优化机会：通过分析因果自注意力的结构，我们发现在优化实现中有两个节省工作的机会：
1. 我们可以避免计算 $QK^\top$ 乘积中大约一半的元素：即对角线以上的元素，这些元素都将被 $C$ 设置为 $-\infty$。
2. 我们可以避免将 $V$ 乘以 $\text{Softmax}(QK^\top + C)$ 中大约一半的条目：即对角线以上的条目，这些条目保证为零。
这些优化是众所周知的，并被用于像FlashAttention【5, Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022】、【4, Flashattention-2: Faster attention with better parallelism and work partitioning, 2023】这样的融合注意力实现中，它们避免在全局内存中物化完整的 $\text{Softmax}(QK^\top + C)$ 矩阵。在理想条件下，这些优化可以将在因果注意力计算中所需的浮点运算次数（FLOPs）减少高达2倍；实际上，硬件限制和固定开销会降低这个因子。

2.2 Ring Attention

分布式策略：最近提出的Ring Attention算法【13, Ring attention with blockwise transformers for near-infinite context, 2023】提供了一种在多个以环形通信拓扑连接的加速器（如GPU或TPU）上分布注意力计算的策略。给定N个可用设备，Ring Attention假设矩阵 $Q, K, V$ 最初沿着序列维度被分片，因此它们可以分解为大小均匀的行块：

这样，设备0初始持有块 $Q_0, K_0, V_0$，设备1初始持有 $Q_1, K_1, V_1$，以此类推。图1的左侧展示了这种分区策略。

算法流程：Ring Attention通过执行N轮计算和通信来计算注意力输出，它将Q块固定在各自的设备上，同时将K和V块以循环方式从一个邻居传递到下一个。在每次迭代中，每个设备计算其拥有的Q块与当前持有的K和V块之间的交互。每个设备拥有输出矩阵的一个固定的行块，对应于其Q块相同的序列位置。设备使用由Rabe和Staats【15, Self-attention does not need o(n2) memory, 2021】引入的惰性softmax策略在迭代中累积部分输出。

算法伪代码与掩码：为便于说明，算法1中提供了描述Ring Attention算法高级结构的伪代码。该算法中特别值得注意的是函数 GetMask(j, k)，它返回控制查询块 $Q_j$ 和键/值块 $K_k, V_k$ 之间交互的注意力掩码瓦片（tile）。当使用Ring Attention计算因果自注意力时，GetMask(j, k) 的输出如下所示：

如果我们将原始的 $n_{seq} \times n_{seq}$ 因果掩码 $C$ 解释为一个 $N \times N$ 的块矩阵，每个块的形状为 $\frac{n_{seq}}{N} \times \frac{n_{seq}}{N}$，那么 GetMask(j, k) 的逻辑等同于选择 $C$ 中位置 $(j, k)$ 的条目。

负载不平衡图示：我们在图2中说明了Ring Attention算法在因果自注意力情况下的行为。我们可以看到，除了第一次迭代外，在所有迭代中，工作负载不平衡阻止了我们有效地利用因果掩码的结构来减少运行时间；尽管某些设备的交互完全被掩码掉，但我们无法通过让这些设备跳过计算来节省任何时间，因为其他设备的交互是完全未被掩码的。

图 2. Ring Attention应用于一个小的因果自注意力问题（nseq = 16），分布在N = 4个设备上。Q块保持静止，而K和V块在每一轮中以循环方式从一个邻居传递到下一个。每轮每个设备下方显示的正方形瓦片表示该设备在该轮计算的查询/键交互的因果掩码；被-∞值掩码掉的元素以黑色表示。我们可以看到，除了第一轮外，工作负载不平衡阻止了我们有效利用因果掩码的结构来减少运行时间。

2.2.1 Tiling

瓦片化计算：为了有效地避免不必要的计算，块被进一步划分为瓦片（tile）。对于每个瓦片，我们会检查它是否完全被掩码掉。如果是，我们完全跳过计算该瓦片。

瓦片化示例：例如，如果我们有一个512×512的瓦片大小，和一个包含九个瓦片（即1536×1536）的块大小，那么在一个带有因果掩码的块中，三个瓦片将是完全未掩码的，三个瓦片是部分掩码的，还有三个瓦片是完全被掩码的。我们会避免计算最后三个瓦片，从而减少33%的计算量。请注意，这少于50%的FLOPs减少量。但是，如果我们要计算的块是完全未掩码的，我们就需要计算所有九个瓦片。

A2 方法细节

Striped Attention的核心思想：Striped Attention通过一种新颖的序列分区方式，减少了Ring Attention在因果Transformer中的工作负载不平衡。它不是将词元（token）划分为连续的块，而是根据它们对设备数量N取模的余数，将它们划分为均匀间隔的条带集合。例如，给定设备数量为4，设备0将拥有词元0, 4, 8, ...，设备1将拥有词元1, 5, 9, ...，依此类推。图1对此进行了说明。

实现方式：在实践中，我们可以通过在模型的第一个嵌入层之前对输入词元进行置换，然后像在Ring Attention中那样将置换后的序列划分为连续的块来实现这种分区方案。对于使用位置嵌入方案（如RoPE【19, Roformer: Enhanced transformer with rotary position embedding, 2021】）的模型，位置ID也必须进行置换；在训练设置中，我们还必须置换用于计算损失的目标词元ID序列。

算法流程：分区后，我们的算法几乎与Ring Attention完全相同（见图3）。我们将设备概念化为一个从0到N索引的环。在每个注意力层的开始，每个设备i初始持有块 $Q_i, K_i, V_i$，这些块继承了应用于原始输入序列的条带化置换。在Striped Attention算法的每一轮中，每个设备i计算其分配的 $Q_i$ 块与前一轮收到的K, V块之间的因果注意力交互。同时，它将其当前的K, V块发送到环中的下一个设备，以隐藏内存延迟。

关键区别：因果性的参照：至关重要的是，Striped Attention中的注意力计算是相对于词元在原始输入序列中的顺序是因果的，而不是它们在置换后序列中的顺序。这对 GetMask(j, k) 函数（见算法1）返回的掩码结构有重要影响：与Ring Attention相反，Striped Attention确保每个设备的因果掩码在每次迭代中都是上三角的（见图4）。下面的算法描述了将Ringed Attention适配为Striped Attention所需的对 GetMask(j, k) 函数的修改。

procedure GETMASKSTRIPEDATTENTION(j, k)
    nseq = nseq
    Initialize MASK ∈ R^(N × N)
    ∀x, y MASK[x, y] = 0
    if j < k then
        ∀x, y | y ≤ x MASK[x, y] = −∞
    else if j ≥ k then
        ∀x, y | y < x MASK[x, y] = −∞
    end if
    return MASK
end procedure

图 3. Striped Attention应用于与上述相同的因果自注意力问题。与Ring Attention不同，每个块包含分布在原始输入序列各个部分的词元。与图2一样，每个设备下方的方形矩阵描绘了每个设备在每一轮中遇到的因果掩码。我们注意到，因果掩码为每个设备在每次迭代中提供了大致相等的可跳过工作部分，从而解决了Ring Attention中的工作负载不平衡问题。

工作负载分析：在Striped Attention中，设备i在第r轮的工作负载（其中$r \in [1, N-1]$）大致是

注意力交互，其中c是每个设备的块大小。相比之下，Ring Attention的计算调度在每次迭代中至少有一个设备必须处理大小为 $c^2$ 的工作负载，这成为每次迭代延迟的限制因素。

负载均衡的精确性：如上式所示，不同设备的工作负载并非完全相同；它们可能在是否包含对角线上有所不同。然而，随着块大小的增长，对角线占据的注意力矩阵比例线性缩小，工作负载不平衡占总运行时间的比例变得可以忽略不计。在块大小和设备数量趋于无穷大的极限情况下，使用Striped Attention相对于Ring Attention的最大理论加速接近2倍。

(b) Striped Attention中的工作负载分布
图 4. Ring Attention和Striped Attention在第0轮和第2轮的工作负载分布。方形矩阵代表所有可能的查询/键成对交互的集合；行索引对应查询，列索引对应键。对角线以上的所有单元格都被因果掩码屏蔽，可以跳过。颜色表示哪些设备负责计算的哪些部分。在第2轮，我们可以看到Ring Attention中的一些设备负责的工作负载完全被屏蔽，而Striped Attention则在所有设备间保持了平衡的工作负载。

Ring Attention算法伪代码：
算法1 由Liu等人【13, Ring attention with blockwise transformers for near-infinite context, 2023】引入的Ring Attention算法。
累加器 $Out_0, \dots, Out_{N-1}$ 存储输出的未归一化部分和，以及运行中的softmax统计信息，“normalize”步骤根据这些统计信息进行归一化，如Rabe和Staats【15, Self-attention does not need o(n2) memory, 2021】所述。“AccumulateAttentionFragment”函数封装了计算输入的片段上的部分注意力结果、累加到输出并更新运行中的softmax统计信息所需的逻辑。关键是，它会跳过每个被掩码的瓦片的计算。这里列出的算法是Ring Attention和Striped Attention的骨干，它们仅在实现GetMask(j, k)函数的方式上有所不同。

A4 实验环境与结果

实验环境

实现细节：
本文在JAX中实现了Striped Attention，作为Liu等人Ring Attention代码库的扩展，并在一些十亿参数规模的因果Transformer训练基准上将其吞吐量与Ring Attention进行了比较。

硬件配置：
实验在以下几种配置上运行：
1. 一台配备8个通过NVLink连接的NVIDIA A100 80GB GPU的服务器。
2. 一个拥有4个TPU v3芯片（即8个Tensor核心）的TPU Pod Slice。
3. 一个拥有16个TPU v4芯片（即32个Tensor核心）的TPU Pod Slice。

软件与算法配置：
- 掩码跳过策略：Ring Attention基线和Striped Attention实现都使用粗粒度的瓦片（tiling）策略来跳过被掩码的注意力计算部分。在每台设备的每一轮中，实现将查询/键交互空间分解为瓦片。
- 瓦片尺寸：GPU和TPU实验的瓦片尺寸不同。通过对Ring Attention进行瓦片尺寸扫描，选择了性能最佳的值，这些值对Striped Attention同样适用。具体来说：
- A100 80GB GPU：瓦片尺寸为2048个查询 × 4096个键。
- v3和v4 TPU：瓦片尺寸为2048个查询 × 2048个键。
- 计算精度：所有计算均以bfloat16格式进行，除了核心注意力计算 $\text{Softmax}(QK^\top + C)V$；在此核心计算中，Q、K、V被提升为float32缓冲区。在NVIDIA A100 GPU上，注意力的矩阵乘法以NVIDIA TensorFloat-32精度执行。在TPU上，这些操作以bfloat16精度执行。

实验变量：
实验中探究了三个方面对Striped Attention相对于Ring Attention吞吐量的影响：
1. 模型规模：研究了原始Ring Attention实现中使用的1B、3B和7B模型配置。这些模型的超参数见表1。
2. 序列长度：研究了总跨设备序列长度从8192（8k）到262144（256k）的训练情况。
3. 分布式网格维度：研究了序列并行度从N=2到N=8个设备，以及模型并行度【18, Megatron-lm: Training multibillion parameter language models using model parallelism, 2019】在1到4之间变化的情况。

表 1. 从原始Ring Attention代码库中获取的1B、3B和7B模型配置的架构超参数。

实验结果

性能总结：
表2总结了在序列长度为256k（262144）时，Striped Attention相对于Ring Attention在每个模型上实现的训练加速。图5绘制了在不同序列长度和网格配置下，在A100 GPU和TPU上实现的加速比。图6绘制了在256k序列长度下，每个模型使用两种方法所达到的绝对训练吞吐量（词元/秒）。

表 2. 在256k序列长度下，使用Striped Attention相对于Ring Attention的端到端训练加速。“Mesh”元组的两个元素分别指模型并行度和序列并行度的程度。由于我们的TPUv4测试平台设备数量是A100测试平台的两倍，我们相应地将序列并行度的程度加倍。

关键趋势分析：
- 小块尺寸下的性能：在所有配置中，第一个数据点代表每个设备块大小为4096×4096。由于GPU实验的瓦片大小为2048×4096，Striped Attention在最小块尺寸下无法节省计算。然而，TPU的瓦片大小为2048×2048，因此即使在最小块尺寸下，TPU也能节省25%的注意力FLOPs。
- 与序列长度的扩展性：最明显的趋势是，我们的方法的优势随序列长度的增加而扩展：
1. 块大小效应：在给定的网格配置下，加速比随着块大小的增加而增加。随着每个块中瓦片数量的增多，被跳过的计算百分比也随之增加，从而带来实际的加速。
2. 序列并行度效应：在给定的块大小下，只要块在查询和键维度上至少有两个瓦片，增加序列并行度就会带来更大的加速。这在A100的图表中最为明显。在给定的块大小下，所有具有相似序列并行度的模型都具有相似的加速比，而与模型并行度的程度无关。然而，具有高度序列并行度的模型加速更大。

图 5. Striped Attention相对于Ring Attention在不同序列长度和网格配置下的端到端训练加速。“网格配置”元组的两个元素分别指模型并行度和序列并行度的程度。在给定的网格配置下，增加序列长度也会增加每个块的瓦片数量。我们通常将实验扩展到系统在测试期间耗尽内存为止。

图 6. 在256k序列长度下，使用Ring Attention和Striped Attention在我们三种因果语言模型配置上的端到端训练吞吐量。所有A100实验均使用(2,4)网格配置，而所有TPUv4实验均使用(2,8)网格配置。

A7 补充细节

理论与实际加速比分析：原则上，Striped Attention可以将Ring Attention中核心注意力计算所需的运行时间减少最多2倍。考虑到除了注意力计算外还有其他成本会影响运行时间，并且Striped Attention只加速了Ring Attention的最后N-1次迭代，因此我们可以为实验中的每种配置量化一个通过用Striped Attention替换Ring Attention可实现的有效加速上限。我们称之为“理论最大加速比”（Theoretical Maximum Speedup, TMS）。

TMS计算方法：为了计算TMS，我们假设所有通信都与计算完美重叠，并且所有非矩阵乘法FLOPs所花费的时间可以忽略不计。对于我们的A100 GPU实验，我们假设注意力FLOPs（以TensorFloat-32执行）的成本是所有其他FLOPs（以bfloat16执行）的两倍。对于我们的TPU实验（其中我们所有的矩阵乘法都使用bfloat16精度），我们假设所有FLOPs的成本相等。

差距分析：我们在表3、表5和表4中比较了我们的理论最大加速比与实验中实现的加速比。我们可以看到，我们实现的Striped Attention所达到的加速比与我们TMS计算设定的有效上限之间存在差距。我们还看到，随着序列长度的增加，这个差距会缩小。我们认为这个差距主要由两个因素造成。

表 3. 在8个A100 GPU上，不同模型尺寸下Striped Attention相对于Ring Attention的理论与实测加速比。TMS代表理论最大加速比。

表 5. 在8个A100 GPU上，不同模型尺寸下Striped Attention相对于Ring Attention的理论与实测加速比。TMS代表理论最大加速比。

表 4. 在16个TPUv4上，扩大模型尺寸对理论与实际加速比的影响。TMS代表理论最大加速比。

差距原因：
1. 瓦片粒度的工作跳过：我们的实现是在瓦片（tile）的粒度上跳过工作的，这意味着在每个设备块大小较小时，我们无法充分利用因果掩码来跳过一半的工作。例如，如3.2节所述，当瓦片大小为2048个查询×2048个键，而每个设备的块大小为4096时，我们的实现只跳过了25%的注意力矩阵计算，尽管由于因果掩码，原则上约有50%的工作是可跳过的。随着序列长度增加，每个设备的块大小相对于瓦片大小变得更大，基于瓦片的工作跳过相对于理想工作跳过的开销就变小了。我们期待未来Striped Attention的实现能采用像FlashAttention【4, Flashattention-2: Faster attention with better parallelism and work partitioning, 2023】这样的融合内核，我们预计这将通过在更细的粒度上跳过工作，从而在更小的块大小下实现更大的加速。
2. 非注意力计算的开销：此外，我们怀疑在我们实验中使用的Ring Attention和Striped Attention实现中，非注意力计算的开销降低了Striped Attention可实现的加速。特别是，我们怀疑我们的Ring Attention和Striped Attention实现可能没有最优地重叠计算和通信，改进重叠可能会增加观察到的加速。这可能有助于解释为什么我们观察到的加速即使在长序列长度下似乎也没有收敛到我们计算的TMS值，因为计算成本和通信开销都随序列长度而扩展。我们将开发更精细优化的Ring Attention和Striped Attention实现留给未来的工作。

A5 结论

本文识别并解决了一个在最近提出的用于分布式长序列注意力的Ring Attention算法中存在的工作负载不平衡问题。我们的实验发现，这个我们称之为Striped Attention的解决方案，在训练因果Transformer处理极长序列时，能带来高达1.65倍的加速。此外，我们的方法作为Ring Attention的扩展易于实现，仅需要在前向传播开始时对输入序列进行一次性置换，并调整注意力掩码的结构。

从更广泛的意义上说，Striped Attention允许以计算高效的方式将大型语言模型扩展到以前未曾探索过的更长序列长度。由于我们的技术专注于因果注意力，它直接关系到生成式语言模型的训练和推理。因为Striped Attention是一种精确的注意力算法，应用程序可以在不牺牲任何准确性的情况下使用它。我们希望看到其他人基于Striped Attention，在新兴的超长上下文生成式Transformer模型领域开发新的应用。

STRIPED ATTENTION: FASTER RING ATTENTION FOR CAUSAL TRANSFORMERS

A1 主要贡献