本文是系列博客的第二篇，专注于如何利用纯原生 PyTorch 加速生成式 AI 模型。我们很高兴分享一系列新发布的 PyTorch 性能特性，并通过实际示例展示 PyTorch 原生性能的极限。在第一篇中，我们展示了如何仅使用纯原生 PyTorch 将 Segment Anything 的速度提升超过 8 倍。本篇博客将重点讨论 LLM（大语言模型）的优化。
在过去的一年中，生成式 AI 的应用场景呈爆发式增长。文本生成是一个特别热门的领域，开源项目中涌现了大量创新，如 llama.cpp、vLLM 和 MLC-LLM。
虽然这些项目性能优异，但它们往往在易用性上有所折衷，例如需要将模型转换为特定格式，或者需要构建和集成新的依赖项。这就引发了一个问题：仅使用纯原生 PyTorch，我们能将 Transformer 推理速度提升到什么程度？
正如我们在最近的 PyTorch 开发者大会上所宣布的那样，PyTorch 团队从零开始构建了一个 LLM，在没有任何精度损失的情况下，利用原生 PyTorch 优化实现了比基准测试快近 10 倍的性能。我们利用了一系列优化手段，包括：
- Torch.compile：用于 PyTorch 模型的编译器
- GPU 量化：通过降低精度运算来加速模型
- 推测解码 (Speculative Decoding)：使用较小的“草稿”模型来预测大型“目标”模型的输出，从而加速 LLM
- 张量并行 (Tensor Parallelism)：通过跨多个设备运行模型来加速模型
更棒的是，我们只需 不到 1000 行的原生 PyTorch 代码 就能做到这一点。
如果你迫不及待想直接查看代码，请访问：https://github.com/pytorch-labs/gpt-fast！
注意：我们所有的基准测试都聚焦于延迟（即 batch size=1）。除非另有说明，所有基准测试均在 A100-80GB GPU 上运行，功耗限制为 330W。
起点（25.5 tok/s）
让我们从一个极其基础和简单的实现开始。
遗憾的是，它的表现并不理想。这是为什么呢？通过查看追踪记录可以发现答案——它严重受限于 CPU 开销 (CPU overhead bound)！这意味着 CPU 无法足够快地向 GPU 发出指令，导致 GPU 无法得到充分利用。
把 GPU 想象成一个拥有巨大算力的超级工厂，而 CPU 是一个在工厂和任务间传递指令的信使。记住，在大规模深度学习系统中，GPU 负责完成 100% 的工作！在这样的系统中，CPU 唯一的角色就是告诉 GPU 该做什么。
所以，CPU 跑过去告诉 GPU 执行一次“加法”，但等 CPU 能传达下一个任务时，GPU 早已完成了之前的工作。
尽管 GPU 需要执行数千次计算，而 CPU 只需处理编排工作，但这种情况却出奇地常见！造成这种情况的原因有很多，从 CPU 可能正在运行某些单线程 Python 代码，到如今 GPU 的处理速度快得惊人等等。
无论原因如何，我们现在处于 开销受限 (overhead-bound) 的状态。那么，我们能做什么呢？一种方法是用 C++ 重写，甚至抛弃框架直接写原生 CUDA。或者……我们可以一次性向 GPU 发送更多任务。
通过一次性发送大量工作，我们可以让 GPU 保持忙碌！虽然在训练期间可以通过增加 batch size 来实现，但在推理期间我们该怎么做呢？
引入 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(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 的最大容量，然后在计算的注意力部分屏蔽掉未使用的值。
第二个障碍是预填充 (prefill) 阶段。Transformer 文本生成最好看作是一个两阶段过程：1. 预填充，处理整个提示词；2. 解码，自回归地生成每个 token。
虽然 kv-cache 静态化后解码阶段可以完全静态化，但预填充阶段由于提示词长度可变，仍然需要更多的动态性。因此，我们需要使用不同的编译策略来编译这两个阶段。
尽管这些细节比较棘手，但实际实现并不困难（参见 gpt-fast）！而且带来的性能提升非常显著。
瞬间，我们的性能提升了 4 倍多！当工作负载受限于开销时，这种性能提升非常常见。
附注：torch.compile 是如何起到帮助的？
值得拆解一下 torch.compile 到底是如何提升性能的。导致其性能提升的主要因素有两个。
第一个因素正如上面提到的，是开销减少。Torch.compile 能够通过各种优化手段减少开销，其中最有效的一种称为 CUDAGraphs。当设置 “reduce-overhead” 时，torch.compile 会自动为你应用此优化，无需你手动编写额外代码。
然而，第二个因素是 torch.compile 能够生成更快的内核 (kernels)。在上面的解码基准测试中，torch.compile 实际上是从零开始生成每一个内核，包括矩阵乘法和注意力机制！更酷的是，这些内核实际上比内置的替代方案（CuBLAS 和 FlashAttention2）更快！
考虑到编写高效的矩阵乘法/注意力内核有多困难，以及 CuBLAS 和 FlashAttention 投入了多少人力，这听起来可能令人难以置信。但这里的关键在于，Transformer 解码具有非常独特的计算属性。特别是由于 KV-cache 的存在，对于 BS=1 的情况，Transformer 中的每一个矩阵乘法实际上都是矩阵-向量乘法。
这意味着这些计算完全是 内存带宽受限的，因此完全在编译器自动生成的范围之内。事实上，当我们对比 torch.compile 的矩阵-向量乘法与 CuBLAS 时，我们发现 torch.compile 的内核实际上要快得多！
第 2 步：通过 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%，这个表现已经相当不错了。
但是……这意味着我们已经非常接近理论极限，而且显而易见，我们的瓶颈就在于将权重从内存加载到 GPU。不管我们做什么——除非以某种方式改变问题设定，否则我们可能只能再挤出 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！
第 3 步：使用推测解码 (Speculative Decoding) 重构问题
即使在使用量化等技术之后，我们仍然面临另一个问题：为了生成 100 个 token，我们必须加载 100 次权重。
即使权重经过了量化，我们仍然必须反复加载它们，每个生成的 token 都要加载一次！有什么办法可以解决这个问题吗？
乍一看，答案似乎是否定的——我们的自回归生成存在严格的串行依赖。然而，事实证明，利用 推测解码 (Speculative Decoding)，我们可以打破这种严格的串行依赖并获得加速！
想象一下，你有一位资深工程师（叫 Verity），她能做出正确的决策但写代码很慢。同时，你有一位初级工程师（叫 Drake），他并不总是能做出正确的决策，但写代码比 Verity 快得多（也更便宜！）。我们如何利用 Drake（初级工程师）来更快地编写代码，同时确保我们仍然做出正确的技术决策？
首先，Drake 进行繁重的编码工作，并在过程中做出技术决策。接下来，我们将代码交给 Verity 审查。
审查代码时，Verity 可能认为 Drake 前 3 个技术决策是正确的，但后 2 个需要重做。于是，Drake 回去，抛弃他最后 2 个决定，并从那里开始重新编码。
值得注意的是，虽然 Verity（资深工程师）只看了一次代码，但我们却能生成 3 个与她原本会写的一模一样的经验证的代码片段！因此，假设 Verity 审查代码的速度比她自己写这 3 个片段要快，那么这种方法就是成功的。
在 Transformer 推理的语境下，Verity 的角色由我们想要获得输出的大模型承担，称为 验证模型 (verifier model)。同样，Drake 的角色由一个能比大模型生成文本快得多的较小模型承担，称为 草稿模型 (draft model)。因此，我们先使用草稿模型生成 8 个 token，然后使用验证模型并行处理所有八个 token，剔除不匹配的部分。
正如上面提到的，推测解码的一个关键属性是 它不会改变输出的质量。只要“使用草稿模型生成 token + 验证 token”所花费的时间少于原本生成这些 token 的时间，我们就是获利的。
在纯原生 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 量化，我们能够在使用 MI250x 的一个 GCD（即一半）的情况下达到 102.5 tok/s！
第 4 步：通过 int4 量化和 GPTQ 进一步减小权重大小（202.1 tok/s）
当然，如果将权重从 16 位减少到 8 位可以通过减少加载的字节数来提速，那么将权重减少到 4 位应该会导致更大的提速！
不幸的是，当将权重降低到 4 位时，模型的准确性开始成为一个更大的问题。根据我们的初步评估，虽然使用 int8 权重仅量化没有可察觉的精度下降，但使用 int4 权重仅量化则会有。
我们可以使用 2 个主要的技巧来限制 int4 量化的精度下降。
第一个是使用更细粒度的缩放因子。考虑缩放因子的一种方式是，当我们拥有量化的张量表示时，它处于浮点张量（每个值都有一个缩放因子）和整数张量（没有值有缩放因子）之间的滑动刻度上。例如，在 int8 量化中，我们每一行使用一个缩放因子。如果想要更高的精度，我们可以将其改为“每 32 个元素使用一个缩放因子”。我们选择 32 作为组大小以最小化精度下降，这也是社区中常见的选择。
另一个是使用比简单的四舍五入权重更先进的量化策略。例如，像 GPTQ 这样的方法利用示例数据来更准确地校准权重。在这种情况下，我们基于 PyTorch 最近发布的 torch.export 在代码库中开发了 GPTQ 的原型实现。
此外，我们需要融合了 int4 反量化与矩阵-向量乘法的内核。在这种情况下，torch.compile 不幸无法从零生成这些内核，因此我们利用了一些 PyTorch 中的手写 CUDA 内核。
这些技术需要一些额外的工作，但将它们结合在一起可以获得更好的性能！
第 5 步：将所有内容结合在一起（244.7 tok/s）
最后，我们可以将所有技术结合起来，以实现更好的性能！
第 6 步：使用张量并行
到目前为止，我们一直将自己局限于在单个 GPU 上最小化延迟。然而，在许多设置中，我们可以访问多个 GPU。这使我们能够进一步改善延迟！
为了直观地理解为什么这能提高延迟，让我们看看之前 MBU 的公式，特别是分母。在多个 GPU 上运行使我们能够访问更多的内存带宽，从而获得更高的潜在性能。
至于选择哪种并行策略，请注意，为了减少单个样本的延迟，我们需要能够同时利用多个设备上的内存带宽。这意味着我们需要将一个 token 的处理过程拆分到多个设备上。换句话说，我们需要使用张量并行。
幸运的是，PyTorch 还提供了可与 torch.compile 组合使用的底层张量并行工具。我们也在开发用于表达张量并行的高级 API，敬请期待！
不过，即使没有高级 API，添加张量并行也相当容易。我们的实现仅需 150 行代码，且不需要任何模型更改。
我们仍然能够利用之前提到的所有优化，这些优化都可以继续与张量并行组合。将这些结合起来，我们能够以 55 tok/s 的速度服务 Llama-70B 模型（使用 int8 量化）！
结论
让我们看看我们实现了什么。
- 简单性：忽略量化，model.py (244 行) + generate.py (371 行) + tp.py (151 行)，总共 766 行代码实现了快速推理 + 推测解码 + 张量并行。
- 性能：使用 Llama-7B，我们能够利用编译 + int4 量化 + 推测解码达到 241 tok/s。使用 Llama-70B，我们还加入了张量并行，达到 80 tok/s。这两者都接近或超过了业界最先进的性能数据！
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-bit 量化性能的极限。