使用 🤗 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")