PyTorch 加速生成式 AI 第二部分：高速 GPT [译]

by Team PyTorch

本篇博客是关于使用纯 PyTorch 加速生成式 AI 模型的系列文章的第二部分，由 PyTorch 团队撰写。我们在这里分享了 PyTorch 的最新性能特性，并通过实际案例，展示了如何最大限度地提升 PyTorch 的性能。在系列的第一篇文章中，我们演示了如何仅用 PyTorch 将“Segment Anything”加速超过 8 倍。本文将聚焦于大语言模型（LLM）的优化技术。

过去一年，生成式 AI 应用的人气迅速上升，尤其是文本生成领域。在这一领域，许多开源项目如 llama.cpp、vLLM、和 MLC-LLM 展示了创新成果。

这些项目虽然表现优异，但在易用性方面常常需要妥协，比如需要将模型转换成特定格式或添加新的依赖项。这就引出了一个问题：仅用原生 PyTorch 运行 Transformer 推理能达到多快的速度？

如我们在近期的 PyTorch 开发者大会上所宣布，我们的团队从头开始构建了一个 LLM，其运行速度几乎是原来的 10 倍，而且完全不损失准确性，这一切都得益于原生 PyTorch 的优化技术。我们采用了多种优化方法，包括：

Torch.compile：一种 PyTorch 模型编译器
GPU 量化：通过降低运算精度来加速模型
推测性解码：使用小型模型预测大型模型的输出，从而加速 LLM
张量并行性：通过多设备运行模型来实现加速

更令人兴奋的是，我们仅使用不到 1000 行的原生 PyTorch 代码就实现了这一切。

如果这让你兴奋到迫不及待想要研究代码，可以在这里找到：https://github.com/pytorch-labs/gpt-fast！

注意：我们在所有这些基准测试中将重点关注延迟（也就是说，批处理大小设为 1）。除非特别指出，所有测试都在配备 80GB 显存的 A100 显卡上进行，其功率被限制在 330W。

起点（每秒处理 25.5 个 Token）

我们从一个非常基础且简单的实现开始探索。

不幸的是，这种方法效果并不理想。原因何在呢？通过分析运行追踪，我们发现问题主要是由于CPU 处理瓶颈造成的！这意味着 CPU 无法足够迅速地向 GPU 发出指令，从而导致 GPU 无法得到充分的利用。

想象一下，GPU 好比一个拥有巨大计算能力的超级工厂，而 CPU 则像是一个忙碌地在 GPU 之间穿梭传递指令的信使。在大型深度学习系统中，所有重要工作都由 GPU 完成。而 CPU 的角色，只是指导 GPU 应该做什么。

比如，CPU 指示 GPU 执行一个加法操作，但在 CPU 能够分配给 GPU 下一项任务之前，GPU 早已完成了先前的任务。

虽然 GPU 需要执行成千上万的运算，而 CPU 只需要负责协调，但这种情况出奇的常见。原因多种多样，从 CPU 可能在运行某些单线程 Python 程序，到 GPU 速度极快。

无论是什么原因，我们现在面临一个处理开销过大的局面。那么我们该怎么办呢？一种方法是使用 C++ 重写程序，或者完全摒弃现有框架，直接编写原生的 CUDA 代码。或者……我们可以一次性向 GPU 发送更多的任务。

通过一次发送大量任务，我们可以确保 GPU 保持忙碌状态！尽管在训练阶段，这通常意味着增加批处理大小，但在推理过程中，我们该如何操作呢？

这时就需要使用 torch.compile 了。

第 1 步：利用 Torch.compile 和静态 KV-Cache 减少 CPU 负载（107.0 TOK/S）

Torch.compile 能够让我们将更广泛的区域编译成一个整体，尤其是在设置模式为“减少开销”（mode="reduce-overhead"）的情况下，它在降低 CPU 负载方面表现出色。此外，我们还设置 fullgraph=True，以确保模型中没有无法被 Torch.compile 编译的“断点”。这样做实际上是为了确保 Torch.compile 能够充分发挥其作用。

应用这一技术的方法是，我们简单地将某个函数或模块用 Torch.compile 包裹起来。

torch.compile(decode_one_token, mode="reduce-overhead", fullgraph=True)

然而，在将 Torch.compile 应用于文本生成过程中，存在一些细节上的难题，这使得用户想要通过这种方式显著提升性能变得不那么直接。

第一个障碍是 kv-cache。kv-cache 是在推理阶段对之前 Token 计算出的激活值进行缓存的一种优化方式（更详细的解释请见这里）。但随着我们生成越来越多的 Token，kv-cache 的“逻辑长度”也在增长。这主要有两个问题：一是每次缓存增长都需要重新分配（并且复制）kv-cache，成本较高；二是这种动态变化使得降低开销变得更加困难，因为我们不能再使用像 cudagraphs 这样的方法。

为了解决这个问题，我们采用了一种“静态”的 kv-cache，即我们为 kv-cache 分配一个固定的最大大小，然后在注意力计算过程中屏蔽掉未使用的部分。

第二个障碍是预填充阶段。可以将 Transformer 文本生成视为两个阶段的过程：1. 预填充阶段，处理整个提示内容；2. 解码阶段，自回归地生成每个 Token。

虽然一旦键值缓存（kv-cache）设为静态后，解码过程可以完全静态化，但由于提示长度的可变性，预填充阶段仍然需要较高的动态处理能力。因此，我们需要为这两个阶段分别采用不同的编译策略。

尽管这些技术细节听起来有些复杂，但实际上的实施过程并不复杂（具体可以参考 gpt-fast）。而且，性能的提升相当惊人。

我们的性能突然提高了超过 4 倍！在工作负载主要受到运行开销限制的情况下，这样的性能提升是比较常见的。

边注：torch.compile 如何提高性能？

探究 torch.compile 提升性能的具体方式是值得的。主要有两个因素促进了 torch.compile 的性能提升。

首先，像上文提到的，减少开销是一个关键因素。通过多种优化措施，Torch.compile 成功减少了开销，其中最有效的一个优化是名为 CUDAGraphs 的技术。当你设置为“reduce-overhead”时，torch.compile 会自动应用这种优化，这样就省去了你在没有使用 torch.compile 时手动执行这些操作所需的额外工作和编码。

然而，第二个提升性能的因素是 torch.compile 生成了更快速的内核。在上述的解码性能测试中，torch.compile 实际上是从头开始为每个过程，包括矩阵乘法和注意力机制，生成内核。更令人兴奋的是，这些内核实际上比现有的替代品（如 CuBLAS 和 FlashAttention2）更快！

虽然这听起来可能难以置信，考虑到编写高效的矩阵乘法/注意力内核的难度以及投入到 CuBLAS 和 FlashAttention 的大量人力资源。但是，transformer 解码有着非常独特的计算特性。特别是因为 KV 缓存，对于批量大小为 1（BS=1）的情况，transformer 中的每一个矩阵乘法实际上都是一个矩阵向量乘法。

这意味着计算过程完全受限于内存带宽，因此，编译器可以自动进行这些计算。事实上，当我们比较 torch.compile 的矩阵 - 向量乘法和 CuBLAS 时，我们发现 torch.compile 的内核实际上更加高效！

第 2 步：通过仅对权重进行 INT8 量化 (每秒处理 157.4 Token) 缓解内存带宽瓶颈

在我们已经通过使用 torch.compile 取得显著加速之后，我们还能做得更好吗？解决这个问题的一个思路是计算我们离理论最高效率有多远。在这种情况下，最大的障碍是把权重从 GPU 的全局内存传输到寄存器的成本。也就是说，每次前向传播都需要处理 GPU 上的每一个参数。那么，理论上我们能多快处理模型中的每一个参数呢？

为此，我们可以使用 模型带宽利用率（MBU） 来衡量。MBU 衡量了我们在推理过程中能够利用多少内存带宽的百分比。

计算方法很简单。我们只需要计算模型的总大小（参数数量 * 每个参数的字节数）乘以我们每秒能进行的推理次数。然后，将这个值除以 GPU 的峰值带宽，就得到了 MBU。

比如，在我们的案例中，模型有 70 亿个参数，每个参数以 fp16 格式存储（每个参数 2 字节），我们实现了每秒处理 107 Token 的速度。最后，我们使用的 A100-80GB GPU 理论上有 2 TB/秒的内存带宽。

综合以上数据，我们得到了 72% 的 MBU！ 这个结果相当不错，因为即便是简单的内存复制也很难超过 85%。

但这也意味着我们已经非常接近理论上的限制，而且明显受制于从内存加载权重。不改变问题的前提下，我们可能只能再提高大约 10% 的性能。

我们再来看看前面提到的方程式。实际上，我们无法改变模型中的参数数量，也无法真正改变 GPU 的内存带宽（除非愿意多花钱）。但是，我们确实可以改变每个参数所占用的字节数！

这就引出了我们的下一个技术策略——int8 量化。其核心思想很简单：如果内存加载权重是主要瓶颈，为何不试着缩减权重的大小呢？

需要注意的是，这里只对权重进行量化——计算本身依然采用 bf16。这种量化方式简单易行，几乎不会影响精确度。

而且，torch.compile 还可以轻松生成高效的 int8 量化代码。让我们再次回顾之前的基准测试，这次我们加入了 int8 的权重量化。

如图中的深蓝线（torch.compile + int8）所示，使用 torch.compile 结合 int8 的权重量化可以显著提升性能！而浅蓝线（未使用 torch.compile 的 int8）的性能甚至不如 fp16！这是因为，要充分利用 int8 量化的性能优势，需要内核融合。这正体现了 torch.compile 的优势——它能为用户自动产生这些内核！

在我们的模型中应用 int8 量化后，我们看到了显著的 50% 性能提升，速度提高到了 157.4 tokens/s！

第 3 步：运用推测性解码重新定义问题

即便我们运用了量化等技术，我们仍面临一个难题。为了生成 100 个 Token，我们不得不重复 100 次加载权重的过程。

权重即便已经量化，每生成一个 Token，就必须再次加载权重。那么，我们有没有办法突破这一限制呢？

初看起来，似乎没有办法 - 因为我们的自回归生成过程中存在着严格的串行依赖。但实际上，通过应用推测性解码，我们可以打破这种串行依赖，从而实现加速！

设想你拥有一位高级工程师 Verity，她能做出正确的技术选择，但编写代码的速度较慢。与此同时，你还有一位初级工程师 Drake，他在技术决策上可能不那么精准，但编写代码的速度比 Verity 快得多（而且成本更低！）。我们怎样才能利用 Drake 的快速编码能力，同时确保技术决策的正确性呢？

首先，Drake 负责编写代码并做出一系列技术决策，这是一个劳动密集的过程。然后，我们将代码交给 Verity 进行审核。

审核过程中，Verity 可能会认为 Drake 做出的前三个技术决策是正确的，但最后两个需要重新处理。于是，Drake 返回，废弃他之前的最后两个决策，并从这个点重新开始编码。

值得关注的是，虽然 Verity（高级工程师）只检查了一次代码，但我们已经生成了三段相当于她亲自编写的代码！因此，如果 Verity 的审核速度比她亲自编写这三段代码的速度要快，这种方法就能取得优势。

在讨论 Transformer 模型的推理过程中，Verity 就像是我们任务所需的大型模型，被称为验证模型。而 Drake 则相当于一个小型模型，它能比大型模型更快速地生成文本，这就是所谓的草稿模型。因此，我们会先用草稿模型生成 8 个 Token，接着用验证模型同时处理这些 Token，剔除不符合的部分。

如前所述，推测性解码的一个重要特点是不会影响输出质量。只要草稿模型生成 Token 和验证所需时间之和少于单独生成这些 Token 的时间，就算是成功。

一个好消息是，使用原生 PyTorch 实现这个过程非常简单！以下是这个技术的完整实现代码，只需大约 50 行代码。

尽管推测性解码在数学上保证了与常规生成完全相同的结果，但它的运行效率会根据生成的文本和草稿模型与验证模型的匹配程度而有所不同。例如，在使用 CodeLlama-34B 加上 CodeLlama-7B 的组合时，我们在代码生成方面实现了 2 倍的 Token/s 提升。但在使用 Llama-7B 加上 TinyLlama-1B 的组合时，提速只有大约 1.3 倍。

附注：在 AMD 上的运行情况

如前所述，解码过程中每个内核都是由 torch.compile 从头开始创建，并转换为 OpenAI Triton。鉴于 AMD 提供了 torch.compile 后端（也支持 Triton），我们可以在 AMD GPU 上应用上述所有优化。通过 int8 量化，我们在一个 MI250x GPU 的一半处理能力上实现了 102.5 Token/s 的速度！

第 4 步：通过 INT4 量化和 GPTQ 进一步缩小权重尺寸，实现更快处理速度（202.1 Token/秒）

显然，如果将数据的权重从 16 位降至 8 位能通过减少加载的数据量来加快速度，那么进一步将权重降至 4 位自然能带来更显著的速度提升！

但是，将权重降至 4 位时，模型的准确度问题就变得更加严重。我们的初步评估显示，虽然使用 int8 权重量化基本上不会影响模型的准确度，但使用 int4 权重量化就会有明显的准确度下降。

为了限制 int4 量化对准确度的影响，我们采用了两种主要策略。

第一种是使用更精细的缩放因子。简单来说，缩放因子就像是在浮点数张量（每个值都有一个缩放因子）和整数张量（没有缩放因子）之间的一个中间值。例如，在 int8 量化中，我们对每行使用一个缩放因子。但若追求更高的准确度，我们可以改为“每 32 个元素使用一个缩放因子”。我们选择 32 个元素作为一组，这样做可以最大限度地减少准确度的损失，这也是业界通常的做法。

第二种策略是采用比简单四舍五入更先进的量化策略。比如，GPTQ 方法利用示例数据来更准确地校准权重。在这种情况下，我们在 PyTorch 的 torch.export 基础上原型实现了 GPTQ。

此外，我们还需要一些特殊的内核，这些内核能将 int4 反量化与矩阵向量乘法结合起来。在这里，torch.compile 无法直接生成这些内核，所以我们使用了 PyTorch 中一些手写的 CUDA 内核。

虽然这些技术需要额外的工作量，但将它们结合使用，最终能带来更出色的性能表现！

步骤 5: 综合应用各种技术 (每秒处理 244.7 Token)

在最后一步，我们将前面提到的所有技术融合起来，目的是实现更高的性能表现！

步骤 6: 引入张量并行技术

之前，我们的工作重点是在单个 GPU 上尽量减少处理延时。但在很多情况下，我们可以使用多个 GPU，这样就可以进一步减少延时。

要理解为什么使用多个 GPU 可以减少延时，我们可以回顾一下之前的 MBU（内存带宽利用率）公式，尤其是分母部分。多 GPU 运行模式意味着我们能够使用更多的内存带宽，从而实现更高的性能。

在选择并行策略时，需要注意的是，为了减少单个样例的延时，我们必须能够同时在更多设备上利用内存带宽。这意味着我们需要将单个 Token 的处理分布到多个设备上。也就是说，我们需要采用张量并行策略。

幸运的是，PyTorch 不仅提供了实现张量并行的底层工具，而且这些工具能与 torch.compile 无缝结合。我们还在开发更高级的张量并行 API，敬请期待！

即便现在还没有更高级的 API，实现张量并行也是相对简单的。我们的实现代码只有 150 行，而且不需要对模型做任何修改。

通过利用之前所有提到的优化措施，并将它们与张量并行结合，我们成功地实现了对 Llama-70B 模型的服务，速度达到每秒处理 55 个 Token，同时还采用了 int8 量化技术！

结论

来看看我们达到的成就。

简洁性：如果不考虑量化，只需 model.py（244 行代码）、generate.py（371 行代码）和 tp.py（151 行代码），共 766 行代码就能实现快速推断、预测解码和张量并行。
性能：使用 Llama-7B 模型，我们结合编译、int4 量化和预测解码技术，达到了每秒 241 个 token 的速度。而使用 Llama-70B 模型，我们还加入了张量并行技术，速度达到每秒 80 个 token。这些成绩都达到或超越了当前最先进的水平！

PyTorch 一直以其简洁、易用和灵活著称。现在，加上 torch.compile 后，它的性能也得到了显著提升。

相关代码可以在这里找到：https://github.com/pytorch-labs/gpt-fast。我们希望社区能从中受益。我们创建这个代码仓库的目的并不是提供一个新的库或框架供人导入使用，而是鼓励用户自行复制、修改和改进这些代码。

致谢

我们非常感谢活跃的开源社区对扩大大语言模型 (LLM) 规模的持续支持，尤其是以下团队和个人：

Lightning AI，支持了 PyTorch 以及在快速注意力机制、int8 量化和 LoRA 微调方面的研究。
GGML，推动了大语言模型在设备上的快速推理技术。
Andrej Karpathy，他在简单、易理解和快速的大语言模型实现方面做出了重要贡献。
MLC-LLM，推动了异构硬件上的 4 位量化性能。