1 摘要
自 OpenAI 發布 ChatGPT 以來,基于 Transformer 架構的大語言模型 (LLM) 在全球范圍內引發了深度的技術關注,并取得了令人矚目的成就。其強大的理解和生成能力,正在深刻改變我們對人工智能的認知和應用。然而大語言模型的推理應用成本過高,高昂的成本大大阻礙了技術落地。因此,大語言模型的推理性能優化成為業界研究的熱點。
大語言模型推理面臨計算資源的巨大需求和計算效率的挑戰。優化推理性能不僅可以減少硬件成本,還可以提高模型的實時響應速度。它使模型能夠更快速地執行自然語言理解、翻譯、文本生成等任務,從而改善用戶體驗,加速科學研究,推動各行業應用的發展。
本文從推理服務系統全局視角介紹典型性能優化技術和各自特點,最后分析未來大語言模型推理優化技術的發展趨勢和演進方向,最終為未來的人工智能應用打開更廣闊的可能性。
2 優化技術
LLM 推理服務重點關注兩個指標:吞吐量和時延:
吞吐量:主要從系統的角度來看,即系統在單位時間內能處理的 tokens 數量。計算方法為系統處理完成的 tokens 個數除以對應耗時,其中 tokens 個數一般指輸入序列和輸出序列長度之和。吞吐量越高,代表 LLM 服務系統的資源利用率越高,對應的系統成本越低。
時延:主要從用戶的視角來看,即用戶平均收到每個 token 所需位時間。計算方法為用戶從發出請求到收到完整響應所需的時間除以生成序列長度。一般來講,當時延不大于 50 ms/token 時,用戶使用體驗會比較流暢。
吞吐量關注系統成本,高吞吐量代表系統單位時間處理的請求大,系統利用率高。時延關注用戶使用體驗,即返回結果要快。這兩個指標一般情況下需要會相互影響,因此需要權衡。例如,提高吞吐量的方法一般是提升 batchsize,即將用戶的請求由串行改為并行。但 batchsize 的增大會在一定程度上損害每個用戶的時延,因為以前只計算一個請求,現在合并計算多個請求,每個用戶等待的時間變長。
LLM 推理性能優化主要以提高吞吐量和降低時延為目的,具體可以劃分為如下六部分,下面詳細展開描述。
2.1 顯存相關優化
2.1.1 KV Cache
大模型推理性能優化的一個最常用技術就是 KV Cache,該技術可以在不影響任何計算精度的前提下,通過空間換時間思想,提高推理性能。目前業界主流 LLM 推理框架均默認支持并開啟了該功能。
Transformer 模型具有自回歸推理的特點,即每次推理只會預測輸出一個 token,當前輪輸出token 與歷史輸入 tokens 拼接,作為下一輪的輸入 tokens,反復執行多次。該過程中,前后兩輪的輸入只相差一個 token,存在重復計算。KV Cache 技術實現了將可復用的鍵值向量結果保存下來,從而避免了重復計算。
具體來講,KV Cache 技術是指每次自回歸推理過程中,將 Transformer 每層的 Attention 模塊中的 和 結果保存保存在一個數據結構(稱為 KV Cache)中,當執行下一次自回歸推理時,直接將 和 與 KV Cache 拼接在一起,供后續計算使用。其中, 代表第 步推理的輸入, 和 分別代表鍵值權重矩陣。
KV Cache 緩存每一輪已計算完畢的鍵值向量,因此會額外增加顯存開銷。以 LLaMA-7B 模型為例,每個 token 對應的 KV Cache 空間 可通過如下公式計算:
公式中第一個因子 2 代表 Key/Value 兩個向量,每層都需存儲這兩個向量, 為 Transformer layer 個數, 代表 KV head 個數(模型為多頭注意力時,該值即注意力頭數,模型為多查詢注意力時,該值為 1), 為每個 KV head 的維度, 為每存放一個數據所需的字節數。模型推理所需的 KV Cache 總量為公式 如下,其中 為輸入和輸出序列長度之和。因此,KV Cache 與 batchsize 和序列長度呈線性關系。
KV Cache 的引入也使得推理過程分為如下兩個不同階段,進而影響到后續的其他優化方法。
預填充階段:發生在計算第一個輸出 token 過程中,計算時需要為每個 Transformer layer 計算并保存 key cache 和 value cache;FLOPs 同 KV Cache 關閉一致,存在大量 GEMM (GEneral Matrix-Matrix multiply) 操作,屬于 Compute-bound 類型計算。
解碼階段:發生在計算第二個輸出 token 至最后一個 token 過程中,這時 KV Cache 已存有歷史鍵值結果,每輪推理只需讀取 Cache,同時將當前輪計算出的新的 Key、Value 追加寫入至 Cache;GEMM 變為 GEMV (GEneral Matrix-Vector multiply) 操作,FLOPs 降低,推理速度相對預填充階段變快,這時屬于 Memory-bound 類型計算。
2.1.2 Paged Attention
LLM 推理服務的吞吐量指標主要受制于顯存限制。研究團隊發現現有系統由于缺乏精細的顯存管理方法而浪費了 60% 至 80% 的顯存,浪費的顯存主要來自 KV Cache。因此,有效管理 KV Cache 是一個重大挑戰。
在 Paged Attention 之前,業界主流 LLM 推理框架在 KV Cache 管理方面均存在一定的低效。HuggingFace Transformers 庫中,KV Cache 是隨著執行動態申請顯存空間,由于 GPU顯存分配耗時一般都高于 CUDA kernel 執行耗時,因此動態申請顯存空間會造成極大的時延開銷,且會引入顯存碎片化。FasterTransformer 中,預先為 KV Cache 分配了一個充分長的顯存空間,用于存儲用戶的上下文數據。例如 LLaMA-7B 的上下文長度為 2048,則需要為每個用戶預先分配一個可支持 2048 個 tokens 緩存的顯存空間。如果用戶實際使用的上下文長度低于2048,則會存在顯存浪費。Paged Attention 將傳統操作系統中對內存管理的思想引入 LLM,實現了一個高效的顯存管理器,通過精細化管理顯存,實現了在物理非連續的顯存空間中以極低的成本存儲、讀取、新增和刪除鍵值向量。
具體來講,Paged Attention 將每個序列的 KV Cache 分成若干塊,每個塊包含固定數量token 的鍵和值。
首先在推理實際任務前,會根據用戶設置的和 預跑一次推理計算,記錄峰值顯存占用量 ,然后根據上面公式獲得當前軟硬件環境下 KV Cache 可用的最大空間,并預先申請緩存空間。其中,為部署環境的硬件顯存一次最多能容納的 token 總量, 為模型推理的最大顯存占用比例, 為物理顯存量, 為塊大小(默認設為 16)。
在實際推理過程中,維護一個邏輯塊到物理塊的映射表,多個邏輯塊可以對應一個物理塊,通過引用計數來表示物理塊被引用的次數。當引用計數大于一時,代表該物理塊被使用,當引用計數等于零時,代表該物理塊被釋放。通過該方式即可實現將地址不連續的物理塊串聯在一起統一管理。
Paged Attention 技術開創性地將操作系統中的分頁內存管理應用到 KV Cache 的管理中,提高了顯存利用效率。另外,通過 token 塊粒度的顯存管理,系統可以精確計算出剩余顯存可容納的 token 塊的個數,配合后文 Dynamic Batching 技術,即可避免系統發生顯存溢出的問題。
2.2 計算相關優化
2.2.1 算子融合
算子融合是深度學習模型推理的一種典型優化技術,旨在通過減少計算過程中的訪存次數和 Kernel 啟動耗時達到提升模型推理性能的目的,該方法同樣適用于 LLM 推理。
以 HuggingFace Transformers 庫推理 LLaMA-7B 模型為例,經分析模型推理時的算子執行分布如下圖所示,該模型有 30 個類型共計 2436 個算子,其中 aten::slice 算子出現頻率為 388 次。大量小算子的執行會降低 GPU 利用率,最終影響推理速度。
目前業界基本都針對 Transformer layer 結構特點,手工實現了算子融合。以 DeepSpeed Inference 為例,算子融合主要分為如下四類:
歸一化層和 QKV 橫向融合:將三次計算 Query/Key/Value 的操作合并為一個算子,并與前面的歸一化算子融合。
自注意力計算融合:將自注意力計算涉及到的多個算子融合為一個,業界熟知的 FlashAttention 即是一個成熟的自注意力融合方案。
殘差連接、歸一化層、全連接層和激活層融合:將 MLP 中第一個全連接層上下相關的算子合并為一個。
偏置加法和殘差連接融合。
由于算子融合一般需要定制化實現算子 CUDA kernel,因此對 GPU 編程能力要求較高。隨著編譯器技術的引入,涌現出 OpenAI Triton 、TVM 等優秀的框架來實現算子融合的自動化或半自動化,并取得了一定的效果。
2.2.2 高性能算子
針對 LLM 推理運行熱點函數編寫高性能算子,也可以降低推理時延。
GEMM 操作相關優化:在 LLM 推理的預填充階段,Self-Attention 和 MLP 層均存在多個 GEMM 操作,耗時占據了推理時延的 80% 以上。GEMM 的 GPU 優化是一個相對古老的問題,在此不詳細展開描述算法細節。英偉達就該問題已推出 cuBLAS、CUDA、CUTLASS 等不同層級的優化方案。例如,FasterTransformer 框架中存在大量基于 CUTLASS 編寫的 GEMM 內核函數。另外,Self-Attention 中存在 GEMM+Softmax+GEMM 結構,因此會結合算子融合聯合優化。
GEMV 操作相關優化:在 LLM 推理的解碼階段,運行熱點函數由 GEMM 變為 GEMV。相比 GEMM,GEMV 的計算強度更低,因此優化點主要圍繞降低訪存開銷開展。
高性能算子的實現同樣對 GPU 編程能力有較高要求,且算法實現中的若干超參數與特定問題規模相關。因此,編譯器相關的技術如自動調優也是業界研究的重點。
2.3 服務相關優化
服務相關優化主要包括 Continuous Batching、Dynamic Batching 和 異步 Tokenize / Detokenize。其中 Continuous Batching 和 Dynamic Batching 主要圍繞提高可并發的 batchsize 來提高吞吐量,異步 Tokenize / Detokenize 則通過多線程方式將 Tokenize / Detokenize 執行與模型推理過程時間交疊,實現降低時延目的。
?
問題分類 | 現象 | 解決方法 | 實現原理 | 特點 |
---|---|---|---|---|
問題一 | 同批次序列推理時,存在“氣泡”,導致 GPU 資源利用率低 | Continuous Batching | 由 batch 粒度的調度細化為 step 級別的調度 | 在時間軸方向動態插入新序列 |
問題二 | 批次大小固定不變,無法隨計算資源負載動態變化,導致 GPU 資源利用率低 | Dynamic Batching | 通過維護一個作業隊列實現 | 在 batch 維度動態插入新序列 |
問題三 | Tokenize / Detokenize 過程在 CPU 上執行,期間 GPU 處于空閑狀態 | 異步 Tokenize / Detokenize | 多線程異步 | 流水線 overlap 實現降低時延 |
?
大語言模型的輸入和輸出均是可變長度的。對于給定問題,模型在運行前無法預測其輸出長度。在實際服務場景下,每個用戶的問題長度各不相同,問題對應的答案長度也不相同。傳統方法在同批次序列推理過程中,存在“氣泡”現象,即必須等同批次內的所有序列完成推理之后,才會執行下一批次序列,這就會引起 GPU 資源的浪費,導致 GPU 利用率偏低。
圖中序列 3 率先結束,但由于其他序列尚未結束,因此需要等待直至所有序列計算完畢。理想情況下,同批次的所有序列的輸入加輸出的長度均相同,這時不存在“氣泡”現象;極端情況下則會出現超過 50% 以上的資源浪費。
另一方面,傳統方法推理時 batchsize 是固定不變的,無法隨計算資源負載動態變化。比如某一段時間內,同批次下的序列長度都偏短,原則上可以增加 batchsize 以充分利用 GPU 計算資源。然而由于固定 batchsize,無法動態調整批次大小。
Continuous Batching 和 Dynamic Batching 思想最早來自論文 Orca: A Distributed Serving System for Transformer-Based Generative Models。針對問題一,提出 Continuous Batching,原理為將傳統 batch 粒度的任務調度細化為 step 級別的調度。首先,調度器會維護兩個隊列,分別為 Running 隊列和 Waiting 隊列,隊列中的序列狀態可以在 Running 和 Waiting 之間轉換。在自回歸迭代生成每個 token 后,調度器均會檢查所有序列的狀態。一旦序列結束,調度器就將該序列由 Running 隊列移除并標記為已完成,同時從 Waiting 隊列中按 FCFS (First Come First Service) 策略取出一個序列添加至 Running 隊列。
圖中,序列 3 率先在 T5 時刻結束,這時調度器會檢測到序列 3 已結束,將序列 3 從 Running 隊列中移除,并從 Waiting 隊列中按 FCFS 策略取出序列 5 添加至 Running 隊列并啟動該序列的推理。通過該方法,即可最大限度地消除“氣泡”現象。
問題一可以理解為在時間軸方向動態插入新序列,問題二則是在 batch 維度動態插入新序列,以盡可能地充分利用顯存空間。具體來講,在自回歸迭代生成每個 token 后,調度器通過當前剩余顯存量,動態調整 Running 隊列的長度,從而實現 Dynamic Batching。例如,當剩余顯存量較多時,會盡可能增加 Running 隊列長度;當待分配的 KV Cache 超過剩余顯存時,調度器會將 Running 隊列中低優先級的序列換出至 Waiting 隊列,并將換出序列占用的顯存釋放。
如上兩個 batching 相關的優化技術可有效提升推理吞吐量,目前已在 HuggingFace Text-Generation-Interface (TGI)、vLLM、OpenPPL-LLM 等多個框架中實現。
2.4 分布式相關優化
由于大語言模型參數量較大,可能無法存放到單一計算設備中,分布式并行可以有效解決該問題。分布式并行中的模型并行和流水線并行已在 LLM 推理中得到應用。由于篇幅有限,本文聚焦模型并行。模型并行通過將權重參數拆分到多個計算設備中,實現分布式計算。
圖中,第一行代表 Column Parallel,即將權重數據按列拆分到多個 GPU 中,每個 GPU 上的本地計算結果需要在列方向拼接為最終結果;第二行代表 Row Parallel,即將權重數據按行拆分到多個 GPU 中,每個 GPU 上的本地計算結果需要 AllReduce 規約為最終結果。
業界最流行的模型并行方案來自 Megatron-LM,其針對 Self-Attention 和 MLP 分別設計了簡潔高效的模型并行方案。
MLP: 第一個全連接層為 Column Parallel,第二個全連接層為 Row Parallel,整個 MLP 只需在 Row Parallel 后執行一次 AllReduce 規約操作即可。
Self-Attention:在計算 Query、Key 和 Value 向量時執行 Column Parallel(按注意力頭個數均分到每個 GPU),在將注意力得分做空間映射時執行 Row Parallel,整個 Self-Attention 只需在 Row Parallel 后執行一次 AllReduce 規約操作即可。
上面分析了 Transformer layer 的模型并行方式。除此之外,LLM 模型中的 Input Embedding 采用 Row Parallel,Output Embedding 采用 Column Parallel;Dropout / Layer Norm / Residual Connections 等操作都沒有做并行拆分。例如 Layer Norm 的權重參數和計算,在每個 GPU 上都是完整的。
?
Layers | Model Parallel Method |
---|---|
Input Embedding | Row Parallel |
Self-Attention | Column Parallel + Row Parallel |
MLP | Column Parallel + Row Parallel |
Output Embedding | Column Parallel |
?
基于以上基礎,以 LLaMA-34B 模型為例進行通信量分析。該模型包含 48 個 Transformer layers,隱藏層大小 8192,每次單 batch 推理共 2 * 48 次 Broadcast 和 248 次 AllReduce 操作,每次通信傳輸的數據量均為 16 KB(此處假設數據類型為半精度浮點,81922/1024=16 KB)。考慮到推理服務一般都是按多 batch 推理執行,假設 batchsize 為 64,每次通信傳輸的數據量也僅為 1 MB。下圖在 A100-PCIE-40GB 機器上測試 NCCL AllReduce 帶寬數據,PCIE 理論帶寬為 32-64 GB/s 左右,實際推理場景下的通信數據量主要集中在 1 MB 以下,對應的實際帶寬約為 1-10 GB/s。NVLink 理論帶寬為 400-600 GB/s,但由于每次的通信量很小,實際帶寬也遠遠小于理論帶寬。因此模型參數量越大、batchsize 越大,通信效率越高,使用模型并行獲得的收益約明顯。
2.5 低比特量化
回歸到 LLM 模型推理吞吐量和時延這兩個重要的性能指標上:吞吐量的提升主要受制于顯存容量,如果降低推理時顯存占用量,就可以運行更大的 batchsize,即可提升吞吐量;LLM 推理具有 Memory-bound 特點,如果降低訪存量,將在吞吐量和時延兩個性能指標上都有收益。低比特量化技術可以降低顯存占用量和訪存量,其能取得加速的關鍵在于顯存量和訪存量的節省以及量化計算的加速遠大于反量化帶來的額外開銷。
?
被量化的對象 | 量化方法 | 特點 |
---|---|---|
權重量化 | LLM.int8(), GPTQ | 顯存占用減半,但由于計算結果需反量化,時延基本無收益 |
權重和激活同時量化 | SmoothQuant | 顯存占用減半,時延有收益,精度幾乎匹配 FP16 |
KV Cache量化 | INT8 或 FP8 量化 | 方法簡單,吞吐量收益明顯 |
基于硬件特點的量化:英偉達 Hopper 架構下的 FP8 | 直接利用 TensorCore FP8 計算指令 | 不需要額外的量化/反量化操作,時延收益明顯 |
?
表中的四類量化方法各有特點,業界在低比特量化方向的研究進展也層出不窮,希望探索出一個適用于大語言模型的、能夠以較高壓縮率壓縮模型、加速端到端推理同時保證精度的量化方法。
2.6 其他新技術
當前,業界在將傳統優化技術引入 LLM 推理的同時,同時也在探索從大模型自回歸解碼特點出發,通過調整推理過程和引入新的模型結構來進一步提升推理性能。
例如,投機采樣(Speculative decoding)針對 LLM 推理串行解碼特點,通過引入一個近似模型來執行串行解碼,原始模型執行并行評估采樣,通過近似模型和原始模型的互相配合,在保證精度一致性的同時降低了大模型串行解碼的次數,進而降低了推理時延。美杜莎頭(Medusa head)則是對投機采樣的進一步改進,其摒棄了近似模型,在原始模型結構上新增了若干解碼頭,每個解碼頭可并行預測多個后續 tokens,然后使用基于樹狀注意力機制并行處理,最后使用典型接收方案篩選出合理的后續 tokens。該方法同樣降低了大模型串行解碼的次數,最終實現約兩倍的時延加速。
3 總結
大語言模型推理性能優化技術正迅速演進,不僅涉及計算機科學和人工智能領域,還融合了多個學科的知識,實現了前所未有的跨學科交叉滲透。演進的動力源自對大規模模型應用的需求,為了充分發揮這些模型的潛力,研究人員正在不斷改進推理性能,包括算法優化、硬件加速、分布式計算等方面的創新。這一快速演進和跨學科滲透的趨勢不僅將提高大語言模型的實用性,還為未來的自然語言處理應用和人工智能技術帶來更大的創新和應用潛力。
參考
W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. H. Yu,J. E. Gonzalez, H. Zhang and I. Stoica, Efficient MemoryManagement for Large Language Model Serving with Page-dAttention, Proceedings of the ACM SIGOPS 29th Sympo-sium on Operating Systems Principles, 2023.
S. Z. Y. S. L. Z. C. Y. J. G. H. Z. Woosuk Kwon, Zhuohan Liand I. Stoica, vLLM: Easy, Fast, and Cheap LLM Servingwith PagedAttention,https://vllm.ai/, 2022.
HuggingFace, ?Transformers: ? State-of-the-art ?MachineLearning for Pytorch, TensorFlow, and JAX.,https://github.com/huggingface/transformers.
NVIDIA, ? ?FasterTransformer,https://github.com/NVIDIA/FasterTransformer, 2021.
R. Y. Aminabadi, S. Rajbhandari, A. A. Awan, C. Li, D. Li,E. Zheng, O. Ruwase, S. Smith, M. Zhang, J. Rasley et al.,DeepSpeed-inference: enabling efficient inference of trans-former models at unprecedented scale, SC22: InternationalConference for High Performance Computing, Networking,Storage and Analysis, 2022, 1–15.
P. Tillet, H.-T. Kung and D. Cox, Triton: an intermedi-ate language and compiler for tiled neural network com-putations, Proceedings of the 3rd ACM SIGPLAN Inter-national Workshop on Machine Learning and ProgrammingLanguages, 2019, 10–19.
T. Chen, T. Moreau, Z. Jiang, L. Zheng, E. Yan, H. Shen,M. Cowan, L. Wang, Y. Hu, L. Ceze et al., TVM: An au-tomated End-to-End optimizing compiler for deep learning,13th USENIX Symposium on Operating Systems Design andImplementation (OSDI 18), 2018, 578–594.
M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casperand B. Catanzaro, Megatron-lm: Training multi-billion pa-rameter language models using model parallelism, arXivpreprint arXiv:1909.08053, 2019.
T. Dettmers, M. Lewis, Y. Belkada and L. Zettlemoyer, Llm.int8 (): 8-bit matrix multiplication for transformers at scale,arXiv preprint arXiv:2208.07339, 2022.
E. Frantar, ?S. Ashkboos, ?T. Hoefler and D. Alistarh,Gptq: ?Accurate post-training quantization for generativepre-trained transformers, arXiv preprint arXiv:2210.17323,2022.
G. Xiao, J. Lin, M. Seznec, H. Wu, J. Demouth and S. Han,Smoothquant: Accurate and efficient post-training quanti-zation for large language models, International Conferenceon Machine Learning, 2023, 38087–38099.
A. C. Elster and T. A. Haugdahl, Nvidia hopper gpu andgrace cpu highlights, Computing in Science & & Engineering, 2022, 24, 95–100.
C. Chen, S. Borgeaud, G. Irving, J.-B. Lespiau, L. Sifreand J. Jumper, Accelerating large language model decodingwith speculative sampling, arXiv preprint arXiv:2302.01318,2023.
Y. Leviathan, M. Kalman and Y. Matias, Fast inference fromtransformers via speculative decoding, International Confer-ence on Machine Learning, 2023, 19274–19286.
T. Cai, Y. Li, Z. Geng, H. Peng and T. Dao, Medusa: SimpleFramework for Accelerating LLM Generation with MultipleDecoding ?Heads,https://github.com/FasterDecoding/Medusa, 2023.
審核編輯:黃飛
?
評論