萬字綜述大模型高效推理:無問芯穹與清華、上交最新聯合研究全面解析大模型推理最佳化

机器之心發表於2024-06-14
近年來,大語言模型(Large Language Models, LLMs)受到學術界和工業界的廣泛關注,得益於其在各種語言生成任務上的出色表現,大語言模型推動了各種人工智慧應用(例如ChatGPT、Copilot等)的發展。然而,大語言模型的落地應用受到其較大的推理開銷的限制,對部署資源、使用者體驗、經濟成本都帶來了巨大挑戰。例如,將包含700億引數量的LLaMA-2-70B模型進行部署推理,至少需要6張RTX 3090Ti顯示卡或2張NVIDIA A100顯示卡,以部署在A100顯示卡上為例,該模型生成512長度的詞塊(token)序列需要耗時超過50秒。

許多研究工作致力於設計最佳化大語言模型推理開銷的技術,最佳化模型的推理延遲、吞吐、功耗和儲存等指標,成為許多研究的重要目標。為了對這些最佳化技術有更全面、更系統的認知,為大語言模型的部署實踐和未來研究提供建議和指南,來自清華大學電子工程系、無問芯穹和上海交通大學的研究團隊對大語言模型的高效推理技術進行了一次全面的調研和整理,在《A Survey on Efficient Inference for Large Language Models》(簡稱LLM Eff-Inference)這篇萬字長文綜述將領域相關工作劃分歸類為三個最佳化層次(即資料層、模型層和系統層),並逐個層次地介紹和總結相關技術工作。此外,該工作還對造成大語言模型推理不高效的根本原因進行分析,並基於對當前已有工作的綜述,深入探討高效推理領域未來應關注的場景、挑戰和路線,為研究者提供可行的未來研究方向。
圖片
圖片
圖注:《A Survey on Efficient Inference for Large Language Models》(LLM Eff-Inference)

一、大模型推理效率瓶頸分析

目前主流的大語言模型都是基於Transformer架構進行設計。通常來說,一個完整的模型架構由多個相同結構的Transformer塊組成,每個Transformer塊則包含多頭自注意力(Multi-Head Self-Attention, MHSA)模組、前饋神經網路(Feed Forward Network, FFN)和層歸一化(Layer Normalization,LN)操作。

大語言模型通常自迴歸(Auto-regressive)的方式生成輸出序列,即模型逐個詞塊生成,且生成每個詞塊時需要將前序的所有詞塊(包括輸入詞塊和前面已生成的詞塊)全部作為模型的輸入。因此,隨著輸出序列的增長,推理過程的開銷顯著增大。為了解決該問題,KV快取技術被提出,該技術透過儲存和複用前序詞塊在計算注意力機制時產生的Key和Value向量,減少大量計算上的冗餘,用一定的儲存開銷換取了顯著的加速效果。基於KV快取技術,通常可以將大語言模型的推理過程劃分為兩個階段(分別如下圖中(a)和(b)所示):

  • 預填充(Prefilling)階段:大語言模型計算並儲存輸入序列中詞塊的Key和Value向量,並生成第一個輸出詞塊。
  • 解碼(Decoding)階段:大語言模型利用KV快取技術逐個生成輸出詞塊,並在每步生成後儲存新詞塊的Key和Value向量。
    圖片
圖注:大語言模型推理的兩個階段,即預填充階段(a)和解碼階段(b)

大語言模型在實際部署應用中,我們通常關注其延時、吞吐、功耗和儲存,而在大語言模型推理過程中,有三個重要因素會直接影響上述效率指標,分別是計算開銷(Computational Cost)、訪存開(Memory Access Cost)和儲存開銷(Memory Cost)。進一步地,本綜述深入分析探究,並總結歸納除出影響上述指標和因素三點根本因素,分別為:
  • 模型規模:主流大語言模型龐大的模型規模會導致巨大的計算量、訪存量和儲存量;
  • 注意力運算元:作為大語言模型的核心運算元,注意力運算元具有與輸入長度呈平方關係增長的計算和儲存複雜度;
  • 解碼方式:主流的自迴歸解碼方式導致極低的計算-訪存比和硬體利用率,同時動態增長的KV快取會導致碎片化的記憶體使用,對訪存開銷和儲存開銷帶來增長。圖片
圖注:大語言模型推理效率瓶頸分析圖示

二、大模型高效推理技術領域縱覽

本綜述將大語言模型高效推理領域的技術劃分為三個層次,分別為:
  • 資料層最佳化技術:指透過最佳化輸入提示詞或規劃模型輸出內容最佳化推理效率。這類最佳化技術通常不需要修改模型本身,因此避免了大量的模型訓練或微調開銷;
  • 模型層最佳化技術:指透過設計高效的模型結構或模型壓縮技術最佳化推理效率。這類技術通常需要對模型進行預訓練或微調來恢復任務精度,同時通常對輸出結果是有損的;
  • 統層最佳化技術:指透過最佳化推理引擎或服務系統最佳化推理效率。這類技術通常不需要額外的模型訓練開銷,同時可以保證對輸出結果是無損的。
圖片
圖注:本綜述的分類體系

2.1 資料層最佳化技術

資料層最佳化技術可以劃分為兩大類:輸入壓縮(Input Compression)和輸出規劃(Output Organization)。

2.1.1 輸入壓縮技術

在實際利用大語言模型做回答時,通常會在輸入提示詞中加入一些輔助內容來增強模型的回答質量,例如,上下文學習技術(In-Context Learning,ICL)提出在輸入中加入多個相關的問答例子來教模型如何作答。然而,這些技術不可避免地會增長輸入詞提示的長度,導致模型推理的開銷增大。為了解決該問題,輸入壓縮技術透過直接減小輸入的長度來最佳化模型的推理效率

本綜述將該類技術進一步劃分為四個小類,分別為:

  • 提示詞剪枝(Prompt Pruning):通常根據設計好的重要度評估指標刪除輸入提示詞中不重要的詞塊、句子或文段,對被壓縮的輸入提示詞執行線上壓縮。
  • 提示詞總結(Prompt Summary):透過對輸入提示詞做文字總結任務,在保證其語義資訊相同地情況下縮短輸入的長度。該壓縮過程通常也是線上執行的。
  • 基於軟提示詞的壓縮(Soft Prompt-based Compression):透過微調訓練的方式得到一個長度較短的軟提示詞,代替原先的輸入提示詞(線上執行)或其中固定的一部分內容(離線執行)。其中,軟提示詞指連續的、可學習的詞塊序列,可以透過訓練的方式學習得到。
  • 檢索增強生成(Retrieval-Augmented Generation):透過檢索和輸入相關的輔助內容,並只將這些相關的內容加入到輸入提示詞中,來降低原本的輸入長度(相比於加入所有輔助內容)。

2.1.2 輸出規劃技術

傳統的生成解碼方式是完全序列的,輸出規劃技術透過規劃輸出內容,並行生成某些部分的的輸出來降低端到端的推理延時。以該領域最早的工作“思維骨架”(Skeleton-of-Thought,以下簡稱SoT)[45](無問芯穹於2023年7月釋出的工作,並被ICLR2024錄用)為例,SoT技術的核心思想是讓大語言模型自行規劃輸出的並行結構,並基於該結構進行並行解碼,提升硬體利用率,減少端到端生成延時。

具體來說,如下圖所示,SoT將大語言模型的生成分為兩個階段:在提綱階段,SoT透過設計的提示詞讓大語言模型輸出答案的大綱;在分點擴充套件階段,SoT讓大語言模型基於大綱中的每一個分點並行做擴充套件,最後將所有分點擴充套件的答案合併起來。SoT技術讓包含LLaMA-2、Vicuna模型在內的9種主流大語言模型的生成過程加速1.9倍以上,最高可達2.39倍。在SoT技術釋出後,一些研究工作透過微調大語言模型、前後端協同最佳化等方式最佳化輸出規劃技術,達到了更好的加速比和回答質量之間的權衡點。
圖片
圖注:輸出規劃技術SoT[45]示意

2.1.3 知識、建議和未來方向
隨著大語言模型需要處理的輸入提示詞和模型的輸出結果越來越長,資料層的最佳化技術越來越重要。在這類技術中,輸入壓縮技術主要最佳化預填充階段中注意力運算元帶來的平方量級的計算和儲存開銷,而輸出規劃技術則主要透過降低解碼階段中大量的訪存開銷,此外,一些研究嘗試利用大語言模型本身做輸出規劃,並取得了不錯的效果。
最近,各種各樣新的提示詞流程(Prompting Pipelines)和大語言模型智慧體(LLM Agents)出現,這些新的框架在增強大語言模型能力的同時,也引入了更長的輸入提示。利用輸入壓縮技術可以有效最佳化此類問題,同時,這些流程和智慧體框架在輸出內容中引入了更多可並行性,因此輸出規劃技術可能會在最佳化這類框架的效率上起到關鍵作用。除此之外,還有一些新工作關注如何直接設計更高效的智慧體框架。

2.2 模型層最佳化技術

模型層最佳化技術可以劃分為兩大類:高效結構設計(Efficient Structure Design)和模型壓縮(Model Compression)。前者通常需要將新設計的模型從頭進行預訓練,而後者通常只需要微調恢復精度即可。

2.2.1 高效結構設計

目前主流的大語言模型大多采用Transformer架構,從結構上看,引數量佔比最多的前饋神經網路(以下簡稱FFN)和平方複雜度的注意力運算元是導致模型推理效率低下的主要原因。基於此,本文將高效結構設計領域的技術進一步劃分為三類:

🔹高效FFN設計(Efficient FFN):該領域的主流方法為混合 專家(Mixture-of-Experts,MoE)技術,其核心是為不同的輸入詞塊分配不同數量的FFN(稱為專家),減少推理過程中被啟用的FFN權重數量。基於MoE的Transformer模型除了包含若干個專家FFN外,還包含一個負責分配專家的路由(Router)模型。該領域的研究工作主要關注三個方向:
  1. 更高效地獲取專家FFN的權重或構建更輕量化的專家FFN;
  2. 最佳化路由模型使其分配更加平衡,避免因分配不平衡導致的精度和效率下降;
  3. 最佳化MoE模型的訓練方式,使得訓練更加穩定。

🔹高效注意力運算元設計(Efficient Attention):該領域的研究工作可以分為:
  1. 多詢問(Multi-Query)注意力技術,即在不同的注意力頭之間共享部分Key和Value來降低訪存開銷和記憶體佔用;
  2. 低複雜度(Low-Complexity)注意力技術,主要包括基於核函式(Kernel-based)的注意力運算元和低秩注意力(Low-Rank)運算元。

🔹Transformer替代架構(Transformer Alternates):最新的研究工作逐漸關注於設計新的模型架構來取代Transformer架構,這些新模型架構大多具有線性或近似線性的複雜度,在處理文字長度較長時有顯著的效能優勢。本綜述將這類工作總結為兩大類:
  1. 狀態空間模型(State Space Models,SSMs)。這類模型的核心思想是將歷史資訊壓縮到一個隱狀態(state)中,透過狀態間的轉移建模新的資訊。聚焦於狀態空間模型的研究可以分為兩個方向,一方面研究狀態轉移矩陣初始化和引數化的方式,另一方面研究與其他模型結構(如Transformer、門控單元等)的融合架構。
  2. 非狀態空間模型。這類工作主要包含兩個分支,分別是採用長卷積做序列建模,以及採用基於注意力機制做改進的線性建模運算元。下表列舉了典型的Transformer替代架構的複雜度,可以總結出該類架構的研究趨勢:
    🔸訓練方面,新架構傾向於採用卷積或類注意力運算元的形式來保持訓練的並行性;
    🔸推理方面,新架構傾向於在預填充階段保證和輸入長度呈線性的計算複雜度,在解碼階段保證與文字長度無關的計算複雜度。
圖片
圖注:典型Transformer替代架構的訓練和推理複雜度(n表示序列長度,d表示隱表徵維度)

2.2.2 模型壓縮技術

本綜述將大語言模型高效推理領域的技術劃分為三個層次,分別為:

🔹模型量化(Model Quantization):模型量化是一類應用廣泛的模型壓縮技術,其透過將模型的權重從高位元數轉換成低位元數來降低模型的計算和儲存開銷。值得注意的是,大模型在推理的不同階段有不一樣的效率瓶頸,因此需要適配不同的量化方法。在預填充階段,推理效率主要受限於較高的計算量,因此通常採用權重啟用量化(Weight-Activation Quantization)的方式;在解碼階段,推理效率主要受限於較高的訪存量,因此通常採用僅權重量化(Weight-only Quantization)的方式。從量化流程來看,模型量化可以劃分為:
  • 訓練後量化(Post-Training Quantization,PTQ):指對預訓練後的模型直接做量化,不需要重新訓練量化後的模型。這類方法對量化後模型的精度損失較大,但因其不需要額外的模型訓練開銷,因此在大語言模型領域應用最為廣泛。下表總結了典型的訓練後量化方法在各個維度的比較。此外,在實際應用中,開發者通常還需要關注量化對大語言模型各種能力的損失,無問芯穹於2024年1月釋出的工作QLLM-Eval[204](已被ICML2024錄用)從不同量化模型、量化引數和量化維度全面研究了量化對模型能力的影響,並給出了針對量化技術的使用指南。
圖片
圖注:典型的訓練後量化演算法的比較

  • 訓練感知量化(Quantization-Aware Training,QAT):指在模型的訓練中加入模型量化過程,並透過訓練減小量化帶來的精度損失。相比於訓練後量化,訓練感知量化方法通常具有更高的精度,但其需要大量的資料和計算量來做模型訓練。因此,目前該子領域的研究主要關注於如何在資料層面和計算量層面降低模型訓練帶來的開銷。

🔹模型稀疏(Model Sparsification)。模型稀疏分為權重稀疏(即權重剪枝)和注意力稀疏(即稀疏注意力):
  • 權重剪枝(Weight Pruning):指透過將模型中不重要的權重和對應結構移除,降低模型的計算和儲存開銷。權重剪枝可以分為非結構化剪枝(Unstructural Pruning)和結構化剪枝(Structural Pruning)兩大類,兩者區別主要在於剪枝粒度不同,如下圖所示。其中,在非結構化剪枝領域,目前的研究工作主要關注於如何加速對大模型的剪枝過程,以及如何設計更有效的重要度分析指標和剪枝率分配策略。而在結構化剪枝領域,目前的研究工作大多關注於如何設計規整的剪枝結構來支援結構化的剪枝。
圖片
圖注:非結構化剪枝和結構化剪枝比較

  • 稀疏注意力(Sparse Attention):指透過減少冗餘的注意力計算,來降低預填充階段的計算開銷和解碼階段中KV cache帶來儲存和訪存開銷。該領域的研究工作主要關注於設計更有效的稀疏模式(Sparse Pattern),包括靜態稀疏(下圖中(a)和(b)所示)以及動態稀疏(下圖中(c)和(d)所示)。無問芯穹於2023年9月釋出的工作SemSA[204],透過對每個注意力頭自動選擇注意力掩膜和掩膜擴充套件方式,在平均稀疏度相同的情況下,大幅提升稀疏注意力大語言模型的有效上下文長達3.9倍。圖片
圖注:典型注意力稀疏模式比較

🔹結構最佳化(Structure Optimization):結構最佳化技術指透過修改模型的架構或結構來達到更好的精度-效率之間的權衡。在該領域有兩類代表性的技術:
  • 神經網路架構搜尋(Neural Architecture Search):指自動化地搜尋出最優的模型架構。然而,這類方法目前只在中等規模的語言模型上得到應用,在大語言模型上還未獲得實際的最佳化效果,原因是該技術通常需要在搜尋過程中對取樣到的架構進行訓練評估,對大模型來說需要花費巨大的訓練開銷。
  • 低秩分解(Low Rank Factorization):指將一個大的權重矩陣近似分解成兩個低質小矩陣的乘積,透過該技術,可以降低大語言模型權重的儲存開銷和訪存開銷。該領域的研究工作主要聚焦於設計分解方式,以及將該技術與其他模型壓縮技術(例如量化、剪枝等)結合獲得更高的效率最佳化效果。

🔹知識蒸餾(Knowledge Distillation):知識蒸餾指用一個大模型(教師模型)來輔助訓練一個小模型(學生模型),從而將大模型的知識傳遞給小模型,透過小模型更小的開銷達到相似的精度效果。知識蒸餾主要分為白盒(White-box)和黑盒(Black-box),前者可以獲得教師模型的架構和權重,因此可以利用更多的資訊(例如特徵、輸出機率等)訓練學生模型,而後者多針對基於API介面訪問的大模型,這類模型的架構和權重無法獲取,因此僅能透過構造資料來訓練學生模型。在知識蒸餾領域,當前的研究工作主要關注設計更有效的蒸餾損失函式,以及構造更有效的蒸餾資料來輔助學生模型的學習。
圖片
圖注:白盒蒸餾和黑盒蒸餾比較

🔹動態推理(Dynamic Inference):動態推理技術指在推理過程中,基於輸入資料的不同,動態決定最合適的推理結構。根據所依賴的輸入資料維度,可以進一步將該領域的技術分為樣本級別(Sample-level)和詞塊級別(Token-level)的動態推理。該領域的研究方向主要聚焦於設計更有效的退出模組,以及解決大語言模型解碼過程中並行推理和KV快取的問題。
圖片
圖注:樣本級別和詞塊級別的動態推理比較

2.2.3 知識、建議和未來方向

在高效結構設計領域,Transformer的替代架構受到了許多研究者們的關注,諸如Mamba[73]、RWKV[60]以及它們的變體架構都在各類任務上驗證了它們高效的推理能力。然而,這些新架構相比Transformer架構在某些方面是否仍然存在缺陷還並未被探究清楚。此外,本文認為將新架構和傳統的Transformer架構做融合,可以達到效率和精度之間更好的權衡,是一個有潛力的研究方向。
在模型壓縮領域:
  • 模型量化是目前最常用也是最推薦使用的模型壓縮技術,一方面,該技術可以很方便快捷地壓縮大語言模型,另一方面,該技術能在有效提升模型效率的同時,儘可能地保證模型的精度損失最小。然而,模型量化技術仍然不可避免地會對模型的能力造成負面影響,因此需要根據任務和應用場景謹慎選取最合適的模型量化演算法。
  • 模型稀疏方面,目前有許多研究工作關注稀疏注意力技術,並將其應用在處理長文字的場景中,然而這類技術通常會損失部分資訊,導致任務精度的下降,因此如何在保留有效資訊的同時還能高效處理長文字是一個值得探索的問題。而在權重剪枝領域,有研究工作[217]指出目前的剪枝方法在不顯著影響任務精度的前提下,僅能達到較低的剪枝率水平,因此設計更有效的模型剪枝演算法是一個值得研究的方向。
  • 結構最佳化方面,神經網路架構搜尋技術受限於其高昂的評估成本,難以被應用在大語言模型的壓縮上,然而這種自動化壓縮的思想仍然值得進一步的探索。另一方面,低秩分解技術受限於其難以同時保證任務精度和壓縮比例。
  • 此外,還有一些工作探究將不同的模型層技術結合,例如將混合專家技術和低秩分解結合、將模型量化和權重稀疏結合、將模型量化和低秩分解結合、將權重剪枝和低秩分解結合。這些方法展示了整合多種壓縮技術以實現大型語言模型更佳最佳化的潛力。

2.3 系統層最佳化技術

大語言模型推理的系統層最佳化主要集中在最佳化模型的前向推理過程。在大語言模型的前向推理計算圖中,注意力運算元和線性運算元佔據了大部分執行時間。系統層最佳化技術包含對這兩種運算元的最佳化,同時還會考慮設計更高效的大語言模型解碼方式。此外,在線上服務的場景中,請求通常來自多個使用者。因此,除了前述最佳化之外,線上服務還面臨與非同步請求相關的記憶體、批處理和排程方面的挑戰。總的來說,本綜述將系統層最佳化技術分為高效的推理引擎(Inference Engine)和服務系統(Serving System)。

2.3.1 推理引擎

圖和運算元最佳化(Graph and Operator Optimization)是高效推理引擎技術中重要的一類方法,其中,運算元最佳化還進一步包含對注意力運算元和線性運算元的最佳化。無問芯穹於2023年11月釋出的工作FlashDecoding++[231](已被MLSys2024錄用),透過對注意力和線性運算元的針對性最佳化和計算圖層面的深度運算元融合技術,顯著提高了大語言模型推理效率,在NVIDIA和AMD GPU上,吞吐率相較HuggingFace提升分別達到4.86和4.35倍。
圖片
圖注:高效推理引擎領域研究總結

猜測解碼(Speculative Decoding)是一類透過並行解碼來加速大模型解碼過程的技術,其核心思想是使用一個小模型來預測未來的若干個詞塊,再用大模型並行地驗證這些詞塊是否準確。具體來說,猜測解碼可以大致分為兩個階段:猜測階段驗證階段
猜測階段,對於當前的輸入文字,猜測解碼方法首先利用一個“便宜”的草稿模型(Draft Model)生成連續的若干個草稿詞塊(Draft Token),值得注意的是,該草稿模型往往具有相比於大語言模型小得多的引數量和計算量,因此這一過程帶來的開銷非常小。
驗證階段,猜測編碼使用大語言模型對草稿詞塊進行並行驗證,並接收透過驗證的草稿詞塊作為解碼結果。猜測編碼的加速效果與草稿模型的預測準確率和驗證接收方式有關,現有大多研究工作關注於設計更好的草稿模型,或設計更有效的驗證接收方式。下表對比了典型的猜測解碼演算法。
圖片
圖注:典型猜測解碼演算法比較

2.3.2 服務系統

服務系統的最佳化旨在提升系統處理非同步請求的效率。這類最佳化技術主要包括:記憶體管理(Memory Management)、連續批處理(Continuous Batching )、排程策略(Scheduling Strategy)和分散式系統( Distributed Systems)。此外,本文還比較了市面上常見的開源推理框架的推理延時和系統吞吐,如下表所示。
圖片
圖注:高效服務系統領域研究總結
圖片
圖注:開源推理框架比較

2.3.3 知識、建議和未來方向

系統層最佳化技術在不降低準確性的情況下提高了推理效率,因此在大語言模型的時間部署中非常普遍。在最近的研究中,運算元最佳化已經與實際的服務場景密切結合,例如,RadixAttention[50]專門被設計用於字首快取,而tree attention技術[225]用於加速猜測解碼的驗證。應用和場景的更迭將繼續為運算元的發展提出新的要求和挑戰。

三、未來研究展望

本文進一步總結了未來的四個關鍵應用場景,並討論了高效性研究在這些場景中的重要性:
  • 能體和多模型框架。在最近的研究中,大語言模型智慧體和多模型協同框架受到了許多關注,這類技術可以提升大語言模型的應用能力,使得模型能更好地服務於人類。然而模型的數量增多,以及輸入到模型中的指令變長,都會使得智慧體或框架系統的推理效率變低。因此需要面向這些框架和場景進行大模型的推理效率最佳化。
  • 長文字場景。隨著輸入模型文字變得越來越長,大語言模型的效率最佳化需求變得愈發提升。目前在資料層、模型層和系統層均有相關的技術來最佳化長文字場景下大語言模型的推理效率,其中設計Transformer的替代新架構受到了許多關注,然而這類架構還未被充分探索,其是否能匹敵傳統的Transformer模型仍未清楚。
  • 邊緣端部署。最近,許多研究工作開始關注於將大語言模型部署到邊緣裝置上,例如移動手機。一類工作致力於設計將大模型變小,透過直接訓練小模型或模型壓縮等途徑達到該目的;另一類工作聚焦於系統層的最佳化,透過運算元融合、記憶體管理等技術,直接將70億引數規模的大模型成功部署到移動手機上。
  • 安全-效率協同最佳化。除了任務精度和推理效率外,大語言模型的安全性也是一個需要被考量的指標。當前的高效性研究均未考慮最佳化技術對模型安全性的影響。若這些最佳化技術對模型的安全性產生了負面影響,一個可能的研究方向就是設計新的最佳化方法,或改進已有的方法,使得模型的安全性和效率能一同被考量。

如果你對文章感興趣,可以聯絡學術作者進一步討論:ningxuefei@mail.tsinghua.edu.cn

相關文章