隨著 ChatGPT 的出現(xiàn)，大規(guī)模語言模型（LLM）研究及其應用得到學術(shù)界和工業(yè)界的廣泛關(guān)注。一方面，開源的 LLM 模型不斷涌現(xiàn)，比如 OPT、BLOOM、LLaMA 等，這些預訓練模型的推出極大地促進了 LLM 的相關(guān)研究，使得 LLM 可以被應用于解決愈發(fā)復雜的實際問題。利用這些開源模型，快速構(gòu)建一套基于 LLM 的應用服務已經(jīng)變得愈發(fā)容易，但 LLM 面臨著高昂的計算和存儲需求，其成本也令人望而卻步。

另一方面，以羊駝家族（如 Alpaca、Vicuna、Guanaco）為代表的，經(jīng)過微調(diào)或蒸餾的小型化 LLM 也成為了當下的研究焦點之一，在多項測評中都展現(xiàn)出了優(yōu)異的表現(xiàn)；此外，以 Quantization、LoRA、Offloading 為代表的多項系統(tǒng)優(yōu)化技術(shù)使得以更低的資源需求部署這些 LLM 成為可能。但天下沒有免費的午餐，有關(guān)證據(jù)表明 [1]，這些小型化的 LLM 以及面向低資源場景的系統(tǒng)優(yōu)化技術(shù)往往都會帶來模型質(zhì)量的下降，影響最終應用的效果。

因此，如何在保證模型輸出質(zhì)量的前提下，讓 LLM 推理變得高效和廉價，已經(jīng)成為了 MLSys 領(lǐng)域非常重要的研究問題。近日，來自卡耐基梅隆大學（CMU）的 Catalyst Group 團隊發(fā)布了一款「投機式推理」引擎 SpecInfer，可以借助輕量化的小模型來幫助大模型，在完全不影響生成內(nèi)容準確度的情況下，實現(xiàn)兩到三倍的推理加速。

• 論文鏈接：https://arxiv.org/abs/2305.09781
• 項目地址：https://github.com/flexflow/FlexFlow/tree/inference

論文作者之一、CMU 助理教授 Zhihao Jia 表示：「生成式大規(guī)模語言模型不僅推理效率低下而且部署成本很高；它們小型化的版本具有速度和價格上的優(yōu)勢，但是也會影響生成內(nèi)容的質(zhì)量；而 SpecInfer 可以實現(xiàn)這兩方面的雙贏?！?/span>

同樣來自 CMU Catalyst Group 的助理教授 Tianqi Chen 也表示：「SpecInfer 可以適用于云上的 LLM 部署等場景，讓 LLM 推理更加可擴展?！?/span>

研究現(xiàn)狀

目前 LLM 推理主要依賴于自回歸式（auto-regressive）的解碼（decoding）方式，每步解碼只能夠產(chǎn)生一個輸出 token，并且需要將歷史輸出內(nèi)容拼接后重新作為 LLM 的輸入，才能進行下一步的解碼。考慮到這種數(shù)據(jù)依賴，現(xiàn)有 LLM 推理系統(tǒng)如 FasterTransformer 會采用一種增量式解碼（incremental decoding）技術(shù)，將已經(jīng)解碼的 token 對應的 key/value 進行緩存，避免重新計算。但是，這類系統(tǒng)仍然面臨兩個關(guān)鍵的缺陷：1）由于逐 token 計算的解碼范式，算子并行度有限，GPU 硬件資源難以被充分利用；2）當序列過長時，KV-cache 空間消耗過大，有限的 GPU 顯存無法承載。因此，當面對超大規(guī)模的 LLM 推理時（如 GPT-4 32K tokens），現(xiàn)有系統(tǒng)往往面臨資源利用低效，推理延遲過高的問題。

Incremental Decoding 示意圖

為了解決上述問題，研究者提出了一種「投機式」推理引擎 SpecInfer，其核心思想是通過計算代價遠低于 LLM 的 “小模型” SSM（Small Speculative Model）替代 LLM 進行投機式地推理（Speculative Inference），每次會試探性地推理多步，將多個 SSM 的推理結(jié)果匯聚成一個 Speculated Token Tree，交由 LLM 進行驗證，通過高效的樹形解碼算子實現(xiàn)并行化推理，驗證通過的路徑將會作為模型的推理結(jié)果序列，進行輸出。

總體上來說，SpecInfer 利用了 SSM 的內(nèi)在知識幫助 LLM 以更低廉的計算成本完成了主要的推理過程，而 LLM 則在一定程度上破除了逐 token 解碼的計算依賴，通過并行計算確保最終輸出的結(jié)果完全符合原始的推理語義。

SpecInfer 工作流程

系統(tǒng)設(shè)計

SpecInfer 系統(tǒng)架構(gòu)

可學習推測器（Learning-based Speculator）

Speculator 的主要作用是利用 SSM 快速產(chǎn)生對 LLM 未來輸出的推測結(jié)果，SSM 可以是（微調(diào)后）小版本的 LLM（如 LLaMA 7B），也可以是量化或蒸餾的小規(guī)模 LLM，還可以是可供檢索的知識庫（如參考文本）亦或是用戶的自定義函數(shù)?？傊琒SM 的輸出結(jié)果越接近 LLM，驗證時才會更容易通過，整體的推理效率才會更高。

為此，SpecInfer 引入集成學習的思想，將多個 SSM 的結(jié)果融合，提高輸出的差異化程度。為了盡可能提高匹配率，Speculator 提出了 Collective Boost-Tuning 方法，即在一個公開的通用數(shù)據(jù)集（如 OpenWebText）上，從一個較弱的 SSM 開始進行微調(diào)，將匹配程度較低的序列不斷從數(shù)據(jù)中過濾，交由新的 SSM 來學習，持續(xù)多次，提高整體的推測質(zhì)量；此外，Speculator 還引入了一個可學習的調(diào)度器（scheduler）來決定選用哪些 SSM 以獲得更長的匹配序列長度。

Token 樹驗證器（Token Tree Verifier）

SSM 的推理速度優(yōu)勢是 SpecInfer 能夠加速推理的前提，但另一個不可或缺的因素就是 LLM 對并行化推理的支持。在 SpecInfer 中，LLM 并不直接作為推理引擎產(chǎn)生輸出 token，但是它需要對 Speculator 中 SSM 產(chǎn)生的 token 進行驗證，確保輸出內(nèi)容符合 LLM 的推理語義。

在 SpecInfer 中，SSM 產(chǎn)生的輸出序列會被組織成 token tree 的樹形結(jié)構(gòu)，避免冗余的存儲開銷。為了能夠在 token tree 上進行并行化的驗證，SpecInfer 提出了一種樹形注意力（Tree Attention）計算方法，通過構(gòu)造的 mask 矩陣和基于深度優(yōu)先的 KV-cache 更新機制，Verifier 可以在不增加額外存儲的同時，盡可能并行化樹中每一條路徑的解碼過程。相比于樸素的逐序列或逐 Token 的解碼方式，樹形解碼可以同時在內(nèi)存開銷和計算效率上達到最優(yōu)。

Tree-based Decoding 示意圖

大規(guī)模 LLM 和小規(guī)模 SSM 協(xié)同工作

Speculative Inference 執(zhí)行 Timeline 對比

大規(guī)模的 LLM 在參數(shù)量上通?？梢赃_到小規(guī)模 SSM 的幾十倍甚至上百倍，而 SSM 相比于 LLM，在推理速度上，基于通常的系統(tǒng)實現(xiàn)，也有數(shù)倍到數(shù)十倍的性能優(yōu)勢，SpecInfer 結(jié)合了 SSM 極低的推理延遲以及 LLM 的并行驗證能力，大幅降低了較為耗時的 LLM 推理次數(shù)，最終可以在保證推理結(jié)果質(zhì)量的情況下顯著提升模型推理速度。

系統(tǒng)實現(xiàn)

SpecInfer 基于 FlexFlow 系統(tǒng)實現(xiàn)，支持用戶自定義模型結(jié)構(gòu)，導入模型參數(shù)，兼容主流深度學習框架的 operator 或 layer 抽象，現(xiàn)已支持常規(guī)的 GPT、LLaMA 等多種主流基礎(chǔ)模型。值得注意的是，F(xiàn)lexFlow 是一款面向分布式場景的深度學習系統(tǒng)，由來自 CMU、Stanford、MIT、NVIDIA 等機構(gòu)的研究人員共同維護，是機器學習系統(tǒng)領(lǐng)域最早提出 “自動并行” 的工作之一 (MLSys’19, ICML’18) [2,3]，也是最早將計算圖優(yōu)化以及自動并行優(yōu)化集成進行聯(lián)合優(yōu)化的工作 (Unity, OSDI’22) [4]。

借助于 FlexFlow 的自動并行能力，SpecInfer 可以自動完成大規(guī)模 LLM 的最優(yōu)分布式部署。與此同時，SpecInfer 還可以支持 Offloading 操作，以較低的成本擴展模型的規(guī)模。SpecInfer 通過獨特的「投機式推理」機制，可以大幅降低 LLM 所需的推理步數(shù)，從而減小分布式場景的網(wǎng)絡(luò)通信開銷，緩解 Offloading 場景下的 PCIe 傳輸帶寬瓶頸。

實驗結(jié)果

端到端推理延遲

端到端實驗：使用 LLaMA-7B 作為 LLM，LLaMA-160M 作為 SSM，在五個對話數(shù)據(jù)集上進行了測試，相比于依賴于增量式解碼的 LLM，SpecInfer 可以使推理延遲降低 1.9-2.8 倍。

單次推理平均步長（LLM：OPT-13B + SSMs：OPT-125M）

單次推理平均步長（LLM：LLaMA-7B + SSMs：LLaMA-160M）

匹配長度測試：分別使用 OPT 和 LLaMA 系列模型，測試 SpecInfer 中 LLM 的平均驗證通過序列長度，可以看出，隨著 SSM 數(shù)量的提升，在各個對話數(shù)據(jù)集上，LLM 的驗證通過長度均會得到提升，以 5 個 SSM 為例，OPT 和 LLaMA 在 5 個數(shù)據(jù)集上平均可達 3.68 和 2.67，相比于僅使用單一 SSM，分別提升 26.4% 和 24.8%。

更多更詳細的實驗結(jié)果可以參考論文原文：https://arxiv.org/abs/2305.09781

總結(jié)

SpecInfer 是首個基于「推測式解碼」的分布式 LLM 推理引擎，通過集成多個小模型，以及基于 token tree 的原創(chuàng)系統(tǒng)實現(xiàn)優(yōu)化，可以幫助現(xiàn)有的主流 LLM 減少內(nèi)存訪問需求，實現(xiàn)兩到三倍的無損推理加速，大幅降低推理成本。

作者介紹

SpecInfer 項目的指導老師是 Zhihao Jia，他目前在卡耐基梅隆大學計算機學院擔任助理教授。他的研究興趣主要包括面向機器學習、量子計算以及大規(guī)模數(shù)據(jù)分析的系統(tǒng)研究。此前他曾畢業(yè)于清華大學的姚班，博士畢業(yè)于 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 共同主持，致力于集成來自于機器學習算法、系統(tǒng)、硬件等多方面的優(yōu)化技術(shù)，構(gòu)造自動化的機器學習系統(tǒng)。此前，該實驗室還推出了 MLC-LLM [5] 等開源項目，推進 LLM 大模型相關(guān)系統(tǒng)的研究和應用。實驗室主頁：https://catalyst.cs.cmu.edu。

論文的共同一作分別是 Xupeng Miao（博士后研究員），Gabriele Oliaro（博一）以及 Zhihao Zhang（博一），均來自于 CMU Catalyst Group 團隊。其中，Xupeng Miao 博士畢業(yè)于北京大學，主要研究方向包括機器學習系統(tǒng)、數(shù)據(jù)管理和分布式計算，曾獲 VLDB2022 最佳可擴展數(shù)據(jù)科學論文獎、2022 年 ACM 中國優(yōu)博獎、2022 年世界人工智能大會（WAIC）云帆獎等榮譽，個人主頁：https://hsword.github.io。