这篇文章是关于如何使用纯、原生的 PyTorch 加速生成式 AI 模型的系列博客的第二部分。我们很高兴地分享一系列新发布的 PyTorch 性能特性，并通过实际示例展示我们可以将 PyTorch 原生性能提升到何种程度。在第一部分中，我们展示了如何仅使用纯、原生的 PyTorch 将 Segment Anything 加速超过 8 倍。在本篇博客中，我们将重点关注LLM优化。
在过去的一年里，生成式 AI 的应用场景在流行度上有了爆炸性的增长。文本生成是一个特别受欢迎的领域，在诸如 llama.cpp、vLLM 和 MLC-LLM等开源项目中，创新层出不穷。
虽然这些项目性能出色，但它们通常在易用性方面有所妥协，例如需要将模型转换为特定格式或构建和分发新的依赖项。这引发了一个问题：我们能否仅使用纯、原生的 PyTorch 实现 Transformer 推理的快速运行？
如在我们最近的 PyTorch 开发者大会上所宣布的，PyTorch 团队从头开始编写了一个LLM，其速度比基线快近 10 倍，且没有损失精度，全部使用原生 PyTorch 优化。我们利用了一系列优化，包括：
- PyTorch 编译器：PyTorch 模型的编译器
- GPU 量化：通过降低精度操作加速模型
- 推测性解码：使用小型“草案”模型预测大型“目标”模型的输出
- 张量并行：通过在多个设备上运行模型来加速模型
甚至更好，我们可以在不到 1000 行的原生 PyTorch 代码中完成它。
如果这足以让你直接跳入代码，请查看 https://github.com/pytorch-labs/gpt-fast!
注意：我们将专注于所有这些基准的延迟（即 batch size=1）。除非另有说明，所有基准均在 A100-80GB 上运行，功率限制为 330W。
起始点（25.5 tok/s）
让我们从一个非常基础和简单的实现开始。
很遗憾，这个实现表现并不好。但是为什么？查看跟踪信息就能找到答案——它严重受 CPU 开销限制！这意味着我们的 CPU 无法快速地告诉 GPU 该做什么，以至于 GPU 无法得到充分利用。
想象一下 GPU 就像一个超级巨大的工厂，拥有大量的计算能力。然后，想象 CPU 就像一个来回传递指令的信使。记住，在大规模深度学习系统中，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 是一个推理时优化，它缓存了之前标记计算出的激活（详见此处获取更深入的说明）。然而，随着我们生成更多的标记，“逻辑长度”的 kv-cache 会增长。这有两个问题。一是每次缓存增长时重新分配（以及复制！）kv-cache 简直太昂贵了。另一个问题是这种动态性使得减少开销变得更加困难，因为我们不再能够利用像 cudagraphs 这样的方法。
为了解决这个问题，我们使用一个“静态”的 kv-cache，这意味着我们静态地分配 kv-cache 的最大大小，然后在计算注意力的部分屏蔽未使用的值。
第二个障碍是预填充阶段。将 Transformer 文本生成视为一个两阶段过程：1. 预填充阶段，处理整个提示，2. 解码阶段，逐个自动回归生成标记。
虽然解码一旦将 kv-cache 设置为静态后就可以完全静态化，但预填充阶段仍然需要更多的动态性，因为提示长度是可变的。因此，我们实际上需要使用不同的编译策略来编译这两个阶段。
虽然这些细节有点棘手，但实际实现并不困难（参见 gpt-fast）！性能提升非常显著。
突然之间，我们的性能提高了 4 倍以上！当一个人的工作负载是瓶颈时，这种性能提升通常是常见的。
简单来说：torch.compile 是如何帮助的？
值得深入探讨 torch.compile 是如何提高性能的。导致 torch.compile 性能提升的主要因素有两个。
首先，正如上文所述，是减少开销。torch.compile 能够通过各种优化来减少开销，但其中最有效的一种叫做 CUDAGraphs。尽管 torch.compile 在设置“reduce-overhead”时自动应用此优化，节省了您手动进行此操作时需要编写的额外工作和代码。
然而，第二个因素是 torch.compile 简单地生成了更快的内核。在上面的解码基准测试中，torch.compile 实际上从头开始生成每个内核，包括矩阵乘法和注意力！更酷的是，这些内核实际上比内置的替代方案（CuBLAS 和 FlashAttention2）更快！
对于很多人来说，这可能听起来有些不可思议，考虑到编写高效的矩阵乘法/注意力内核有多难，以及 CuBLAS 和 FlashAttention 投入了多少人力。然而，关键在于，transformer 解码具有非常独特的计算特性。特别是，由于 KV 缓存，对于 BS=1，transformer 中的每一次矩阵乘法实际上都是矩阵向量乘法。
这意味着计算完全受内存带宽限制，因此非常适合编译器自动生成。事实上，当我们将 torch.compile 的矩阵向量乘法与 CuBLAS 进行基准测试时，我们发现 torch.compile 的内核实际上要快得多！
第 2 步：通过 int8 权重量化缓解内存带宽瓶颈（157.4 tok/s）
那么，鉴于我们已经从应用 torch.compile 中看到了巨大的速度提升，是否还能做得更好呢？思考这个问题的方法之一是计算我们距离理论峰值有多近。在这种情况下，最大的瓶颈是从 GPU 全局内存到寄存器的权重加载成本。换句话说，每次前向传播都需要我们“接触”GPU 上的每个参数。那么，理论上我们以多快的速度“接触”模型中的每个参数呢？
为了测量这一点，我们可以使用模型带宽利用率（MBU）。这衡量了我们能在推理过程中使用多少内存带宽。
计算它相当简单。我们只需将模型的总体大小（参数数量 * 每个参数的字节数）乘以我们每秒可以进行的推理次数。然后，我们将这个结果除以 GPU 的峰值带宽，以获得我们的 MBU。
例如，在我们的上述案例中，我们有一个 70 亿参数的模型。每个参数以 fp16 存储（每个参数 2 字节），我们实现了每秒 107 个 token。最后，我们的 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 个 token！
第 3 步：使用推测性解码重新定义问题
即使使用了量化等技术，我们仍然面临另一个问题。为了生成 100 个标记，我们必须加载我们的权重 100 次。
即使权重已经量化，我们仍然必须反复加载我们的权重，每次生成一个标记时都要加载一次！有没有什么解决办法？
初看之下，答案可能是否定的——在我们的自回归生成中存在严格的序列依赖。然而，实际上，通过利用推测性解码，我们能够打破这种严格的序列依赖，并获得加速！
想象一下，你有一个资深工程师（称为 Verity），她在做出正确的技术决策方面做得很好，但编写代码的速度较慢。然而，你还有一个初级工程师（称为 Drake），他并不总是做出正确的技术决策，但他的代码编写速度（和成本）比 Verity 快得多（更便宜！）。我们如何利用 Drake（初级工程师）来更快地编写代码，同时确保我们仍然做出正确的技术决策？
首先，Drake 通过劳动密集型的过程编写代码，并在过程中做出技术决策。然后，我们将代码交给 Verity 进行审查。
在审查代码后，Verity 可能会决定 Drake 做出的前三个技术决策是正确的，但最后两个需要重做。因此，Drake 回到原点，放弃了他最后两个决策，并从那里重新开始编码。
值得注意的是，尽管资深工程师 Verity 只审查了一次代码，但我们已经生成了 3 段经过验证的代码，与她亲自编写的代码完全相同！因此，假设 Verity 能够比她亲自编写这 3 段代码更快地审查代码，这种方法是更有利的。
在 Transformer 推理的背景下，Verity 将扮演我们想要用于任务的输出的大模型角色，称为验证器模型。同样，Drake 将扮演一个较小的模型，该模型能够比大模型更快地生成文本，称为草案模型。因此，我们将使用草案模型生成 8 个标记，然后并行处理这 8 个标记，丢弃那些不匹配的标记。
如上所述，推测解码的一个关键属性是它不会改变输出的质量。只要使用草案模型生成标记和验证标记所需的时间少于生成这些标记所需的时间，我们就处于领先地位。
在原生 PyTorch 中完成所有这些事情的一个好处是，这种技术实际上非常容易实现！以下是整个实现的全部内容，大约 50 行原生 PyTorch 代码。
虽然推测性解码保证了与常规生成相比，我们在数学上具有相同的结果，但它确实有一个属性，即运行时性能会根据生成的文本以及草案和验证器模型的匹配程度而变化。例如，当运行 CodeLlama-34B + CodeLlama-7B 时，我们能够获得每秒 2 倍的 token 提升用于生成代码。另一方面，当使用 Llama-7B + TinyLlama-1B 时，我们只能获得大约 1.3 倍的 token/s 提升。
旁注：在 AMD 上运行此操作
如上所述，解码过程中的每个内核都是通过 torch.compile 从头生成的，并转换为 OpenAI Triton。由于 AMD 有 torch.compile 后端（以及 Triton 后端），我们可以在 AMD GPU 上简单地执行上述所有优化……但使用 int8 量化，我们能够达到每秒 102.5 个 token，即 MI250x 的一半 GCD（即一半）！
第 4 步：使用 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 内核。
这些技术需要一些额外的工作，但将它们全部结合起来，结果会带来更好的性能！
第 5 步：将所有内容组合在一起（244.7 tok/s）
最后，我们可以将所有技术组合起来，以实现更好的性能！
步骤 6：使用张量并行
到目前为止，我们一直在限制自己，以最小化单个 GPU 上的延迟。然而，在许多情况下，我们能够访问多个 GPU。这使得我们可以进一步提高延迟！
为了直观地理解为什么这会使我们提高延迟，让我们看看 MBU 的先前方程式，特别是分母部分。在多个 GPU 上运行使我们能够访问更多的内存带宽，从而具有更高的潜在性能。
关于选择哪种并行策略，请注意，为了降低我们处理一个示例的延迟，我们需要能够同时利用我们的内存带宽跨越更多设备。这意味着我们需要将一个标记的处理分配给多个设备。换句话说，我们需要使用张量并行。
幸运的是，PyTorch 也为张量并行提供了与 torch.compile 兼容的低级工具。我们还在开发用于表达张量并行的更高级 API，敬请期待！
即使没有高级 API，实际上添加张量并行仍然相当简单。我们的实现只有 150 行代码，并且不需要对模型进行任何更改。
我们仍然能够利用之前提到的所有优化，这些优化都可以与张量并行一起使用。结合这些，我们能够以 55 个标记/秒的速度使用 int8 量化来服务 Llama-70B！
结论
让我们看看我们能完成什么。
- 简单性：忽略量化，model.py（244 行代码）+ generate.py（371 行代码）+ tp.py（151 行代码）共计 766 行代码实现了快速推理、推测解码和张量并行。
- 性能：使用 Llama-7B，我们能够通过编译+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 注意力、int8 量化、LoRA 微调等工作。
- GGML 助力快速推进，设备端LLMs推理
- 安德烈·卡帕西领导简单、可解释且快速的LLM实现
- MLC-LLM推动在异构硬件上实现 4 位量化性能
