作者：Lucas Pasqualin, Driss Guessous, Christian Puhrsch, Bertrand Maher, Michael Gschwind

太长不看：我们展示了如何使用加速 PyTorch 2.0 Transformer 模型和新引入的 torch.compile() 方法，以 nanoGPT 为例加速大型语言模型，nanoGPT 是 Andrej Karpathy 对 GPT 模型的一种精简开源实现。通过使用加速 PT2 Transformer 模型中引入的新的 scaled dot product attention 算子，我们选择了 flash_attention 定制内核，实现了更快的每批次训练时间（使用 Nvidia A100 GPU 测量），从基准的约 143ms/批次提升到约 113 ms/批次。此外，使用 SDPA 算子的增强实现提供了更好的数值稳定性。最后，通过填充输入实现了进一步的优化，结合 flash attention 后可达到约 87ms/批次。

近年来，大型语言模型 (LLM) 和生成式 AI 在日常生活中得到了指数级增长的应用。与这些不断增长的模型紧密相关的是不断上升的训练成本——包括时间和硬件利用率。PyTorch 团队通过 加速 PyTorch 2 Transformer 模型（先前称为“Better Transformer”）以及 PyTorch 2.0 中的 JIT 编译功能，正面应对了这些挑战。

在这篇博文中，我们探讨了通过利用 SDPA（即 scaled dot product attention，Transformer 模型中的关键层）的定制内核实现所获得的训练优化。SDPA 的定制内核将多个离散的顺序操作替换为一个全局优化的内核，从而避免分配大量中间 CUDA 内存。这种方法带来了许多优势，包括但不限于：通过减少内存带宽瓶颈提高 SDPA 的计算性能，减小内存占用以支持更大的批次大小，以及通过预缩放输入张量增加数值稳定性。这些优化在 nanoGPT 上进行了演示，nanoGPT 是 Andrej Karpathy 对 GPT 模型的一种开源实现。

背景

Scaled dot product attention 是多头注意力机制的基本构建块，如 “Attention is All You Need” 论文中介绍的，并在 LLM 和生成式 AI 模型中有广泛应用。

The Transformer model architecture

图 1： 基于 “Attention is All You Need” 的 Transformer 模型架构。通过新的 PyTorch SDPA 算子，多头注意力机制可以通过用于输入投影的线性层、SDPA 算子和用于输出投影的线性层高效实现。

通过新的 scaled_dot_product_attention 算子，多头注意力机制仅需 3 个步骤即可实现：通过线性层进行输入投影、SDPA 和通过线性层进行输出投影。

# In Projection
# variable descriptions:
# q,k,v = Query, Key, Value tensors
# bsz = batch size
# num_heads = Numner of heads for Multihead Attention
# tgt_len = Target length
# src_len = Source Length
# head_dim: Head Dimension
    q, k, v = _in_projection(query, key, value, q_proj_weight, k_proj_weight, v_proj_weight, b_q, b_k, b_v)
    q = q.view(bsz, num_heads, tgt_len, head_dim)
    k = k.view(bsz, num_heads, src_len, head_dim)
    v = v.view(bsz, num_heads, src_len, head_dim)

    # Scaled Dot Product Attention
    attn_output = scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)

    # Out Projection
    attn_output = attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim)
    attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
    attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))

PyTorch 2. 支持针对特定用例优化并具有特定要求的多种不同内核。内核选择器会为特定的输入参数组合选择最佳内核。如果无法为特定的输入参数组合找到优化的“定制内核”，内核选择器将选择一个可以处理所有输入组合的通用内核。

虽然未来版本可能会扩展这组算子，但 PyTorch 2.0 推出了 SDPA 算子的 3 种实现

  1. 一种通用内核，在 sdpa_math() 函数中实现了 SDPA 的数学公式
  2. 一种优化内核，基于论文“Flash Attention”，支持在计算架构 SM80 (A100) 上使用 16 位浮点数据类型评估 SDPA。
  3. 一种优化内核，基于论文“Self-Attention Does Not Need O(n^2) Memory”并在 xFormer 中实现，它在更广泛的架构（SM40 及更高版本）上支持 32 位和 16 位浮点数据类型。这篇博文将此内核称为 mem_efficient 内核。

请注意，这两种优化内核（上面列出的第二种和第三种）都支持 key padding mask，并将支持的注意力 mask 限制为因果注意力 (causal attention)。加速 PyTorch 2.0 Transformer 模型目前仅在使用 is_causal 布尔值指定时支持因果 mask。指定 mask 时，将选择通用内核，因为分析提供的 mask 内容以确定它是否为因果 mask 的成本太高。关于每个内核的约束的更多解释可以在 加速 PT2 Transformer 博文 中找到。

在 nanoGPT 中启用加速 Transformer 模型

鉴于 SDPA 算子是 GPT 模型的关键组成部分，我们认为开源的 nanoGPT 模型是演示 PyTorch 2.0 加速 Transformer 模型易于实现性和其优势的绝佳选择。以下展示了在 nanoGPT 中启用加速 Transformer 模型的具体过程。

此过程主要围绕将现有的 SDPA 实现替换为从 functional.py 新增的 F.scaled_dot_product_attention 算子展开。此过程可以轻松应用于在许多其他 LLM 中启用此算子。或者，用户也可以选择直接调用 F.multi_head_attention_forward() 或在适用时使用 nn.MultiHeadAttention 模块。以下代码片段改编自 Karpathy 的 nanoGPT 仓库。

步骤 1：识别现有的 SDPA 实现

在 nanoGPT 中，SDPA 在模型的 CausalSelfAttention 类中实现。撰写本文时，原始实现已改编如下，以供参考。

The original implementation at time of writing

步骤 2：替换为 Torch 的 scaled_dot_product_attention

此时我们可以注意到以下几点

  • 第 36 - 42 行定义了我们正在替换的 SDPA 数学实现
  • 第 39 行应用的 mask 不再相关，因为我们使用的是 scaled_dot_product_attention 的 is_causal 标志。
  • 第 41 行使用的 dropout 层现在也不再需要了。

将 SDPA 实现替换为 torch 的 scaled_dot_product_attention 并移除现有的冗余代码后，得到以下实现。

Swapping out the SDPA implementation for torch’s scaled_dot_product_attention and removing the now redundant code yields the following implementation.

或者，也可以将原始 mask 传入 attn_mask 字段，但由于前面提到的内核约束，这将限制实现仅支持通用 sdpa_math 内核。

步骤 3（额外）：通过填充实现更快的矩阵乘法 (matmul)

除了 SDPA 带来的性能提升之外，我们的分析还带来了另一项不错的额外收获。用 Andrej 的话来说：“迄今为止，对 nanoGPT 最显著的优化（约 25% 的加速）就是简单地将词汇表大小 (vocab size) 从 50257 增加到 50304（最接近 64 的倍数）。”

Tweet by Andrej Karpathy

词汇表大小决定了 GPT 输出层中矩阵乘法 (matmuls) 的维度，这些维度非常大，以至于占据了整个训练循环的大部分时间！我们发现它们的性能远低于 A100 GPU 上可实现的峰值吞吐量，并从 NVIDIA 的矩阵乘法文档 中猜测，64 元素对齐会产生更好的结果。事实上，对这些矩阵乘法进行填充实现了近 3 倍的加速！根本原因在于非对齐内存访问显著降低了效率。更深入的分析可以在 此 Twitter 帖子 中找到。

通过这项优化，我们将训练时间从约 113 ms/批次（使用 flash attention）进一步缩短到约 87 ms/批次。

结果

下图展示了使用 PyTorch 定制内核获得的性能提升。具体数据如下

  • 基准（nanoGPT 实现）：约 143ms
  • sdpa_math（通用）：约 134ms（快 6.71%）
  • mem_efficient 内核：约 119ms（快 20.16%）
  • flash_attention 内核：约 113ms（快 26.54%）
  • flash_attention + 填充词汇表：约 87ms（快 64.37%）

所有代码均在配备 80 GB HBM [A100 SXM4 80GB] 的 8 个 NVIDIA Corporation A100 服务器上运行，并且出于本次实验的目的，dropout 设置为 0。

Using scaled dot product attention with custom kernels and torch.compile delivers significant speedups for training large language models

图 2： 使用 scaled dot product attention 配合定制内核和 torch.compile() 为训练大型语言模型提供了显著加速，例如此处展示的 nanoGPT。

增强数值模型稳定性

除了更快之外，PyTorch 的实现通过避免在许多执行场景中损失精度，提供了更高的数值稳定性。在此 这里 有一个很好的解释，但本质上 PyTorch 实现是在乘法之前对 Query 和 Key 矩阵进行缩放，这被认为更稳定并避免损失精度。由于 SDPA 融合的定制内核架构，这种缩放不会在注意力结果的计算中引入额外的开销。相比之下，由单独计算组件实现的版本需要额外的预缩放，从而产生额外成本。如需更多解释，请参见附录 A。

优化内存占用

使用 torch SDPA 内核的另一个巨大优势是内存占用减少，这使得可以使用更大的批次大小。下图比较了 flash attention 和基准因果注意力实现的训练一小时后的最佳验证损失。可以看出，使用基准因果注意力实现（在配备 80 GB HBM 的 8 个 NVIDIA Corporation A100 服务器上）实现的最大批次大小为 24，远小于使用 flash attention 实现的最大值 39。

Using Flash Attention enables the usage of larger batch sizes

图 3： 使用 Flash Attention 可以使用更大的批次大小，使用户在训练一小时后获得更低的验证损失（值越小越好）。

结论

加速 PyTorch 2 Transformer 模型旨在使最先进的 Transformer 模型的训练和生产部署更经济实惠，并与 PyTorch 2.0 模型 JIT 编译集成。新引入的 PyTorch SDPA 算子为训练 Transformer 模型提供了改进的性能，对于昂贵的大型语言模型训练尤其有价值。在这篇博文中，我们在示例性的 nanoGPT 模型上展示了许多优化，包括

  • 训练速度提升超过 26%，与使用恒定批次大小的基准相比
  • 通过填充词汇表实现的额外加速，使总优化幅度相对于基准达到约 64%
  • 额外的数值稳定性

附录 A：分析注意力数值稳定性

在本节中，我们更深入地解释了前面提到的通过预缩放 SDPA 输入向量获得的增强数值稳定性。以下是 nanoGPT 的 SDPA 数学实现的简化版本。这里需要注意的重要一点是，query 在未经缩放的情况下进行矩阵乘法。

# nanoGPT implementation of SDPA
# notice q (our query vector) is not scaled !
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)

# Dropout is set to 0, so we can safely ignore this line in the implementation# att = self.attn_dropout(att) 

y_nanogpt = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)

以下是 torch 的 scaled_dot_product_attention 中的等效数学实现。

# PyTorch implementation of SDPA
embed_size = q.size(-1)
scaling_factor = math.sqrt(math.sqrt(embed_size))
q = q / scaling_factor 	# notice q _is_ scaled here !

# same as above, but with scaling factor
att = q @ (k.transpose(-2, -1) / scaling_factor)
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
att = F.softmax(att0, dim=-1)

# Dropout is set to 0, so we can safely ignore this line in the implementation# att = self.attn_dropout(att) 

y_scale_before = att @ v

从数学上讲，这两种方法应该是等效的，但我们的实验表明，在实践中，我们从每种方法中得到了不同的结果。

使用上述方法，我们验证了 y_scale_before 与使用 scaled_dot_product_attention 方法得到的预期输出相匹配，而 y_nanogpt 不匹配。

使用 torch.allclose 方法来测试等效性。具体来说，我们展示了

y_sdpa = torch.nn.functional._scaled_dot_product_attention(
	q,
	k,
	v,
	attn_mask=self.bias[:,:,:T,:T] != 0,
	dropout_p=0.0,
	need_attn_weights=False,
	is_causal=False,
)

torch.allclose(y_sdpa, y_nanogpt) # False, indicating fp issues
torch.allclose(y_sdpa, y_scale_before) # True, as expected

附录 B：重现实验结果

希望重现这些结果的研究人员应从 Andrej 的 nanoGPT 仓库的以下 commit 开始 - b3c17c6c6a363357623f223aaa4a8b1e89d0a465。此 commit 被用作测量每批次速度提升时的基准。对于包含填充词汇表优化的结果（这些优化对批次速度提升最显著），请使用以下 commit - 77e7e04c2657846ddf30c1ca2dd9f7cbb93ddeab。从任一 checkout 版本，使用 torch.backends API 选择用于实验的内核都变得轻而易举。

可以通过上下文管理器选择所需的内核

with torch.backends.cuda.sdp_kernel (
    enable_math = False,
    enable_flash = False,
    enable_mem_efficient = True
):
    train(model)