跳转到主要内容
博客

使用加速 Transformer 加速大型语言模型

TL;DR. 我们将展示如何使用加速的PyTorch 2.0 Transformer和新引入的torch.compile()方法来加速大型语言模型,以Andre Karpathy的GPT模型紧凑开源实现nanoGPT为例。通过使用加速PT2 Transformer引入的新缩放点积注意力运算符,我们选择了flash_attention自定义内核,并实现了更快的每批训练时间(使用Nvidia A100 GPU测量),从基线的约143毫秒/批提高到约113毫秒/批。此外,使用SDPA运算符增强的实现提供了更好的数值稳定性。最后,通过使用填充输入实现了进一步的优化,当与flash attention结合时,可达到约87毫秒/批。

近来,大型语言模型(LLMs)和生成式AI在日常生活中得到了指数级的普及。与这些不断增长的模型紧密相关的是不断增长的训练成本——无论是时间还是硬件利用率。PyTorch团队已经通过加速PyTorch 2 Transformer(以前称为“Better Transformer”)和PyTorch 2.0中的JIT编译正面应对了这些挑战。

在本篇博客文章中,我们探讨了通过利用SDPA(也称为缩放点积注意力)的自定义内核实现所获得的训练优化,SDPA是Transformer模型中的关键层。SDPA的自定义内核用一个全局优化的内核取代了几个离散的顺序操作,从而避免了分配大量的中间CUDA内存。这种方法具有多项优点,包括但不限于:通过减少内存带宽瓶颈实现SDPA的更高性能计算、减少内存占用以支持更大的批次大小,以及通过对输入张量进行预缩放来增加数值稳定性。这些优化在Andrej Karpathy的GPT开源实现nanoGPT上得到了验证。

背景

缩放点积注意力是多头注意力的基本组成部分,如“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内核。

请注意,这两种优化内核(上面列出的第二种和第三种)都支持键填充掩码,并将支持的注意力掩码限制为因果注意力。目前,加速的PyTorch 2.0 Transformer仅在通过is_causal布尔值指定时才支持因果掩码。当指定掩码时,将选择通用内核,因为分析提供的掩码内容以确定它是否是因果掩码的成本太高。有关每个内核的约束的更多解释可以在加速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行应用的掩码不再相关,因为我们使用的是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.

或者,原始掩码可以传递到attn_mask字段,但是由于前面提到的内核限制,这将使实现仅支持通用的sdpa_math内核。

步骤3(附加):使用填充加速矩阵乘法

除了SDPA带来的性能提升外,我们的分析还发现了一个不错的附加优势。用Andrej的话来说,“迄今为止,nanoGPT最显著的优化(约25%的加速)是简单地将词汇量从50257增加到50304(最近的64的倍数)。”

Tweet by Andrej Karpathy

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

通过这项优化,我们将每批次的训练时间从约113毫秒(使用闪存注意力)进一步缩短至约87毫秒。

结果

下图展示了使用Pytorch自定义内核所获得的性能。以下是具体数据:

  • 基线(nanoGPT实现):约143毫秒
  • sdpa_math(通用):约134毫秒(快6.71%)
  • mem_efficient内核:约119毫秒(快20.16%)
  • flash_attention内核:约113毫秒(快26.54%)
  • flash_attention + 填充词汇:约87毫秒(快64.37%)

所有代码均在配备8个NVIDIA Corporation A100服务器(80 GB HBM [A100 SXM4 80GB])上运行,本次实验中,dropout设置为0。

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

图2:使用带自定义内核和torch.compile的缩放点积注意力,可显著加速大型语言模型(如此处所示的nanoGPT)的训练。

增强模型数值稳定性

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

改进内存消耗

使用torch SDPA内核的另一个巨大优势是内存占用减少,这允许使用更大的批量大小。下表比较了flash attention和因果注意力的基线实现训练一小时后的最佳验证损失。可以看出,基线因果注意力实现(在8个NVIDIA Corporation A100服务器上,配备80 GB HBM)所能达到的最大批量大小为24,远小于flash attention所能达到的最大批量大小39。

Using Flash Attention enables the usage of larger batch sizes

图3:使用Flash Attention可以启用更大的批量大小,允许用户在训练一小时后获得更低的验证损失(越低越好)。

结论

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

  • 与基线相比,在批大小不变的情况下,训练速度提升超过26%。
  • 通过填充词汇量实现了额外的加速,使总优化比基线提高了约64%。
  • 额外的数值稳定性

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

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

# 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仓库中的以下提交开始——b3c17c6c6a363357623f223aaa4a8b1e89d0a465。此提交在测量每批次速度提升时用作基线。对于包含填充词汇优化(对批次速度提升最显著)的结果,请使用以下提交——77e7e04c2657846ddf30c1ca2dd9f7cbb93ddeab。从任一检出,使用torch.backends API选择用于实验的内核变得非常简单。

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

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