PyTorch 团队亲自教你怎么加快大模型推理。
在曩昔的一年里，生成式 AI 发展迅猛，在这当中，文本生成一直是一个特别受欢迎的范畴，很多开源项目如 llama.cpp、vLLM 、 MLC-LLM 等，为了取得更好的作用，都在进行不断的优化。
作为机器学习社区中最受欢迎结构之一的 PyTorch，天然也是抓住了这一新的机遇，不断优化。为此让大家更好的了解这些立异，PyTorch 团队专门设置了系列博客，重点介绍怎么运用纯原生 PyTorch 加快生成式 AI 模型。
在第一篇博客中，PyTorch 团队展示了仅运用纯原生 PyTorch 重写 Segment Anything（SAM）模型，比原始完成快 8 倍。在本博客中，他们又为我们带来了新的内容，即怎么加快 LLM 推理。
我们先来看看结果，该团队重写 LLM，推理速度比基线足足快了 10 倍，并且没有丢失准确率，只用了不到 1000 行的纯原生 PyTorch 代码！
一切基准测试都在 A100-80GB 上运转的，功率约束在 330W。
这些优化包含：
-
Torch.compile：PyTorch 模型编译器， PyTorch 2.0 加入了一个新的函数，叫做 torch.compile ()，可以经过一行代码对已有的模型进行加快；
-
GPU 量化：经过下降运算精度来加快模型；
-
Speculative Decoding：一种大模型推理加快方法，运用一个小的「draft」模型来猜测大的「目标」模型的输出；
-
张量并行：经过在多个设备上运转模型来加快模型推理。
接下来，我们看看每一步都是怎么完成的。
6 步加快大模型推理
该研究表明，在没有优化之前，大模型的推理功能为 25.5 tok/s，作用不是很好：
经过一番探究后总算找到了原因：CPU 开支过大。然后就有了下面的 6 步优化进程。
第一步：经过 Torch.compile 和静态 KV 缓存削减 CPU 开支，完成 107.0 TOK/S
torch.compile 答应用户将更大的区域捕获到单个编译区域中，特别是在 mode=”reduce-overhead” 时（参阅下面的代码），这一功能关于削减 CPU 开支非常有效，除此以外，本文还指定 fullgraph=True，用来验证模型中没有「图形中断」（即 torch.compile 无法编译的部分）。
然而，即使有 torch.compile 的加持，仍是会遇到一些妨碍。
第一个妨碍是 kv 缓存。即当用户生成更多的 token 时， kv 缓存的「逻辑长度（logical length）」会增加。出现这种问题有两个原因：一是每次缓存增加时重新分配（和仿制）kv 缓存的本钱非常高；其次，这种动态分配使得削减开支变得更加困难。
为了处理这个问题，本文运用静态 KV 缓存，静态分配 KV 缓存的大小，然后屏蔽掉注意力机制中未运用的值。
第二个妨碍是 prefill 阶段。用 Transformer 进行文本生成可视为一个两阶段进程：1. 用来处理整个提示的 prefill 阶段 2. 解码 token.
尽管 kv 缓存被设置为静态化，但由于提示长度可变 ，prefill 阶段依然需求更多的动态性。因而，需求运用单独的编译战略来编译这两个阶段。
尽管这些细节有点棘手，但完成起来并不困难，而且功能的提升是巨大的。这一通操作下来，功能提高了 4 倍多，从 25 tok/s 提高到 107 tok/s。
第二步：经过 int8 权重量化缓解内存带宽瓶颈，完成 157.4 tok /s
经过上文，我们现已看到应用 torch.compile 、静态 kv 缓存等带来的巨大加快，但 PyTorch 团队并不满足于此，他们又找了其他角度进行优化。
他们以为加快生成式 AI 训练的最大瓶颈是将权重从 GPU 全局内存加载到寄存器的代价。换句话说，每次前向传播都需求「触摸（touch）」GPU 上的每个参数。那么，理论上我们可以以多快的速度「触摸」模型中的每个参数？
为了衡量这一点，本文运用模型带宽运用率（MBU），计算它非常简略，如下所示：
举例来说，关于一个 7B 参数模型，每个参数都存储在 fp16 中（每个参数 2 字节），可以完成 107 tokens/s。A100-80GB 理论上有 2 TB/s 的内存带宽。
如下图所示，将上述公式带入详细的数值，可以得到 MBU 为 72%！这个结果是适当不错的，因为很多研究很难打破 85%。
但 PyTorch 团队还想将这个数值在提高一些。他们发现无法改变模型中参数的数量，也无法改变 GPU 的内存带宽。但他们发现可以更改每个参数存储的字节数！
因而，他们打算用 int8 量化。
请注意，这仅是量化权重，计算自身依然在 bf16 中完成。此外，有了 torch.compile，可以轻松生成 int8 量化的高效代码。
就像上图所展示的，从深蓝色线（torch.compile + int8）可以看出，运用 torch.compile + int8 仅权重量化时，功能有明显提升。
将 int8 量化应用于 Llama-7B 模型，功能提高了约 50%，达到 157.4 tokens/s。
**第三步：运用 Speculative Decoding
**
即使在运用了 int8 量化等技能之后，该团队依然面临着另一个问题，即为了生成 100 个 token，有必要加载权重 100 次。
即使权重被量化，一遍又一遍地加载权重也避免不了，这种问题该怎么处理呢？事实证明，运用 speculative decoding 可以打破这种严厉的串行依靠性并取得加快。
该研究运用草稿（draft）模型生成 8 个 token，然后运用验证器模型并行处理，丢弃不匹配的 token。这一进程打破了串行依靠。整个完成进程大约 50 行原生 PyTorch 代码。
第四步：运用 int4 量化和 GPTQ 方法进一步减小权重，完成 202.1 tok/s
本文发现，当权重为 4-bits 时，模型的准确率开端下降。
为了处理这个问题，本文运用两个技巧来处理：第一个是拥有更细粒度的缩放因子；另一种是运用更先进的量化战略。将这些操作组合在一起，得到如下：
第五步：将一切内容组合在一起，得到 244.7 tok/s
最终，将一切技能组合在一起以取得更好的功能，得到 244.7 tok/s。
第六步：张量并行性
到目前为止，本文一直是在单个 GPU 上最大极限地削减延迟。其实，运用多个 GPU 也是可以的，这样一来，延迟现象会得到进一步改善。
非常庆幸的是，PyTorch 团队供给了张量并行的低级东西，只需 150 行代码，并且不需求任何模型更改。
前面说到的一切优化都可以持续与张量并行性组合，将这些组合在一起，能以 55 tokens/s 的速度为 Llama-70B 模型供给 int8 量化。
最终，简略总结一下文章主要内容。在 Llama-7B 上，本文运用「compile + int4 quant + speculative decoding」这一套组合拳，完成 240+ tok/s。在 Llama-70B，本文还经过引入张量并行性以达到约 80 tok/s，这些都挨近或超越 SOTA 功能。
