作为 PyTorch 2.0 发布的一部分，注意力机制的加速实现作为“Better Transformer”项目的一部分（在 PyTorch 中称为 Accelerated Transformers）已原生添加到 PyTorch 中，即 torch.nn.functional.scaled_dot_product_attention。此实现利用了 FlashAttention 和 Memory-efficient attention 的融合内核，并支持训练和推理。
我们还发布了一个笔记本，展示了此集成的示例 此处
在看到 扩散模型的推理速度提升 20-30% 后，我们继续通过 🤗 Optimum 库 实现了与 🤗 Transformers 模型的集成。与 之前针对编码器模型的集成 类似，此集成将 Transformers 的模块替换为使用 torch.nn.functional.scaled_dot_product_attention 的高效实现。用法如下
from optimum.bettertransformer import BetterTransformer
from transformers import AutoModelForCausalLM
with torch.device(“cuda”):
model = AutoModelForCausalLM.from_pretrained(“gpt2-large”, torch_dtype=torch.float16)
model = BetterTransformer.transform(model)
# do your inference or training here
# if training and want to save the model
model = BetterTransformer.reverse(model)
model.save_pretrained(“fine_tuned_model”)
model.push_to_hub(“fine_tuned_model”)
下面总结了我们关于 torch.nn.functional.scaled_dot_product_attention 的发现
- 它最有助于在给定硬件上拟合更大的模型、序列长度或批次大小进行训练。
- 训练期间 GPU 上的内存占用节省从 20% 到 110%+ 不等。
- 训练期间加速从 10% 到 70% 不等。
- 推理期间加速从 5% 到 20% 不等。
- 对于小型头部维度，
scaled_dot_product_attention的独立加速可高达 3 倍，内存节省可高达 40 倍（取决于序列长度）。
您可能会对内存节省和加速的范围之广感到惊讶。在这篇博文中，我们讨论了我们的基准测试、此功能的亮点以及未来 PyTorch 版本中即将进行的改进。
在下一个 Transformer 版本中，您只需安装适当的 Optimum 版本并运行
model = model.to_bettertransformer()
即可使用 BetterTransformer API 转换您的模型。您现在可以通过从源代码安装 Transformer 来试用此功能。
基准测试和与 🤗 Transformers 的用法
torch.nn.functional.scaled_dot_product_attention 可用于任何使用标准注意力的架构，并主要取代样板代码
# native scaled_dot_product_attention is equivalent to the following:
def eager_sdpa(query, key, value, attn_mask, dropout_p, is_causal, scale):
scale_factor = 1 / math.sqrt(Q.size(-1)) if scale is None else scale
attn_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0) if is_causal else attn_mask
attn_mask = attn_mask.masked_fill(not attn_mask, -float('inf')) if attn_mask.dtype==torch.bool else attn_mask
attn_weight = torch.softmax((Q @ K.transpose(-2, -1) * scale_factor) + attn_mask, dim=-1)
attn_weight = torch.dropout(attn_weight, dropout_p)
return attn_weight @ V
在 🤗 Optimum 与 Transformer 模型的集成中，目前支持以下架构：gpt2、gpt-neo、gpt-neox、gptj、t5、bart、codegen、pegasus、opt、LLaMA、blenderbot、m2m100。预计此列表将在不久的将来扩展！
为了验证原生缩放点积注意力的优势，我们进行了推理和训练基准测试，结果如下所示。
在单个 A10G GPU、AWS g5.4xlarge 实例上进行的推理基准测试
在单个 A10G GPU、AWS g5.4xlarge 实例上进行的训练基准测试
在单个 A100-SXM4-80GB、Nvidia DGX 上进行的训练基准测试
在此基准测试中，最有趣的发现是原生 SDPA 允许使用更长的序列长度和批次大小，而不会出现内存不足问题。此外，推理期间可实现高达 20% 的加速，训练期间甚至更大。
正如训练基准测试所示，较小的头部维度似乎带来了更高的加速和内存节省，我们将在下一节中讨论。
借助 🤗 Accelerate 库，通过将 device_map=”auto” 传递给 from_pretrained 方法，该实现还支持多 GPU 设置。以下是在两个 A100-SXM4-80GB 上进行训练的一些结果。
在两个 A100-SXM4-80GB、Nvidia DGX 上进行的训练基准测试，使用 🤗 Accelerate 库进行分布式训练
请注意，某些内核仅支持 sm_80 计算能力（即 A100 GPU 的计算能力），这限制了在各种硬件上的可用性，尤其是当头部维度不是 2 的幂时。例如，截至 PyTorch 2.0.0 训练期间，opt-2.7b (headim=80) 和 gpt-neox-20b (headdim=96) 无法调度到使用 Flash Attention 的内核，除非在 A100 GPU 上运行。未来可能会开发出更好的内核：https://github.com/pytorch/pytorch/issues/98140#issuecomment-1518101895
Flash Attention、内存高效注意力与数学差异
原生 scaled_dot_product_attention 依赖于三种可能的后端实现：Flash Attention、内存高效注意力，以及所谓的数学实现，它为所有 PyTorch 平台提供硬件无关的备用方案。
当给定问题大小存在融合内核时，将使用 Flash Attention 或内存高效注意力，从而有效降低内存占用，因为在内存高效注意力的情况下，GPU 全局内存上执行 O(N) 内存分配，而不是传统急切注意力实现的经典 O(N^2)。通过 Flash Attention，预计内存访问（读写）次数会减少，因此两者都能带来加速和内存节省。
“数学”实现只是一个使用 PyTorch C++ API 的实现。此实现中值得注意的是，查询和键张量会单独缩放以提高数值稳定性，因此会启动两个 aten::div 操作，而不是在不包含此数值稳定性优化的急切实现中可能只启动一个操作。
头部维度对加速和内存节省的影响
在对 torch.nn.functional.scaled_dot_product_attention 进行基准测试时，我们注意到随着头部维度的增加，加速/内存增益会降低。这对于某些架构来说是一个问题，例如 EleutherAI/gpt-neo-2.7B，它的头部维度相对较大，为 128，或者 EleutherAI/gpt-j-6B（以及派生模型，如 PygmalionAI/pygmalion-6b），它的头部维度为 256（实际上目前由于头部维度过大而无法调度到融合内核）。
这种趋势可以在下图中看到，其中 torch.nn.scaled_dot_production 与上述急切实现进行独立基准测试。此外，我们使用 torch.backends.cuda.sdp_kernel 上下文管理器来强制使用数学、Flash Attention 和内存高效注意力实现。
使用内存高效注意力 SDP 内核（仅向前），A100
使用数学（无 dropout），A100
使用 Flash Attention SDP 内核（无 dropout），A100
使用内存高效注意力 SDP 内核（无 dropout），A100
我们看到，对于相同的问题规模，无论是仅推理还是训练，加速都随着头部维度的增加而降低，例如，使用 Flash Attention 内核时，从 headdim=8 的 3.4 倍降至 headdim=128 的 1.01 倍。
随着头部维度的增加，内存节省的减少是预期的。回想标准注意力计算
由于中间计算，此标准分步计算中的全局内存占用为 2 * N * N + N * d。内存高效注意力建议迭代更新 softmax 归一化常数并将其计算移到最后，从而仅进行恒定输出内存分配 N * d。
因此，内存节省比率为 2 * N / d + 1，它随着头部维度的增加而降低。
在 Flash Attention 中，权衡在于头部维度 d 和 GPU 流式多处理器共享内存大小 M 之间，总内存访问次数为 O(N² * d²/M)。因此，内存访问与头部维度呈二次方关系，与标准注意力呈线性关系相反。原因是，在 Flash Attention 中，对于较大的头部维度 d，键和值 K、V 需要分成更多块以适应共享内存，反过来，每个块需要加载完整的查询 Q 和输出 O。
因此，Flash Attention 的最高加速发生在比率 d² / M 足够小的状态下。
PyTorch 2.0.0 的当前限制
缺少比例参数
截至 PyTorch 2.0.0，torch.nn.functional.scaled_dot_product_attention 没有比例参数，并使用隐藏大小的默认平方根 sqrt(d_k)。
然而，一些架构（如 OPT 或 T5）在注意力中不使用缩放，这在 PyTorch 2.0.0 中强制其在调用 scaled_dot_product_attention 之前进行人工重新缩放。这引入了不必要的开销，因为除了注意力中不必要的除法之外，还需要额外的乘法。
此问题的修复已合并到 PyTorch 存储库中。
Flash Attention / 内存高效注意力与自定义掩码的支持
截至 PyTorch 2.0.0，当传递自定义注意力掩码时，无法使用 Flash Attention 和内存高效注意力。在这种情况下，scaled_dot_product_attention 会自动调度到 C++ 实现。
然而，正如我们所见，一些架构需要自定义注意力掩码，例如使用位置偏差的 T5。此外，在批次大小大于 1 且某些输入可能被填充的情况下，还需要传递自定义注意力掩码。对于后一种情况，替代方法是使用 SDPA 支持的 NestedTensor。
因此，对自定义掩码的有限支持限制了 SDPA 在这些特定情况下的优势，尽管我们希望 未来能获得更广泛的支持。
请注意，xformers（PyTorch 的 SDPA 部分受其启发）目前支持任意注意力掩码：https://github.com/facebookresearch/xformers/blob/658ebab39545f180a6075385b3897921623d6c3b/xformers/ops/fmha/cutlass.py#L147-L156。HazyResearch 的 Flash Attention 实现也支持等效的填充实现，因为使用了累积序列长度数组以及打包的查询/键/值——本质上类似于 NestedTensor。
总结
使用 torch.nn.functional.scaled_dot_product_attention 是一种免费的优化，它使您的代码更具可读性，占用更少的内存，并且在大多数常见情况下速度更快。
尽管 PyTorch 2.0.0 中的实现仍有一些小限制，但推理和训练在大多数情况下已经从 SDPA 中大量受益。我们鼓励您使用此原生实现来训练或部署您的 PyTorch 模型，并将其作为 🤗 Transformers 模型的一行转换！
未来，我们希望调整 API，以便用户也能在基于编码器的模型中使用 SDPA。
我们感谢 Benjamin Lefaudeux、Daniel Haziza 和 Francisco Massa 在头部维度影响方面的建议，以及 Michael Gschwind、Christian Puhrsch 和 Driss Guessous 对这篇博文的反馈！
基准测试复现
本文中提出的基准测试是使用 torch==2.0.0、transformers==4.27.4、accelerate==0.18.0 和 optimum==1.8.0 完成的。