闪光注意力：革命性地提高Transformer效率

</pre>
import torch

<p>def standard_attention(Q, K, V)：</p>

<p># Q, K, V形状：(batch_size, seq_len, d_model)</p>

<p> d_k = K.size(-1)</p>

<p> scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k))</p>

<p> attention_weights = torch.softmax(scores, dim=-1)</p>

<p> return torch.matmul(attention_weights, V)</p>

这种实现虽然直接，但遭受了上述低效率。得分张量的形状为（batch_size，seq_len，seq_len），对于长序列来说可能会变得非常大。

闪光注意力

闪光注意力，由Tri Dao和他的同事在2022年的论文中引入，是一种计算注意力的方法，它大大减少了内存使用和计算效率。闪光注意力的关键思想是：

1. 分块：将大型注意力矩阵分解为可以适合快速芯片SRAM的较小块。
2. 重新计算：不要存储整个注意力矩阵，而是在反向传播过程中需要时重新计算其部分。
3. IO感知实现：优化算法以最小化不同GPU内存层次之间的数据移动。

闪光注意力算法

在其核心，闪光注意力重新想象了我们如何计算注意力机制。它不再一次性计算整个注意力矩阵，而是以块为单位处理，利用现代GPU的内存层次结构。

以下是算法的高级概述：

1. 输入：矩阵Q、K、V在HBM（高带宽内存）和SRAM中，大小为M。
2. 根据可用的SRAM计算块大小。
3. 初始化输出矩阵O和辅助向量l和m。
4. 算法将输入矩阵分解为适合SRAM的块。
5. 两个嵌套循环处理这些块：
• 外循环加载K和V块。
• 内循环加载Q块并执行计算。
6. 芯片上的计算包括矩阵乘法、softmax和输出计算。
7. 结果在处理每个块后写回HBM。

这种块级计算使闪光注意力能够保持较小的内存占用，同时仍然计算出精确的注意力。

闪光注意力的数学原理

使闪光注意力工作的关键是一个数学技巧，它允许我们以块级方式计算softmax。该论文引入了两个关键公式：

1. softmax分解：
   softmax(x) = exp(x - m) / Σexp(x - m)

   其中m是x中的最大值。

2. softmax合并：
   softmax(x ∪ y) = softmax(softmax(x) * e^(m_x - m), softmax(y) * e^(m_y - m))

   其中m = max(m_x, m_y)。

这些公式使闪光注意力能够计算每个块的部分softmax结果，然后正确地合并它们以获得最终结果。

实现细节

让我们深入了解闪光注意力的简化实现，以说明其核心概念：

import torch

<p>def flash_attention(Q, K, V, block_size=256)：</p>

<p> batch_size, seq_len, d_model = Q.shape</p>

<p> # 初始化输出和运行统计</p>

<p> O = torch.zeros_like(Q)</p>

<p> L = torch.zeros((batch_size, seq_len, 1))</p>

<p> M = torch.full((batch_size, seq_len, 1), float('-inf'))</p>

<p> for i in range(0, seq_len, block_size)：</p>

<p> Q_block = Q[:, i:i+block_size, :]</p>

<p> for j in range(0, seq_len, block_size)：</p>

<p> K_block = K[:, j:j+block_size, :]</p>

<p> V_block = V[:, j:j+block_size, :]</p>

<p> # 计算此块的注意力得分</p>

<p> S_block = torch.matmul(Q_block, K_block.transpose(-2, -1)) / (d_model ** 0.5)</p>

<p> # 更新运行最大值</p>

<p> M_new = torch.maximum(M[:, i:i+block_size], S_block.max(dim=-1, keepdim=True).values)</p>

<p> # 计算指数</p>

<p> exp_S = torch.exp(S_block - M_new)</p>

<p> exp_M_diff = torch.exp(M[:, i:i+block_size] - M_new)</p>

<p> # 更新运行和</p>

<p> L_new = exp_M_diff * L[:, i:i+block_size] + exp_S.sum(dim=-1, keepdim=True)</p>

<p> # 计算此块的输出</p>

<p> O[:, i:i+block_size] = (exp_M_diff * O[:, i:i+block_size] + torch.matmul(exp_S, V_block)) / L_new</p>

<p> # 更新运行统计</p>

<p> L[:, i:i+block_size] = L_new</p>

<p> M[:, i:i+block_size] = M_new</p>

<p> return O</p>

这种实现虽然简化，但抓住了闪光注意力的本质。它以块为单位处理输入，保持运行统计（M和L）以正确地跨块计算softmax。

闪光注意力的影响

闪光注意力的引入对机器学习领域产生了深远的影响，特别是在大型语言模型和长上下文应用中。一些关键的好处包括：

1. 减少内存使用：闪光注意力将内存复杂度从O(N^2)减少到O(N)，其中N是序列长度。这使得可以使用相同的硬件处理更长的序列。
2. 提高速度：通过最小化数据移动和更好地利用GPU计算能力，闪光注意力实现了显著的加速。作者报告称，GPT-2的训练速度比标准实现快了3倍。
3. 精确计算：与其他一些注意力优化技术不同，闪光注意力计算精确的注意力，而不是近似值。
4. 可扩展性：减少的内存占用使得可以扩展到更长的序列，可能长达数百万个标记。

现实世界的影响

闪光注意力的影响超出了学术研究。它已经被许多流行的机器学习库和模型快速采用：

• Hugging Face Transformers：流行的Transformers库已经集成了闪光注意力，允许用户轻松利用其优点。
• GPT-4及以后：虽然没有得到确认，但人们猜测高级语言模型，如GPT-4，可能正在使用类似于闪光注意力的技术来处理长上下文。
• 长上下文模型：闪光注意力使得能够处理极长上下文的新一代模型成为可能，例如可以处理整个书籍或长视频的模型。

闪光注意力：最近的发展

标准注意力与闪光注意力

闪光注意力-2

在原始闪光注意力的成功基础上，同一团队在2023年推出了闪光注意力-2。此更新版本带来了几项改进：

1. 进一步优化：闪光注意力-2实现了更好的GPU利用率，在A100 GPU上达到70%的理论峰值FLOPS。
2. 改进的反向传播：反向传播被优化为几乎与前向传播一样快，从而导致训练速度大幅加快。
3. 支持不同注意力变体：闪光注意力-2扩展了对各种注意力变体的支持，包括分组查询注意力和多查询注意力。

闪光注意力-3

2024年发布的闪光注意力-3代表了这一研究线的最新进展。它引入了几种新技术以进一步提高性能：

1. 异步计算：利用新的GPU指令的异步性来重叠不同的计算。
2. FP8支持：利用低精度FP8计算以实现更快的处理。
3. 非相干处理：一种减少使用低精度格式时的量化误差的技术。

以下是闪光注意力-3如何利用异步计算的简化示例：

import torch
from torch.cuda.amp import autocast

<p>def flash_attention_3(Q, K, V, block_size=256)：</p>

<p> with autocast(dtype=torch.float8)： # 使用FP8进行计算</p>

<p> # ... （类似于之前的实现）</p>

<p> # 异步计算示例</p>

<p> with torch.cuda.stream(torch.cuda.Stream())：</p>

<p> #异步计算GEMM</p>

<p> S_block = torch.matmul(Q_block, K_block.transpose(-2, -1)) / (d_model ** 0.5)</p>

<p> # 同时，在默认流中：</p>

<p> # 为softmax计算做准备</p>

<p> # 同步流</p>

<p> torch.cuda.synchronize()</p>

<p> # 继续softmax和输出计算</p>

<p> # ...</p>

<p> return O</p>

此代码片段演示了闪光注意力-3如何利用异步计算和FP8精度。请注意，这是一个简化的示例，实际实现将更加复杂和特定于硬件。

在您的项目中实现闪光注意力

如果您对在自己的项目中利用闪光注意力感到兴奋，您有几个选项：

1. 使用现有库：像Hugging Face Transformers这样的许多流行库现在都包含闪光注意力的实现。只需更新到最新版本并启用相应的标志可能就足够了。
2. 自定义实现：对于更复杂的用例或自定义实现，您可能希望自己实现闪光注意力。xformers库提供了一个良好的参考实现。
3. 硬件特定优化：如果您使用特定的硬件（例如NVIDIA H100 GPU），您可能希望利用硬件特定的功能以获得最佳性能。

以下是如何使用Hugging Face Transformers库实现闪光注意力的示例：

from transformers import AutoModel, AutoConfig

<p># 启用闪光注意力</p>

<p>config = AutoConfig.from_pretrained("bert-base-uncased")</p>

<p>config.use_flash_attention = True</p>

<p># 加载带有闪光注意力的模型</p>

<p>model = AutoModel.from_pretrained("bert-base-uncased", config=config)</p>

<p># 像往常一样使用模型</p>

<p># ...</p>

挑战和未来方向

虽然闪光注意力在提高注意力机制的效率方面取得了显著进步，但仍然存在挑战和需要研究的领域：

1. 硬件特异性：当前的实现通常针对特定的GPU架构进行优化。跨不同硬件的优化仍然是一个挑战。
2. 与其他技术的集成：将闪光注意力与其他优化技术（如剪枝、量化和模型压缩）相结合是一个活跃的研究领域。
3. 扩展到其他领域：虽然闪光注意力在NLP中表现出色，但将其优点扩展到其他领域（如计算机视觉和多模态模型）仍是一个正在进行的工作。
4. 理论理解：对闪光注意力为什么如此有效的理论理解将有助于开发出更强大的优化技术。

结论

通过巧妙地利用GPU内存层次结构和数学技巧，闪光注意力在不牺牲准确性的情况下实现了显著的速度和内存使用改进。

正如我们在本文中所探讨的，闪光注意力的影响远远超出了简单的优化技术。它使得开发出更强大和高效的模型成为可能。