PyTorch 2.0 版本包含 PyTorch Transformer API 的新高性能实现，旨在降低最先进 Transformer 模型的训练和部署成本。继“fastpath”推理执行（“Better Transformer”）成功发布后，此版本通过针对缩放点积注意力 (SPDA) 的自定义内核架构，引入了高性能的训练和推理支持。
您可以通过直接调用新的 SDPA 运算符（如 SDPA 教程中所述），或者通过透明地集成到现有 PyTorch Transformer API 中来利用新的融合 SDPA 内核。PyTorch Transformer API 的所有功能将继续兼容，许多功能映射到高性能 SDPA 内核，而其他功能则无法实现更高性能（例如，need_weights，如下所述），同时对其他功能的高性能支持可能仍在积极开发中。
与“fastpath”架构类似，自定义内核完全集成到 PyTorch Transformer API 中——因此，使用原生 Transformer 和 MultiHeadAttention API 将使用户能够透明地看到显著的速度提升。与“fastpath”架构不同，新引入的“自定义内核”支持更多用例，包括使用交叉注意力、Transformer 解码器和用于训练模型的模型，此外还支持现有 fastpath 推理，适用于固定和可变序列长度的 Transformer 编码器和自注意力用例。
为了充分利用不同的硬件模型和 Transformer 用例，支持多个 SDPA 自定义内核，并带有自定义内核选择逻辑，该逻辑将为给定模型和硬件类型选择最高性能的内核。特别是，PyTorch 2.0 版本中包含的第一个自定义内核是 Flash Attention 内核（sdpa_flash，用于 Nvidia GPU 上具有 SM80+ 架构级别的 16 位浮点训练和推理）和 xFormers 内存高效注意力内核（sdpa_mem_eff，用于各种 Nvidia GPU 上的 16 位和 32 位浮点训练和推理）。当自定义内核不适用时，通用内核 sdpa_math 提供了一种实现方式。
如前所述，自定义内核提供了更广泛的执行场景支持。为确保高效执行（例如，使用 GPU 张量核），模型配置需要满足少量要求。此要求列表将随时间推移而演变，有望放宽限制当前支持的自定义内核使用的约束，或者将来提供额外的内核。
有关自定义内核和调度约束的最新列表，您可以参考 sdp_utils.h。截至 PyTorch 2.0，现有融合 SDPA 内核具有以下约束：
- Flash Attention 仅支持 16 位浮点数据类型（float16 和 bfloat16）。
- 对于 16 位浮点数，头部维度必须是 8 的倍数；对于 32 位浮点数，头部维度必须是 4 的倍数。目前，Flash Attention 自定义内核支持的最大 head_dim 为 128。
- mem_efficient 内核的 CUDA 架构级别必须是 sm5x 或更高版本，Flash Attention 的 CUDA 架构级别必须是 sm80。
- Flash Attention 支持任意 dropout，在 PyTorch 2.0 中，mem_efficient 内核不支持 dropout（即，对于此内核，在 PyTorch 2.0 中必须将 dropout 设置为零）。
- 为了支持可变序列长度批次，所有 SDPA 内核都支持 Nested Tensor 输入，这些输入使用可变序列长度张量结合输入数据和填充信息进行前向传播。（您可以在 Nested Tensor 教程中找到有关 Nested Tensors 的更多信息。）
- 您可以通过在传递给 SDPA 运算符之前组合 `key_padding_mask` 和 `attn_mask` 来同时指定它们。特别是，您可以使用 nn.Transformer API 的每个批次元素的键填充掩码来实现批次中可变序列长度输入的训练。
- 目前，融合内核实现唯一支持的注意力掩码是训练中常用的因果掩码。要在自定义内核中指定因果掩码，必须使用 `is_causal` 布尔值进行指定，并且 `attn_mask` 必须为 None。
- 对嵌套张量的支持仍在开发中。具体来说，在 PyTorch 2.0 中，只有 sdpa_math 内核支持使用嵌套张量进行训练。此外，PyTorch 2.0 不支持将嵌套张量作为使用 torch.compile() 编译的代码的一部分。
- SDPA 运算符不支持返回平均注意力权重，因为计算它们会破坏使融合内核能够更有效执行的优化。`torch.nn.MultiheadAttention` 的 `forward` 函数的参数 `need_weights` 默认为 True。为了使用融合内核，`need_weights` 需要设置为 `need_weights=False`。
我们发现注意力掩码在实际应用中很少使用，除了训练期间的因果掩码。因此，我们通过内置使用因果掩码作为注意力掩码的选项来降低内核复杂度和计算成本，并通过与新的 SDPA 运算符一起引入的 `is_causal` 参数选择此新功能。
为常用因果掩码提供 `is_causal` 布尔标志也避免了昂贵且内存密集型的因果掩码分配，通过允许更多内存用于大批次大小来提高训练内存效率，并减少内存带宽和缓存争用——这两者在 GPU 加速器中都非常宝贵——因为不需要加载注意力掩码张量。
如果现有自定义内核的约束均未满足，则训练将退回到使用默认的 sdpa_math 内核，该内核使用一系列 PyTorch 运算符来实现 SDPA，从而实现缩放点积注意力的数学方程。这是最通用的“全能”回退内核，可确保所有模型都能成功训练。
除了现有的 Transformer API 之外，模型开发者还可以通过调用新的 `scaled_dot_product_attention()` 运算符直接使用缩放点积注意力内核。如 SDPA 教程中所述，该运算符可与输入投影和输出投影结合使用，以高效实现多头注意力。
除了添加自定义内核之外，加速 PyTorch 2 Transformers 还与 PyTorch 2.0 编译集成。要在受益于 PT2 编译的额外加速（用于推理或训练）的同时使用您的模型，请使用以下方法预处理模型：
model = torch.compile(model)
我们通过结合自定义内核和 `torch.compile()`，使用加速 PyTorch 2 Transformers 在训练 Transformer 模型，特别是大型语言模型方面，取得了显著的加速。
图：使用带自定义内核和 `torch.compile` 的缩放点积注意力，显著加速了大型语言模型的训练，例如此处所示的 nanoGPT。
最后，由于自定义内核的内存效率更高，请尝试增加训练批次的大小，以通过增加批次大小实现更快的训练。
除了自动内核选择之外，上下文管理器还允许开发人员覆盖内核选择算法——这对于日常操作不是必需的，但它使开发人员能够调试其代码，并使性能工程师能够覆盖内核选择。SDPA 教程提供了有关使用 SDPA 上下文管理器的更多信息。
除了作为 nn.Transformer API 的一部分提供外，加速 PyTorch 2 Transformer 自定义内核也随着 PyTorch 2.0 的发布，与 torchtext、torchvision 和 fairseq 领域库结合提供。