这篇文章是一个多系列博客的第二部分，重点介绍如何使用纯原生 PyTorch 加速生成式 AI 模型。我们很高兴能分享一系列新发布的 PyTorch 性能特性，并通过实际示例来展示我们能将 PyTorch 原生性能推向多远。在第一部分中，我们展示了如何仅使用纯原生 PyTorch 将 Segment Anything 加速超过 8 倍。在这篇博客中，我们将专注于大语言模型（LLM）的优化。
在过去的一年里，生成式 AI 的应用场景大受欢迎。文本生成是一个特别热门的领域，许多开源项目如 llama.cpp、 vLLM 和 MLC-LLM 都进行了大量创新。
虽然这些项目性能出色，但它们通常在易用性方面有所妥协，例如需要将模型转换为特定格式，或者需要构建和发布新的依赖项。这就引出了一个问题：仅使用纯原生 PyTorch，我们能以多快的速度运行 Transformer 推理？
正如我们在最近的 PyTorch 开发者大会上宣布的那样，PyTorch 团队从零开始编写了一个 LLM，其速度比基准快了近 10 倍，且没有任何精度损失，所有这些都使用了原生的 PyTorch 优化。我们利用了多种优化技术，包括：
- Torch.compile：一个 PyTorch 模型的编译器
- GPU 量化：通过降低精度的运算来加速模型
- 推测解码（Speculative Decoding）：使用一个小的“草稿”模型来预测大型“目标”模型的输出，从而加速 LLM
- 张量并行（Tensor Parallelism）：通过在多个设备上运行模型来加速模型。
而且，更棒的是，我们用不到 1000 行原生 PyTorch 代码就实现了这一切。
如果这让你兴奋得想直接看代码，请在此处查看： https://github.com/pytorch-labs/gpt-fast！
注意：所有这些基准测试都将关注延迟（即批量大小=1）。除非另有说明，所有基准测试均在 A100-80GB 上运行，功率限制为 330W。
起点（25.5 tok/s）
让我们从一个极其基础和简单的实现开始。
遗憾的是，这个实现的性能并不好。但为什么呢？查看追踪信息揭示了答案——它严重受到 CPU 开销的限制！这意味着我们的 CPU 无法足够快地告诉 GPU 该做什么，导致 GPU 未能被充分利用。
想象一下，GPU 就像一个拥有海量计算能力的超级工厂。然后，想象 CPU 像一个信使，来回给 GPU 传递指令。请记住，在大型深度学习系统中，GPU 负责完成 100% 的工作！在这类系统中，CPU 的唯一作用就是告诉 GPU 它应该做什么工作。
所以，CPU 跑过去告诉 GPU 做一个“加法”运算，但当 CPU 能给 GPU 下一个工作块时，GPU 早已完成了前一个工作块。
尽管 GPU 需要执行数千次计算，而 CPU 只需做协调工作，但这种情况出奇地普遍！原因多种多样，从 CPU 可能在运行某些单线程 Python，到如今 GPU 速度实在是太快了。
无论原因如何，我们现在发现自己处于开销限制（overhead-bound）的状态。那么，我们能做什么呢？一个方法是，我们可以用 C++ 重写我们的实现，甚至可能完全抛弃框架，直接编写原始的 CUDA 代码。或者……我们可以一次性向 GPU 发送更多的工作。
通过一次性发送一个巨大的工作块，我们就能让 GPU 保持忙碌！虽然在训练期间，这可能只需要增加批量大小就能实现，但在推理期间我们该怎么做呢？
进入 torch.compile。
第一步：通过 torch.compile 和静态 kv-cache 减少 CPU 开销（107.0 tok/s）
Torch.compile 允许我们将一个更大的区域捕获到一个单一的编译区域中，特别是在以 `mode=”reduce-overhead”` 模式运行时，它在减少 CPU 开销方面非常有效。在这里，我们还指定了 `fullgraph=True`，这可以验证模型中没有“图中断”（即 torch.compile 无法编译的部分）。换句话说，它确保 torch.compile 正在发挥其最大潜力。
要应用它，我们只需用它包装一个函数（或模块）。
torch.compile(decode_one_token, mode="reduce-overhead", fullgraph=True)
然而，这里有几个细微之处，使得人们在将 torch.compile 应用于文本生成时，要获得显著的性能提升并非易事。
第一个障碍是 kv-cache。kv-cache 是一种推理时优化，它会缓存先前词元计算出的激活值（更深入的解释请见这里）。然而，随着我们生成更多的词元，kv-cache 的“逻辑长度”会增长。这在两个方面存在问题。一是在每次缓存增长时重新分配（并复制！）kv-cache 的成本很高。另一个是，这种动态性使得减少开销变得更加困难，因为我们无法再利用像 cudagraphs 这样的方法。
为了解决这个问题，我们使用一个“静态”kv-cache，这意味着我们静态地分配 kv-cache 的最大尺寸，然后在计算的注意力部分掩码掉未使用的值。
第二个障碍是预填充（prefill）阶段。Transformer 文本生成最好被看作是一个两阶段过程：1. 预填充阶段，处理整个提示；2. 解码阶段，每个词元被自回归地生成。
尽管一旦 kv-cache 被设为静态，解码可以变得完全静态，但预填充阶段由于提示长度可变，仍然需要显著的动态性。因此，我们实际上需要用不同的编译策略来编译这两个阶段。
尽管这些细节有点棘手，但实际的实现并不困难（参见 gpt-fast）！而且性能提升是巨大的。
突然之间，我们的性能提升了超过 4 倍！当工作负载受开销限制时，这种性能提升通常很常见。
旁注：torch.compile 是如何提供帮助的？
有必要弄清楚 torch.compile 究竟是如何提升性能的。有两个主要因素导致了 torch.compile 的性能提升。
第一个因素，如上所述，是减少开销。Torch.compile 通过多种优化来减少开销，其中最有效的一种叫做CUDAGraphs。当设置 “reduce-overhead” 时，torch.compile 会自动应用此优化，省去了在没有 torch.compile 的情况下手动操作时需要编写的额外工作和代码。
然而，第二个因素是 torch.compile 能够生成更快的核心（kernel）。在上面的解码基准测试中，torch.compile 实际上从头开始生成了每一个核心，包括矩阵乘法和注意力！更酷的是，这些核心实际上比内置的替代品（CuBLAS 和 FlashAttention2）更快！
这可能对你们中的许多人来说听起来难以置信，考虑到编写高效的矩阵乘法/注意力核心是多么困难，以及在 CuBLAS 和 FlashAttention 上投入了多少人力。然而，这里的关键是，Transformer 解码具有非常不寻常的计算特性。特别是因为 KV-cache，对于 BS=1 的情况，*Transformer 中的每一次矩阵乘法实际上都是矩阵向量乘法*。
这意味着计算完全受*内存带宽限制*，因此，它们完全在编译器可以自动生成的范围内。事实上，当我们用 torch.compile 的矩阵向量乘法与 CuBLAS 进行基准测试时，我们发现 torch.compile 的核心实际上快得多！
第二步：通过 int8 权重-仅量化缓解内存带宽瓶颈（157.4 tok/s）
那么，既然我们已经通过应用 torch.compile 看到了巨大的速度提升，是否有可能做得更好呢？思考这个问题的一个方法是计算我们距离理论峰值有多近。在这种情况下，最大的瓶颈是从 GPU 全局内存加载权重到寄存器的成本。换句话说，每次前向传播都要求我们“触及”GPU 上的每一个参数。那么，理论上我们能以多快的速度“触及”模型中的每一个参数呢？
为了衡量这一点，我们可以使用**模型带宽利用率（Model Bandwidth Utilization, MBU）**。这衡量了我们在推理过程中能够使用多少内存带宽的百分比。
计算它非常简单。我们只需将模型的总大小（# 参数 * 每个参数的字节数）乘以我们每秒可以执行的推理次数。然后，我们将这个结果除以 GPU 的峰值带宽，就得到了我们的 MBU。
例如，在上述案例中，我们有一个 7B 参数的模型。每个参数都以 fp16（每个参数 2 字节）存储，我们达到了 107 tok/s。最后，我们的 A100-80GB 的理论内存带宽为 2 TB/s。
将这些数据综合起来，我们得到 **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 tok/s！
第三步：使用推测解码重新构建问题
即使使用了像量化这样的技术，我们仍然面临另一个问题。为了生成 100 个词元，我们必须加载我们的权重 100 次。
即使权重被量化了，我们仍然必须一遍又一遍地加载我们的权重，每生成一个词元就要加载一次！有什么办法可以绕过这个吗？
乍一看，答案似乎是否定的——在我们的自回归生成中存在严格的串行依赖。然而，事实证明，通过利用推测解码，我们能够打破这种严格的串行依赖并获得速度提升！
想象你有一位高级工程师（名叫 Verity），她能做出正确的技术决策，但写代码的速度比较慢。然而，你还有一位初级工程师（名叫 Drake），他并不总能做出正确的技术决策，但写代码的速度比 Verity 快得多（也更便宜！）。我们如何利用 Drake（初级工程师）来更快地写代码，同时确保我们仍然在做出正确的技术决策？
首先，Drake 经历编写代码的艰苦过程，并在途中做出技术决策。接下来，我们把代码交给 Verity 审查。
审查代码后，Verity 可能认为 Drake 做出的前 3 个技术决策是正确的，但最后 2 个需要重做。于是，Drake 回去，扔掉他最后的 2 个决策，并从那里重新开始编码。
值得注意的是，虽然 Verity（高级工程师）只看过一次代码，但我们能够生成 3 段与她自己会写的一模一样的经过验证的代码！因此，假设 Verity 审查代码的速度比她自己写那 3 段代码要快，这种方法就占了上风。
在 Transformer 推理的背景下，Verity 的角色将由一个更大的模型扮演，我们希望得到它的输出来完成我们的任务，这个模型被称为**验证模型（verifier model）**。同样，Drake 的角色将由一个更小的模型扮演，它能够比大模型快得多地生成文本，这个模型被称为**草稿模型（draft model）**。所以，我们会用草稿模型生成 8 个词元，然后用验证模型并行处理所有这 8 个词元，丢弃那些不匹配的。
如上所述，推测解码的一个关键特性是**它不改变输出的质量**。只要使用草稿模型生成词元和验证词元所花费的时间，少于直接生成这些词元所需的时间，我们就能获得优势。
在原生 PyTorch 中做这一切的一大好处是，这项技术实际上非常容易实现！这是完整的实现，大约 50 行原生 PyTorch 代码。
虽然推测解码保证了我们得到与常规生成在数学上相同的结果，但它的确有一个特性，即运行时性能会根据生成的文本以及草稿模型和验证模型的对齐程度而变化。例如，当运行 CodeLlama-34B + CodeLlama-7B 时，我们能够在生成代码时获得 2 倍的 tok/s 提升。另一方面，当使用 Llama-7B + TinyLlama-1B 时，我们只能获得大约 1.3 倍的 tok/s 提升。
旁注：在 AMD 上运行此程序
如上所述，解码中的每一个核心都是由 torch.compile 从头生成，并转换为 OpenAI Triton。由于 AMD 有一个torch.compile 后端（以及一个 Triton 后端），我们可以简单地将上述所有优化……在 AMD GPU 上实现！通过 int8 量化，我们能够用一个 GCD（即 MI250x 的一半）达到 102.5 tok/s！
第四步：通过 int4 量化和 GPTQ 进一步减小权重大小（202.1 tok/s）
当然，如果将权重从 16 位减少到 8 位能通过减少需要加载的字节数来提速，那么将权重减少到 4 位将会带来更大的速度提升！
不幸的是，当将权重降低到 4 位时，模型的精度成了一个更大的问题。从我们初步的评估来看，我们发现虽然使用 int8 权重-仅量化没有可察觉的精度下降，但使用 int4 权重-仅量化确实有。
我们可以使用两个主要技巧来限制 int4 量化的精度下降。
第一个是使用更精细的缩放因子。理解缩放因子的一种方式是，当我们有一个量化的张量表示时，它处于一个浮点张量（每个值都有一个缩放因子）和一个整数张量（没有值有缩放因子）之间的滑动尺度上。例如，对于 int8 量化，我们每行有一个缩放因子。然而，如果我们想要更高的精度，我们可以将其更改为“每 32 个元素一个缩放因子”。我们选择 32 的组大小以最小化精度下降，这也是社区中的一个常见选择。
另一个是使用比简单地对权重进行四舍五入更高级的量化策略。例如，像 GPTQ 这样的方法利用示例数据来更准确地校准权重。在这种情况下，我们在仓库中基于 PyTorch 最近发布的 torch.export 实现了一个 GPTQ 的原型。
此外，我们需要能将 int4 反量化与矩阵向量乘法融合的核心。在这种情况下，torch.compile 不幸无法从头生成这些核心，所以我们利用了 PyTorch 中一些手写的 CUDA 核心。
这些技术需要一些额外的工作，但将它们全部结合起来会带来更好的性能！
第五步：将所有技术结合起来（244.7 tok/s）
最后，我们可以将所有技术组合起来，以实现更好的性能！
第六步：使用张量并行
到目前为止，我们一直局限于在单个 GPU 上最小化延迟。然而，在许多设置中，我们有多个 GPU 可用。这使我们能够进一步改善延迟！
为了直观地理解为什么这能让我们改善延迟，让我们看看之前关于 MBU 的公式，特别是分母部分。在多个 GPU 上运行让我们能够获得更多的内存带宽，从而获得更高的潜在性能。
至于选择哪种并行策略，请注意，为了减少单个示例的延迟，我们需要能够同时利用更多设备上的内存带宽。这意味着我们需要将一个词元的处理分割到多个设备上。换句话说，我们需要使用张量并行。
幸运的是，PyTorch 也提供了可与 torch.compile 组合的低级张量并行工具。我们也在致力于开发用于表达张量并行的高级 API，敬请期待！
然而，即使没有高级 API，添加张量并行实际上仍然相当容易。我们的实现在150 行代码内完成，并且不需要任何模型更改。
我们仍然能够利用前面提到的所有优化，这些优化都可以继续与张量并行组合。将这些结合起来，我们能够以 int8 量化在 Llama-70B 上达到 55 tok/s 的服务速度！
结论
让我们来看看我们取得了哪些成就。
- 简洁性：忽略量化，model.py (244 LOC) + generate.py (371 LOC) + tp.py (151 LOC) 总共 766 行代码，实现了快速推理 + 推测解码 + 张量并行。
- 性能：对于 Llama-7B，我们能够使用编译 + int4 量化 + 推测解码达到 241 tok/s。对于 llama-70B，我们还能加入张量并行，达到 80 tok/s。这两个数字都接近或超过了 SOTA（业界顶尖）的性能数据！
PyTorch 一直以来都具备简洁性、易用性和灵活性。然而，通过 torch.compile，我们也能兼顾性能。
代码可以在这里找到：https://github.com/pytorch-labs/gpt-fast。我们希望社区会觉得它有用。我们这个仓库的目标不是提供另一个供人们导入的库或框架。相反，我们鼓励用户复制粘贴、复刻和修改仓库中的代码。
致谢
我们要感谢充满活力的开源社区对扩展 LLM 的持续支持，包括：
- Lightning AI 对 pytorch 以及在 flash attention、int8 量化和 LoRA 微调方面工作的支持。
- GGML 推动了 LLM 在设备上的快速推理。
- Andrej Karpathy 率先实现了简洁、可解释且快速的 LLM。
- MLC-LLM 推动了异构硬件上的 4 位量化性能。