使用 Optimum Intel 在 Intel Xeon 上加速 StarCoder

Hugging Face 與 Intel 已在第 4 代 Intel Xeon 可擴展處理器上為 StarCoder-15B 模型實現超過 7 倍的推論加速。此效能提升是透過結合 8 位元與 4 位元量化技術以及透過 optimum-intel 套件的輔助生成（推測解碼）來實現的。

推論效能基準

量化與輔助生成的結合顯著降低每個輸出標記的時間（TPOT），同時在 HumanEval 資料集上保持模型準確度。

StarCoder 量化 精度 HumanEval (pass@1) 首次標記時間 (ms) 首次標記加速 每標記時間 (ms) 每標記加速
Baseline None 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（1.64 億參數），其規模約為目標 StarCoder-15B 模型的 95 倍小。
  • 最佳配置： 透過對草稿模型與目標模型皆套用 8 位元 SmoothQuant，可獲得約 7.30 倍的每標記時間（TPOT）加速，為最佳結果。

輔助生成的技術取捨

在輔助生成中選擇 8 位元量化而非 4 位元，主要是因為瓶頸的轉移：

  1. 草稿模型： 8 位元量化的 1.64 億參數模型大部分可容納於 CPU 快取中，消除記憶體頻寬瓶頸，且避免了 4 位元 WOQ 所帶來的反量化開銷。
  2. 目標模型： 由於目標模型同步驗證多個標記，瓶頸從記憶體頻寬轉移至計算。8 位元量化模型因避免了計算密集的反量化過程，表現優於 4 位元模型。

透過 Optimum Intel 的實作

使用者可透過將標準的 AutoModelForCausalLM 類別替換為 optimum-intel 套件中的 IPEXModelForCausalLM，來實作這些最佳化。此套件利用第 4 代 Xeon 處理器上的 Intel Extension for PyTorch（IPEX）與 Intel Advanced Matrix Extensions（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