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

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

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

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

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

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

研究現(xiàn)狀

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

Incremental Decoding 示意圖

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

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

SpecInfer 工作流程

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

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

可學(xué)習(xí)推測(cè)器（Learning-based Speculator）

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

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

Token 樹驗(yàn)證器（Token Tree Verifier）

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

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

Tree-based Decoding 示意圖

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

Speculative Inference 執(zhí)行 Timeline 對(duì)比

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

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

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

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

實(shí)驗(yàn)結(jié)果

端到端推理延遲

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

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

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

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

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

總結(jié)

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

作者介紹

SpecInfer 項(xiàng)目的指導(dǎo)老師是 Zhihao Jia，他目前在卡耐基梅隆大學(xué)計(jì)算機(jī)學(xué)院擔(dān)任助理教授。他的研究興趣主要包括面向機(jī)器學(xué)習(xí)、量子計(jì)算以及大規(guī)模數(shù)據(jù)分析的系統(tǒng)研究。此前他曾畢業(yè)于清華大學(xué)的姚班，博士畢業(yè)于 Stanford 大學(xué)，師從 Alex Aiken 和 Matei Zaharia，曾獲 Stanford Arthur Samuel Best Doctoral Thesis Award，NSF CAREER Asward 以及來自 Amazon, Google, Meta, Oracle, 以及 Qualcomm 的多項(xiàng)研究獎(jiǎng)項(xiàng)，個(gè)人主頁：https://www.cs.cmu.edu/~zhihaoj2/。

孵化 SpecInfer 項(xiàng)目的主要是 CMU 的 Catalyst Group 實(shí)驗(yàn)室，該實(shí)驗(yàn)室由 Zhihao Jia 與 Tianqi Chen（陳天奇）在 CMU 共同主持，致力于集成來自于機(jī)器學(xué)習(xí)算法、系統(tǒng)、硬件等多方面的優(yōu)化技術(shù)，構(gòu)造自動(dòng)化的機(jī)器學(xué)習(xí)系統(tǒng)。此前，該實(shí)驗(yàn)室還推出了 MLC-LLM [5] 等開源項(xiàng)目，推進(jìn) LLM 大模型相關(guān)系統(tǒng)的研究和應(yīng)用。實(shí)驗(yàn)室主頁：https://catalyst.cs.cmu.edu。

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