使用 🤗 Optimum Intel 在 Xeon 上加速 StarCoder：Q8/Q4 和推测解码

引言

最近，代码生成模型变得非常流行，特别是随着 BigCode 的 StarCoder 和 Meta AI 的 Code Llama 等最先进的开源模型的发布。越来越多的工作致力于优化大型语言模型 (LLM) 并使其更易于访问。在此博客中，我们很高兴分享 LLM 在 Intel Xeon 上的最新优化结果，重点关注流行的代码生成 LLM StarCoder。

StarCoder 模型是一款尖端 LLM，专为帮助用户完成各种编码任务而设计，例如代码补全、错误修复、代码摘要，甚至从自然语言描述生成代码片段。StarCoder 模型是 StarCoder 家族的成员，其中还包括 StarCoderBase 变体。这些用于代码的大型语言模型 (Code LLM) 基于 GitHub 上许可的数据进行训练，包括 80 多种编程语言、Git 提交、GitHub 问题和 Jupyter 笔记本。在这项工作中，我们展示了通过将 8 位和 4 位量化与 辅助生成相结合，在 Intel 第 4 代 Xeon 上，StarCoder-15B 模型的推理加速了 7 倍以上。

在 Hugging Face Spaces 上试用我们的 演示，该演示正在第 4 代 Intel Xeon 可扩展处理器上运行。

步骤 1：基线和评估

我们使用 StarCoder (15B) 与 PyTorch 和 Intel Extension for PyTorch (IPEX) 建立基线。有几个数据集旨在评估自动代码补全的质量。在这项工作中，我们使用流行的 HumanEval 数据集来评估模型的质量和性能。HumanEval 包含 164 个编程问题，以函数签名和文档字符串的形式呈现，模型补全函数的代码。提示的平均长度为 139。我们使用 Bigcode Evaluation Harness 测量质量并报告 pass@1 指标。我们通过测量 HumanEval 测试集上的首个 Token 时间 (TTFT) 和每个输出 Token 时间 (TPOT) 来测量模型性能，并报告平均 TTFT 和 TPOT。第 4 代 Intel Xeon 处理器具有 AI 注入加速功能，称为 Intel® Advanced Matrix Extensions (Intel® AMX)。具体来说，每个核心都内置了 BFloat16 (BF16) 和 Int8 GEMM 加速器，以加速深度学习训练和推理工作负载。AMX 加速推理通过 PyTorch 2.0 和 Intel Extension for PyTorch (IPEX) 引入，此外还有针对 LLM 推理中使用的各种常见操作（例如层归一化、SoftMax、缩放点积）的其他优化。作为起点，我们使用 PyTorch 和 IPEX 中的开箱即用优化来使用 BF16 模型执行推理。图 1 显示了基线模型的延迟，表 1 和表 2 显示了其延迟和精度。

图 1. 基线模型的延迟。

LLM 量化

LLM 中的文本生成以自回归方式执行，因此需要将整个模型从内存加载到 CPU 以生成每个新 Token。我们发现，片外内存 (DRAM) 和 CPU 之间的带宽是 Token 生成过程中最大的瓶颈。量化是缓解此问题的常用方法。它减小了模型大小，从而缩短了模型权重加载时间。

在这项工作中，我们重点关注两种量化类型

1. 仅权重 量化 (WOQ) - 模型的权重被量化，但激活未量化，而计算以更高精度（例如 BF16）执行，这需要反量化。
2. 静态量化 (SQ) - 权重和激活都经过量化。此量化过程包括通过校准步骤预先计算量化参数，这使得计算能够以较低精度（例如 INT8）执行。图 2 显示了 INT8 静态量化计算过程。

步骤 2：8 位量化 (INT8)

SmoothQuant 是一种后训练量化算法，用于对 LLM 进行 INT8 量化，且精度损失最小。由于激活的特定通道中存在大值异常点，静态量化方法在 LLM 上的表现不佳。由于激活是按 Token 量化的，静态量化会导致截断异常值或低幅度激活下溢。SmoothQuant 算法通过引入预量化阶段解决了这个问题，在该阶段，将额外的平滑缩放因子应用于激活和权重，从而平滑激活中的异常值并确保更好地利用量化级别。

图 2. INT8 静态量化的计算图。

使用 IPEX，我们将 SmoothQuant 应用于 StarCoder 模型。我们使用 MBPP 数据集的测试拆分作为校准数据集，并引入了 Q8-StarCoder。我们的评估表明，Q8-StarCoder 相对于基线没有精度损失（事实上，甚至略有改进）。在性能方面，Q8-StarCoder 在 TTFT 方面实现了 ~2.19 倍的加速，在 TPOT 方面实现了 ~2.20 倍的加速。图 3 显示了 Q8-StarCoder 与 BF16 基线模型相比的延迟 (TPOT)。

图 3. 8 位量化模型的延迟加速。

步骤 3：4 位量化 (INT4)

尽管 INT8 将模型大小比 BF16 减小了 2 倍（每个权重 8 位，而 BF16 为 16 位），但内存带宽仍然是最大的瓶颈。为了进一步缩短模型从内存加载的时间，我们使用 WOQ 将模型的权重量化为 4 位。请注意，4 位 WOQ 在计算之前需要反量化为 16 位（图 4），这意味着存在计算开销。

图 4. 量化为 INT4 的模型的计算图。

张量级非对称最近舍入 (RTN) 量化是一种基本的 WOQ 技术，它带来了挑战，并且通常会导致精度降低，然而 文献 (Zhewei Yao, 2022) 表明，对模型权重进行分组量化有助于保持精度。为了避免精度下降，我们以组（例如 128）为单位，沿着输入通道对连续值执行 4 位量化，并根据每组计算缩放因子。我们发现分组 4 位 RTN 足以在 HumanEval 数据集上保持 StarCoder 的精度。与 BF16 基线相比，4 位模型在 TPOT 方面实现了 3.35 倍的加速（图 5），但由于计算前将 4 位反量化为 16 位的开销，它在 TTFT 方面出现了预期的 0.84 倍减速（表 1）。

图 5. 4 位量化模型的延迟加速。

生成第一个 Token 和后续 Token 之间的不同瓶颈

生成第一个 Token 的初始步骤涉及对整个输入提示进行并行处理，当提示长度较高时，这需要大量的计算资源。因此，计算成为此阶段的瓶颈。因此，将精度从 BF16 切换到 INT8 可改善此过程的性能，与基线（以及涉及计算开销形式的反量化的 4 位 WOQ）相比。然而，从第二步开始，当系统以自回归方式逐个生成其余 Token 时，模型会一次又一次地从内存加载以生成每个新 Token。结果，瓶颈变为内存带宽，而不是执行的计算量 (FLOPS)，因此 INT4 的性能优于 INT8 和 BF16。

步骤 4：辅助生成 (AG)

另一种缓解高推理延迟和内存带宽瓶颈问题的方法是 辅助生成 (AG)，它是 推测解码 的实际实现。AG 通过更好地平衡内存和计算操作来缓解此问题。它依赖于这样一种前提：较小且较快的辅助草稿模型通常会生成与较大目标模型相同的 Token。

AG 使用一个小型快速草稿模型贪婪地生成 K 个候选 Token。这些输出 Token 生成速度快得多，但其中一些可能与原始目标模型的输出 Token 不相似。因此，在下一步中，目标模型并行地在一次前向传递中检查所有 K 个候选 Token 的有效性。此过程加快了解码速度，因为并行解码 K 个 Token 的延迟小于自回归生成 K 个 Token 的延迟。

为了加速 StarCoder，我们使用 bigcode/tiny_starcoder_py 作为草稿模型。该模型与 StarCoder 共享相似的架构，但仅包含 1.64 亿个参数 - 比 StarCoder 小 ~95 倍，因此速度快得多。为了实现更大的加速，除了量化目标模型外，我们还对草稿模型进行量化。我们考虑了草稿模型和目标模型的 8 位 SmoothQuant 和 4 位 WOQ 量化选项。在评估草稿模型和目标模型的两种量化选项时，我们发现两种模型都采用 8 位 SmoothQuant 产生了最佳结果：TPOT 加速了 ~7.30 倍（图 6）。

这些量化选择得到了以下观察结果的支持

1. 草稿模型量化：当使用 8 位量化的 StarCoder（1.64 亿参数）作为草稿模型时，模型大部分适合 CPU 缓存。结果，内存带宽瓶颈得到缓解，因为 Token 生成无需为每个 Token 重复从片外内存读取目标模型。在这种情况下，没有内存瓶颈，我们看到 8 位量化的 StarCoder-164M 比 4 位 WOQ 量化的 StarCoder-164M 具有更好的加速。我们注意到，4 位 WOQ 在内存带宽是瓶颈的情况下具有优势，因为它内存占用较小，但是 4 位量化会带来计算开销，因为在计算之前需要执行 4 位到 16 位的反量化。
2. 目标模型量化：在辅助生成中，目标模型处理由草稿模型生成的一系列 K 个 Token。通过目标模型一次（并行）转发 K 个 Token，而不是应用“标准”顺序自回归处理，将平衡从内存带宽转移到计算瓶颈。因此，我们观察到，使用 8 位量化目标模型比使用 4 位模型产生更高的加速，因为将每个值从 4 位反量化到 16 位会产生额外的计算开销。

图 6. 优化模型的延迟加速。

表 1：StarCoder 模型在 Intel 第 4 代 Xeon 上的精度和延迟测量

要加载生成的模型并运行推理，只需将您的 AutoModelForXxx 类替换为 optimum-intel 中相应的 IPEXModelForXxx 类即可。

在开始之前，请确保已安装所有必要的库

pip install --upgrade-strategy eager optimum[ipex]

- from transformers import AutoModelForCausalLM
+ from optimum.intel import IPEXModelForCausalLM
  from transformers import AutoTokenizer, pipeline

- model = AutoModelForCausalLM.from_pretrained(model_id)
+ model = IPEXModelForCausalLM.from_pretrained(model_id)
  tokenizer = AutoTokenizer.from_pretrained(model_id)
  pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
  results = pipe("He's a dreadful magician and")