本文是系列博客的第二部分，重点介绍如何使用纯原生 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 量化：使用降低精度的操作加速模型
- 推测解码：使用一个小的“草稿”模型预测大的“目标”模型的输出，从而加速 LLMs
- 张量并行：通过在多个设备上运行模型来加速它们。
更好的是，我们可以在不足 1000 行的原生 PyTorch 代码中实现这一点。
如果您对此感到兴奋并想立即查看代码，请访问 https://github.com/pytorch-labs/gpt-fast！
注意：所有这些基准测试都将侧重于延迟（即批处理大小=1）。除非另有说明，所有基准测试都在一块 A100-80GB GPU 上运行，功耗限制为 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 实在太快了。
无论原因如何，我们现在都处于开销受限状态。那么，我们能做什么呢？第一，我们可以用 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(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-cache，对于 BS=1 的情况，Transformer 中的每一个矩阵乘法实际上都是一个矩阵向量乘法。
这意味着计算完全是内存带宽受限的，因此编译器完全有能力自动生成相应的内核。事实上，当我们对 torch.compile 的矩阵向量乘法与 CuBLAS 进行基准测试时，我们发现 torch.compile 的内核实际上要快不少！
步骤 2：通过 int8 仅权重量化缓解内存带宽瓶颈 (157.4 tok/s)
那么，既然我们已经看到了应用 torch.compile 带来的巨大加速，是否还有可能做得更好呢？思考这个问题的一种方式是计算我们离理论峰值有多近。在这种情况下，最大的瓶颈是将权重从 GPU 全局内存加载到寄存器的成本。换句话说，每一次前向传递都需要我们“触碰”GPU 上的每一个参数。那么，理论上我们可以多快地“触碰”模型中的每一个参数呢？
为了衡量这一点，我们可以使用模型带宽利用率 (MBU)。它衡量我们在推理期间能够使用内存带宽的百分比。
计算它非常简单。我们只需计算模型的总大小（# 参数 * 每个参数的字节数），然后乘以每秒可以进行的推理次数。最后，将结果除以 GPU 的峰值带宽，即可得出我们的 MBU。
例如，在上面的例子中，我们有一个 7B 参数的模型。每个参数以 fp16 存储（每个参数 2 字节），我们实现了 107 tokens/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 tokens/s！
步骤 3：使用推测解码重新定义问题
即使使用了量化等技术后，我们仍然面临另一个问题。为了生成 100 个 token，我们必须加载权重 100 次。
即使权重经过量化，我们仍然必须一次又一次地加载权重，每生成一个 token 就要加载一次！有没有什么办法可以避免这种情况？
乍一看，答案似乎是否定的——我们的自回归生成存在严格的序列依赖。然而，事实证明，通过利用推测解码，我们能够打破这种严格的序列依赖并获得加速！
想象一下你有一个高级工程师（姑且称之为 Verity），他能做出正确的技术决定，但写代码比较慢。然而，你还有一个初级工程师（姑且称之为 Drake），他并不总是能做出正确的技术决定，但他写代码比 Verity 快得多（也便宜得多！）。我们如何利用 Drake（这位初级工程师）更快地写代码，同时确保我们仍然做出正确的技术决定？
首先，Drake 经历编写代码的劳动密集型过程，在此过程中做出技术决定。接下来，我们将代码交给 Verity 进行审查。
审查代码后，Verity 可能会决定 Drake 做出的前 3 个技术决定是正确的，但最后 2 个需要重做。于是，Drake 回去，放弃他最后的 2 个决定，然后从那里重新开始编码。
值得注意的是，尽管 Verity（这位高级工程师）只看了一次代码，但我们能够生成 3 段经过验证的代码，这些代码与她自己编写的完全相同！因此，假设 Verity 审查代码的速度比她自己编写这 3 段代码的速度快，那么这种方法就更具优势。
在 Transformer 推理的背景下，Verity 扮演的是我们希望为其任务获取输出的更大模型的角色，称为验证模型。类似地，Drake 扮演的是一个能够比更大模型更快生成文本的小模型的角色，称为草稿模型。因此，我们将使用草稿模型生成 8 个 token，然后使用验证模型并行处理这 8 个 token，并丢弃不匹配的 token。
如上所述，推测解码的一个关键特性是它不会改变输出的质量。只要使用草稿模型生成 token + 验证 token 所花费的时间少于直接生成这些 token 所花费的时间，我们就能获得优势。
在原生 PyTorch 中完成所有这些工作的一个巨大优势是，这项技术实际上非常容易实现！这是整个实现的全部代码，大约 50 行原生 PyTorch 代码。
虽然推测解码保证了结果与常规生成在数学上完全相同，但其运行时性能会因生成的文本以及草稿模型和验证模型的一致程度而异。例如，在使用 CodeLlama-34B + CodeLlama-7B 组合时，生成代码的 token/s 性能可以提升 2 倍。另一方面，在使用 Llama-7B + TinyLlama-1B 组合时，token/s 性能仅能提升约 1.3 倍。
旁注：在 AMD 平台上运行
如前所述，解码中的每个内核都是由 torch.compile 从零开始生成的，并转换为 OpenAI Triton。由于 AMD 具有torch.compile 后端（以及 Triton 后端），我们可以简单地应用上述所有优化……但在 AMD GPU 上运行！通过 int8 量化，我们可以在一个 MI250x 的一个 GCD（即一半）上达到 102.5 tokens/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 行，并且不需要对模型进行任何更改。
我们仍然能够利用之前提到的所有优化，所有这些优化都可以继续与张量并行结合使用。将这些结合在一起，我们能够在进行 int8 量化的情况下，以 55 tokens/s 的速度提供 Llama-70B 的服务！
结论
让我们看看我们能够实现什么。
- 简洁性：忽略量化，model.py (244 行代码) + generate.py (371 行代码) + tp.py (151 行代码) 共计 766 行代码，实现了快速推理 + 推测解码 + 张量并行。
- 性能：对于 Llama-7B，我们能够使用 compile + int4 量化 + 推测解码达到 241 tok/s。对于 llama-70B，我们还加入了张量并行，达到 80 tok/s。这两个性能指标都接近或超越了 SOTA 性能！
PyTorch 一直以来都以简洁性、易用性和灵活性著称。然而，有了 torch.compile，我们还可以同时获得高性能。
代码可以在这里找到：https://github.com/pytorch-labs/gpt-fast。我们希望社区能从中受益。我们创建这个仓库的目的不是提供另一个供人们导入的库或框架。相反，我们鼓励用户复制、Fork 和修改仓库中的代码。
致谢
我们要感谢充满活力的开源社区对 LLMs 扩展的持续支持，包括
- Lightning AI，感谢其对 PyTorch 以及在 flash attention、int8 量化和 LoRA 微调方面的工作的支持。
- GGML，感谢其推动 LLMs 的快速设备端推理。
- Andrej Karpathy，感谢其率先提出简单、可解释且快速的 LLM 实现。
- MLC-LLM，感谢其在异构硬件上提升 4 位量化性能。
