使用 Optimum Intel 在 Intel Xeon 上加速 StarCoder

Hugging Face 与 Intel 已在第 4 代 Intel Xeon 可扩展处理器上实现了 StarCoder-15B 模型超过 7 倍的推理加速。该性能提升通过在 optimum-intel 库中结合 8 位和 4 位量化技术与辅助生成（投机解码）实现。

推理性能基准

量化与辅助生成的结合显著降低了每输出令牌时间（TPOT），同时在 HumanEval 数据集上保持模型准确性。

StarCoder 量化 精度 HumanEval（pass@1） TTFT（ms） TTFT 加速比 TPOT（ms） TPOT 加速比
基线 无 A16W16 33.54 357.9 1.00x 181.0 1.00x
INT8 SmoothQuant A8W8 33.96 163.4 2.19x 82.4 2.20x
INT4 RTN (g128) A16W4 32.80 425.1 0.84x 54.0 3.35x
INT8 + AG SmoothQuant A8W8 33.96 183.6 1.95x 24.8 7.30x

针对 Intel Xeon 的量化策略

为了优化 StarCoder，团队使用两种不同的量化方法来解决内存带宽瓶颈——自回归令牌生成的主要限制：

8 位静态量化（INT8）

使用 SmoothQuant 算法，模型被量化为 INT8，精度损失极小。SmoothQuant 对激活和权重都应用平滑缩放因子，以减轻激活中大幅度异常值的影响。该方法在 TPOT 上实现约 2.20 倍加速，在首次令牌时间（TTFT）上实现约 2.19 倍加速。

4 位仅权重量化（INT4）

为了进一步降低内存占用和加载时间，模型权重使用分组最近取整（RTN）量化（每组 128）压缩至 4 位。虽然在 TPOT 上实现了 3.35 倍加速，但由于在计算前需要将 4 位权重反量化回 16 位，导致 TTFT 下降至 0.84 倍。

辅助生成（投机解码）

辅助生成（AG）通过使用一个小且快速的草稿模型预测候选令牌，然后并行地由更大的目标模型进行验证，从而加速推理。

  • 草稿模型： 团队使用了 bigcode/tiny_starcoder_py（164M 参数），其规模约为目标 StarCoder-15B 模型的 95 倍小。
  • 最佳配置： 通过对草稿模型和目标模型均使用 8 位 SmoothQuant，可获得约 7.30 倍 TPOT 加速的最佳结果。

辅助生成中的技术权衡

在辅助生成中选择 8 位量化而非 4 位，源于瓶颈的转移：

  1. 草稿模型： 8 位量化的 164M 参数模型大部分能够放入 CPU 缓存，消除了内存带宽瓶颈，并避免了 4 位仅权重量化（WOQ）带来的反量化开销。
  2. 目标模型： 由于目标模型并行验证多个令牌，瓶颈从内存带宽转向计算。8 位量化模型优于 4 位模型，因为它避免了计算密集的反量化过程。

通过 Optimum Intel 实现

用户可以通过将标准的 AutoModelForCausalLM 类替换为 optimum-intel 库中的 IPEXModelForCausalLM 来实现这些优化，该库利用了第 4 代 Xeon 处理器上的 Intel PyTorch 扩展（IPEX）和 Intel 高级矩阵扩展（AMX）。

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

from optimum.intel import IPEXModelForCausalLM
from transformers import AutoTokenizer, pipeline

model = IPEXModelForCausalLM.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

Sources