機器之心專欄
機器之心編輯部
近日，來自卡耐基梅隆大學（CMU）的 Catalyst Group 團隊發布了一款「投機式推理」引擎 SpecInfer，可以藉助輕量化的小模型來幫助大模型，在完全不影響生成內容準確度的情況下，實現兩到三倍的推理加速。
隨著 ChatGPT 的出現，大規模語言模型（LLM）研究及其應用得到學術界和工業界的廣泛關注。一方面，開源的 LLM 模型不斷湧現，比如 OPT、BLOOM、LLaMA 等，這些預訓練模型的推出極大地促進了 LLM 的相關研究，使得 LLM 可以被應用於解決愈發複雜的實際問題。利用這些開源模型，快速構建一套基於 LLM 的應用服務已經變得愈發容易，但 LLM 面臨著高昂的計算和存儲需求，其成本也令人望而卻步。
另一方面，以羊駝家族（如 Alpaca、Vicuna、Guanaco）為代表的，經過微調或蒸餾的小型化 LLM 也成為了當下的研究焦點之一，在多項測評中都展現出了優異的表現；此外，以 Quantization、LoRA、Offloading 為代表的多項系統優化技術使得以更低的資源需求部署這些 LLM 成為可能。但天下沒有免費的午餐，有關證據表明 [1]，這些小型化的 LLM 以及面向低資源場景的系統優化技術往往都會帶來模型質量的下降，影響最終應用的效果。
因此，如何在保證模型輸出質量的前提下，讓 LLM 推理變得高效和廉價，已經成為了 MLSys 領域非常重要的研究問題。近日，來自卡耐基梅隆大學（CMU）的 Catalyst Group 團隊發布了一款「投機式推理」引擎 SpecInfer，可以藉助輕量化的小模型來幫助大模型，在完全不影響生成內容準確度的情況下，實現兩到三倍的推理加速。
論文連結：https://arxiv.org/abs/2305.09781
項目地址：https://github.com/flexflow/FlexFlow/tree/inference
論文作者之一、CMU 助理教授 Zhihao Jia 表示：「生成式大規模語言模型不僅推理效率低下而且部署成本很高；它們小型化的版本具有速度和價格上的優勢，但是也會影響生成內容的質量；而 SpecInfer 可以實現這兩方面的雙贏。」
同樣來自 CMU Catalyst Group 的助理教授 Tianqi Chen 也表示：「SpecInfer 可以適用於雲上的 LLM 部署等場景，讓 LLM 推理更加可擴展。」
研究現狀
目前 LLM 推理主要依賴於自回歸式（auto-regressive）的解碼（decoding）方式，每步解碼只能夠產生一個輸出 token，並且需要將歷史輸出內容拼接後重新作為 LLM 的輸入，才能進行下一步的解碼。考慮到這種數據依賴，現有 LLM 推理系統如 FasterTransformer 會採用一種增量式解碼（incremental decoding）技術，將已經解碼的 token 對應的 key/value 進行緩存，避免重新計算。但是，這類系統仍然面臨兩個關鍵的缺陷：1）由於逐 token 計算的解碼範式，算子並行度有限，GPU 硬體資源難以被充分利用；2）當序列過長時，KV-cache 空間消耗過大，有限的 GPU 顯存無法承載。因此，當面對超大規模的 LLM 推理時（如 GPT-4 32K tokens），現有系統往往面臨資源利用低效，推理延遲過高的問題。
Incremental Decoding 示意圖
為了解決上述問題，研究者提出了一種「投機式」推理引擎 SpecInfer，其核心思想是通過計算代價遠低於 LLM 的 「小模型」 SSM（Small Speculative Model）替代 LLM 進行投機式地推理（Speculative Inference），每次會試探性地推理多步，將多個 SSM 的推理結果匯聚成一個 Speculated Token Tree，交由 LLM 進行驗證，通過高效的樹形解碼算子實現並行化推理，驗證通過的路徑將會作為模型的推理結果序列，進行輸出。
總體上來說，SpecInfer 利用了 SSM 的內在知識幫助 LLM 以更低廉的計算成本完成了主要的推理過程，而 LLM 則在一定程度上破除了逐 token 解碼的計算依賴，通過並行計算確保最終輸出的結果完全符合原始的推理語義。
SpecInfer 工作流程
系統設計
SpecInfer 系統架構
可學習推測器（Learning-based Speculator）
Speculator 的主要作用是利用 SSM 快速產生對 LLM 未來輸出的推測結果，SSM 可以是（微調後）小版本的 LLM（如 LLaMA 7B），也可以是量化或蒸餾的小規模 LLM，還可以是可供檢索的知識庫（如參考文本）亦或是用戶的自定義函數。總之，SSM 的輸出結果越接近 LLM，驗證時才會更容易通過，整體的推理效率才會更高。
為此，SpecInfer 引入集成學習的思想，將多個 SSM 的結果融合，提高輸出的差異化程度。為了儘可能提高匹配率，Speculator 提出了 Collective Boost-Tuning 方法，即在一個公開的通用數據集（如 OpenWebText）上，從一個較弱的 SSM 開始進行微調，將匹配程度較低的序列不斷從數據中過濾，交由新的 SSM 來學習，持續多次，提高整體的推測質量；此外，Speculator 還引入了一個可學習的調度器（scheduler）來決定選用哪些 SSM 以獲得更長的匹配序列長度。
Token 樹驗證器（Token Tree Verifier）
SSM 的推理速度優勢是 SpecInfer 能夠加速推理的前提，但另一個不可或缺的因素就是 LLM 對並行化推理的支持。在 SpecInfer 中，LLM 並不直接作為推理引擎產生輸出 token，但是它需要對 Speculator 中 SSM 產生的 token 進行驗證，確保輸出內容符合 LLM 的推理語義。
在 SpecInfer 中，SSM 產生的輸出序列會被組織成 token tree 的樹形結構，避免冗餘的存儲開銷。為了能夠在 token tree 上進行並行化的驗證，SpecInfer 提出了一種樹形注意力（Tree Attention）計算方法，通過構造的 mask 矩陣和基於深度優先的 KV-cache 更新機制，Verifier 可以在不增加額外存儲的同時，儘可能並行化樹中每一條路徑的解碼過程。相比於樸素的逐序列或逐 Token 的解碼方式，樹形解碼可以同時在內存開銷和計算效率上達到最優。
Tree-based Decoding 示意圖
大規模 LLM 和小規模 SSM 協同工作
Speculative Inference 執行 Timeline 對比
大規模的 LLM 在參數量上通常可以達到小規模 SSM 的幾十倍甚至上百倍，而 SSM 相比於 LLM，在推理速度上，基於通常的系統實現，也有數倍到數十倍的性能優勢，SpecInfer 結合了 SSM 極低的推理延遲以及 LLM 的並行驗證能力，大幅降低了較為耗時的 LLM 推理次數，最終可以在保證推理結果質量的情況下顯著提升模型推理速度。
系統實現
SpecInfer 基於 FlexFlow 系統實現，支持用戶自定義模型結構，導入模型參數，兼容主流深度學習框架的 operator 或 layer 抽象，現已支持常規的 GPT、LLaMA 等多種主流基礎模型。值得注意的是，FlexFlow 是一款面向分布式場景的深度學習系統，由來自 CMU、Stanford、MIT、NVIDIA 等機構的研究人員共同維護，是機器學習系統領域最早提出 「自動並行」 的工作之一 (MLSys』19, ICML』18) [2,3]，也是最早將計算圖優化以及自動並行優化集成進行聯合優化的工作 (Unity, OSDI』22) [4]。
藉助於 FlexFlow 的自動並行能力，SpecInfer 可以自動完成大規模 LLM 的最優分布式部署。與此同時，SpecInfer 還可以支持 Offloading 操作，以較低的成本擴展模型的規模。SpecInfer 通過獨特的「投機式推理」機制，可以大幅降低 LLM 所需的推理步數，從而減小分布式場景的網絡通信開銷，緩解 Offloading 場景下的 PCIe 傳輸帶寬瓶頸。
實驗結果
端到端推理延遲
端到端實驗：使用 LLaMA-7B 作為 LLM，LLaMA-160M 作為 SSM，在五個對話數據集上進行了測試，相比於依賴於增量式解碼的 LLM，SpecInfer 可以使推理延遲降低 1.9-2.8 倍。
單次推理平均步長（LLM：OPT-13B + SSMs：OPT-125M）
單次推理平均步長（LLM：LLaMA-7B + SSMs：LLaMA-160M）
匹配長度測試：分別使用 OPT 和 LLaMA 系列模型，測試 SpecInfer 中 LLM 的平均驗證通過序列長度，可以看出，隨著 SSM 數量的提升，在各個對話數據集上，LLM 的驗證通過長度均會得到提升，以 5 個 SSM 為例，OPT 和 LLaMA 在 5 個數據集上平均可達 3.68 和 2.67，相比於僅使用單一 SSM，分別提升 26.4% 和 24.8%。
更多更詳細的實驗結果可以參考論文原文：https://arxiv.org/abs/2305.09781
總結
SpecInfer 是首個基於「推測式解碼」的分布式 LLM 推理引擎，通過集成多個小模型，以及基於 token tree 的原創系統實現優化，可以幫助現有的主流 LLM 減少內存訪問需求，實現兩到三倍的無損推理加速，大幅降低推理成本。
作者介紹
SpecInfer 項目的指導老師是 Zhihao Jia，他目前在卡耐基梅隆大學計算機學院擔任助理教授。他的研究興趣主要包括面向機器學習、量子計算以及大規模數據分析的系統研究。此前他曾畢業於清華大學的姚班，博士畢業於 Stanford 大學，師從 Alex Aiken 和 Matei Zaharia，曾獲 Stanford Arthur Samuel Best Doctoral Thesis Award，NSF CAREER Asward 以及來自 Amazon, Google, Meta, Oracle, 以及 Qualcomm 的多項研究獎項，個人主頁：https://www.cs.cmu.edu/~zhihaoj2/。
孵化 SpecInfer 項目的主要是 CMU 的 Catalyst Group 實驗室，該實驗室由 Zhihao Jia 與 Tianqi Chen（陳天奇）在 CMU 共同主持，致力於集成來自於機器學習算法、系統、硬體等多方面的優化技術，構造自動化的機器學習系統。此前，該實驗室還推出了 MLC-LLM [5] 等開源項目，推進 LLM 大模型相關系統的研究和應用。實驗室主頁：https://catalyst.cs.cmu.edu。
論文的共同一作分別是 Xupeng Miao（博士後研究員），Gabriele Oliaro（博一）以及 Zhihao Zhang（博一），均來自於 CMU Catalyst Group 團隊。其中，Xupeng Miao 博士畢業於北京大學，主要研究方向包括機器學習系統、數據管理和分布式計算，曾獲 VLDB2022 最佳可擴展數據科學論文獎、2022 年 ACM 中國優博獎、2022 年世界人工智慧大會（WAIC）雲帆獎等榮譽，個人主頁：https://hsword.github.io。
參考文獻：
[1] Gudibande, A., Wallace, E., Snell, C., Geng, X., Liu, H., Abbeel, P., Levine, S., & Song, D. (2023). The False Promise of Imitating Proprietary LLMs.
[2] Jia, Z., Lin, S., Qi, C. R., & Aiken, A. (2018, July). Exploring Hidden Dimensions in Parallelizing Convolutional Neural Networks. In ICML (pp. 2279-2288).
[3] Jia, Z., Zaharia, M., & Aiken, A. (2019). Beyond Data and Model Parallelism for Deep Neural Networks. Proceedings of Machine Learning and Systems, 1, 1-13.
[4] Unger, C., Jia, Z., Wu, W., Lin, S., Baines, M., Narvaez, C. E. Q., ... & Aiken, A. (2022). Unity: Accelerating Training Through Joint Optimization of Algebraic Transformations and Parallelization. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22) (pp. 267-284).
[5] https://github.com/mlc-ai/mlc-llm