一、背景

最近我們在生產環境批量部署了大模型專用推理集群，并成功讓包括70B在內的大模型推理速度提升50%，大幅縮減部署成本，穩定應用于生產環境。本文基于我們在部署大模型推理集群時的一些經驗，分享一些有效提升大模型的推理速度方法。最后，我們在結尾處推薦了幾個經過我們評測且表現優異的大模型推理框架。希望這些建議能幫助讀者在項目中選擇適合自己的推理框架。

OpenAI的科學家Hyung Won Chung在2023年的公開演講《Large Language Models》[8]中指出，大模型的某些能力僅在達到特定規模時才能顯現，可見未來大模型的參數量肯定會越來越大，這也是大模型的發展趨勢。隨著參數量的增加，對大模型的推理速度要求越來越高，有哪些方法可以提高大模型的推理速度或吞吐量？

首先我們將探討大模型的加速優化方向，隨后文章將依據時間線，介紹一些業界內較為經典的實用大模型加速技術，包括但不限于“FlashAttention[1]”和“PageAttention[3]”等技術。

以下為按時間順序業界的一些經典大模型推理加速技術，本文試圖為讀者提供一個按時間發展順序的大模型加速方法綜述。

除了上面提到的技術外，提高大模型推理速度的還有大模型的量化技術等，這里先不探討，后面有機會，我們會單獨發文章來介紹。

二、大模型發展面臨的挑戰

未來大模型的參數量肯定會越來越大，這也是大模型的發展趨勢，對推理加速的要求會越來越高。

OpenAI在其論文《Scaling Laws for Neural Language Models》[7]中介紹了大模型的擴展規則，這些規則闡釋了模型能力與其規模之間的關系。具體來說，模型的能力強烈依賴于其規模，包括模型參數的數量，數據集的大小，以及訓練過程中所需的計算量。此外，OpenAI的科學家Hyung Won Chung在2023年的公開演講《Large Language Models》[8]中指出，大模型的某些能力僅在達到特定規模時才能顯現。

上圖摘自Hyung Won Chung演講中的ppt[8]。圖中主要表達一個觀點，隨著模型規模的增大，比如GPT3到GPT4，模型的能力變的越來越強，甚至會出現新的能力。

但是隨著模型的規模增大，大模型的推理速度將會逐漸降低，這是因為更多的參數量需要更多的GPU計算。推理速度的下降進一步帶來更差的用戶體驗，因此如何對大模型推理加速變得越來越重要了。

三、大模型推理加速的優化方向

Llama2的模型結構

我們先簡單了解一下Llama 2模型系列的結構，參考自Llama 2的論文[9]。目前，像Llama系列這樣的大多數生成式語言模型，主要采用了Transformer架構中的Decoder模塊。在Huggingface平臺上，這類模型結構通常被稱作CausalLM，即因果語言模型。

上圖為Llama2大模型的結構，其中最核心的是注意力計算(Llama Attention)。這也是整個推理過程中最耗費時間的模塊，后面的優化大部分都是基于Attention去實施的。為了更好的理解Llama 2大模型的結構，我們先簡單對Llama2模型的整個推理過程進行拆解，不感興趣同學可以直接跳過。

1. 用戶向模型提交Prompt后，模型首先進行的操作是預測下一個字符(Token)，并將預測出的字符添加到輸入中繼續進行預測。這個過程會一直持續，直到模型輸出一個停止符號(STOP token)，此時預測停止，模型輸出最終結果。
2. 在生成下一個字符(Token)的過程中，模型需要執行N次的Llama解碼器層(Llama Decoder Layer)計算。具體來說，Llama-2-7B模型執行32次計算，而Llama-2-13B模型執行40次。
3. Llama解碼器層(Llama Decoder Layer)中最關鍵的計算環節是注意力(Llama Attention)的計算。大部分推理時間都消耗在Attention的計算上，因此多種優化技巧都旨在提高Attention計算的效率。

大模型推理的加速方向有哪些

從Llama 2模型的結構分析中，我們可以總結出大模型在推理計算過程中表現出以下特點：

1. 在整個推理過程中，最耗時的部分為注意力(Attention)計算。針對Attention的計算進行速度優化，可以顯著提高整體推理性能。
2. 注意力(Attention)計算過程中，鍵值對緩存(KV Cache)占用了大量顯存資源。以13B模型為例，處理一個Prompt序列大約需要3GB額外顯存，并且這部分顯存會被頻繁地分配和釋放，產生大量碎片，如果能減少顯存碎片，也能提升大模型的吞吐。
3. 推理過程GPU需要處理和計算大量的參數。7B模型擁有70億參數，而13B模型則包含130億參數，最新全球最強大模型DBRX更是高達1300億參數，這需要高效地處理這些參數。這里也可以有優化空間。

針對上述三個特性，目前業界提出了多種有效的優化方法，典型如下：

1. FlashAttention-Attention計算速度優化

FlashAttention[1]在不改變Attention算子的計算結果的前提下，提升Attention算子的計算速度。FlashAttention在各種模型和任務上展示了顯著的性能提升。例如，在BERT-large、GPT-2等模型上，相比于基線實現，FlashAttention能夠實現15%到3倍的端到端加速。

2. PageAttention-KV Cache顯存管理優化

PageAttention[3]的目標是減少顯存碎片，基于PageAttention的VLLM系統能夠將流行的大型語言模型（LLM）的吞吐量提高到10倍以上，同時保持耗時分布平穩。

3. MOE-縮減推理時模型參數

MOE(Mixture of Experts)[4]目標是減少模型推理時參與計算的參數量。

實驗效果：Mixtral模型在多數基準測試中表現優于Llama 2 70B模型，并且其推理速度比后者快了6倍。該模型支持多種語言，具有強大的代碼生成能力，并可以細化配置以遵循具體指令，從而在MT-Bench基準測試中取得了高分。

后面我們將針對上面的每個方向詳細介紹。

四、FlashAttention-Attention算子計算優化

FlashAttention先后發表了兩篇論文闡述對Attention算子的優化，包括FlashAttention-1[1]與FlashAttention-2[2]，我們以FlashAttention-1[1]為例了解下他的優化原理。

我們先了解下GPU的內存分層結構，參考下圖，圖片來自論文FlashAttention-1[1]。

GPU的內存層次結構由三個主要部分組成：SRAM、HBM和DRAM，下面為A100GPU的參考配置。

SRAM（靜態隨機訪問存儲器）具有最快的訪問速度（19TB/s），但其容量相對較小（僅20MB）。

HBM（高帶寬存儲器）提供較大的存儲空間（40GB）和高速的數據訪問（1.5TB/s）。

DRAM（動態隨機訪問存儲器），在這里特指GPU外部的主存，容量最大（超過1TB），但訪問速度最慢（12.8GB/s）。

從上述配置中可以看出，內存容量越小，處理速度就越快。

在傳統的Attention計算過程中，大量的輸入/輸出操作都是通過訪問HBM來完成的。FlashAttention算法通過優化Attention計算流程，減少了對HBM的訪問次數，以提高計算效率，所以它是一種IO感知的優化算法。

下圖為FlashAttention的加速方法，來自論文FlashAttention-1[1]

FlashAttention利用了一個聰明的技巧來快速且內存高效地計算注意力機制，即它通過將輸入數據分塊（tiling）來避免一次性處理整個巨大的注意力矩陣，這通常需要大量的內存和計算資源。想象一下，我們有一個巨大的圖書館(矩陣)，而FlashAttention的方法就像是把圖書館里的書分成幾個小堆，然后每次只處理一堆書。這樣，我們就不需要一次性把所有書都拿出來放在桌子上（這需要很大的桌子和很多時間）。

具體來說，FlashAttention在做矩陣計算的時候，通過將數據分塊并利用GPU上的快速但容量較小的存儲（SRAM）去計算，有效減少了對慢速但容量大的存儲（HBM）的訪問。這樣不僅加快了計算速度，而且大幅減少了顯存的需求。

通過減少對慢速存儲的依賴，FlashAttention能夠顯著提高模型訓練的速度，同時保持或甚至提高模型的性能。例如，讓BERT-large的訓練比MLPerf 1.1的記錄快15%，GPT-2訓練速度是HuggingFace和Megatron-LM基線的三倍，長序列領域訓練速度提升至2.4倍。

下圖來自huggingface 對flash attention介紹的blog[14]，可以更好的理解Flash Attention對矩陣拆分的方式。

既然Flash Attention可以加速計算，那么支持Flash Attention計算的框架包括都有哪些，文章后半部我們會推薦一些比較優秀的推理框架。

五、PageAttention-顯存管理優化

PageAttention[3]的概念最初由VLLM的作者Woosuk Kwon提出，它也是VLLM推理框架的最主要的優化策略。Woosuk Kwon在其論文中介紹了如何通過PageAttention來解決大型語言模型（LLM）服務中的一個關鍵問題——在不增加延遲的情況下有效管理內存以提升吞吐量。

我們先了解下大模型在推理的情況下的內存結構分布，下圖來自論文[3]。

這是一個在NVIDIA A100上服務一個擁有13B參數的大型語言模型的內存布局，13B LLM 推理顯存占用分部，13B LLM的參數占用26G顯存，每個請求，KV Cache會占用12G顯存，隨著QPS的增加，KVCache會快速上升，并且會被頻繁的分配與釋放，系統會產生大量的顯存碎片，如果不加處理，系統就會慢慢崩掉。

那么VLLM是如何通過PageAttention解決顯存碎片的問題的呢？下圖來自文章[14]，為VLLM的顯存管理技術。

PageAttention的工作原理是通過將鍵值緩存（KV緩存）分割成固定大小的塊（或“頁面”），并允許這些塊在內存中非連續地存儲。這種方法靈感來源于操作系統的虛擬內存和分頁技術，目的是為了更靈活和高效地管理內存資源。

在傳統的注意力機制中，一個請求的KV緩存需要在內存中連續存儲，這會導致兩個主要問題：內存碎片化和無法高效共享內存。內存碎片化限制了批處理的大小，而無法共享內存則導致重復數據，浪費寶貴的內存資源。

PageAttention通過以下步驟工作來解決這些問題：

1. 分割KV緩存：將每個請求的KV緩存劃分為多個較小的塊，這些塊的大小是固定的，可以根據模型和硬件的具體需求進行調整。
2. 非連續存儲：與傳統KV緩存塊在內存中連續存儲不同，PageAttention允許這些塊在物理內存中非連續地分布。這樣，就可以根據實際需要動態地分配和回收內存塊，減少內存浪費。
3. 動態管理：通過類似于操作系統中虛擬內存管理的方式，PageAttention動態地管理這些內存塊。系統可以根據當前的內存使用情況，按需分配或釋放KV緩存塊，從而優化內存使用。
4. 內存共享：PageAttention還支持在不同請求之間或同一個請求中的不同序列之間共享KV緩存塊。這種共享是靈活的，可以基于塊級別進行，進一步減少內存使用和提高效率。

通過這種方式，PageAttention允許LLM服務系統在保持相同延遲的情況下，通過減少內存浪費和提高內存共享，顯著提高處理請求的吞吐量。

通過PageAttention的優化，VLLM對LLaMA 7B與13B的吞吐量提升了10倍以上，下圖來自文章[11]。

六、MOE-縮減推理時模型參數

最近發布的全球最強開源大模型1300億參數的DBRX，以及Mistral的8x7B開源大模型都是基于MOE架構的。為什么參數量越大的模型越要使用MOE架構呢？我們以Mistral的8x7B開源大模型為例，介紹下MOE架構在性能方面的優勢。

說到MOE大模型，我們先對比下普通大模型與MOE大模型在結構上的區別，參考上圖。在MOE大模型中，把大模型的參數分成了8個小組外加一個路由器，每個小組我們稱作專家組。當請求過來的時候，MOE大模型則先有路由器從8個專家組中選擇兩個，只有這兩個專家組參與了計算。而對比普通大模型，則需要所有參數都參加GPU計算。

所以MOE大模型要比同等級的普通大模型推理速度快四倍左右。

我們來看下Mistral MOE的實現，Mistral MOE是由mistral.ai發布的8*7B大模型[12]，下圖來自論文[12]，是其8*7B大模型的專家層的結構。

Mixtral 8x7B是一個稀疏混合專家（Sparse Mixture of Experts, SMoE）語言模型，它基于Mistral 7B的架構，但每一層都由8個前饋塊（即專家）組成。在處理每個令牌時，每層的一個路由網絡會選擇兩個專家來處理當前狀態并結合它們的輸出。雖然每個令牌只與兩個專家交互，但在每個時間步驟中選取的專家可以不同，因此每個令牌可以接觸到47B的參數，但在推理過程中只使用13B的活躍參數。

Mixtral在多項基準測試上展示了其卓越的性能，尤其是在數學、代碼生成和多語言理解方面。相比于Llama 2 70B和GPT-3.5，Mixtral在大多數評估指標上表現出類似或更優的性能。具體來說，Mixtral使用的活躍參數（13B）比Llama 2 70B（70B）少5倍，但在幾乎所有類別中的表現都更好或相當。

MOE大模型實現了增加參數量的同時，推理速度并不降低，是未來大模型的發展趨勢。

七、Tensor parallelize-張量并行

如果你有多卡GPU，可以采用張量并行進一步加速大模型的推理速度。

想象一下，你有一本非常厚的書，你想一次性復印整本書，但是你的復印機一次只能復印幾頁。這時，你可以把這本書分成幾個部分，每個部分分別復印，最后再把所有復印好的部分按順序拼接起來，這樣就完成了整本書的復印。

在張量并行中，我們要處理的大模型就像是那本厚書，而GPU則像是復印機。因為單個GPU無法一次處理整個大模型，我們就需要把模型（在這個例子中是權重張量）分成幾個部分，讓不同的GPU分別處理（相當于復印書的不同部分）。在處理輸入數據時，就像是把書的每一頁分別復印，然后再把復印好的各個部分拼接起來，形成完整的輸出結果。

這樣，通過分擔工作，多個GPU協同完成了一個本來單個GPU無法完成的大任務。這就是張量并行的工作方式，它讓我們能夠處理那些非常大的模型。

圖片來自文章[13]

張量并行技術用于將大模型分布式地部署在多個GPU上。以矩陣乘法來打個比方，當輸入張量與第一個權重張量進行矩陣相乘時，這個操作可以視作先將權重張量按列進行分割，接著將分割后的每列分別與輸入張量相乘，然后把這些乘積結果合并。這些合并后的輸出會被從GPU中導出，并聚合以形成最終的輸出結果，過程上圖，參考文章[13]。

八、推理框架推薦

在前文中，我們探討了幾種加速和優化技術，諸如Flash Attention、Page Attention、MOE以及張量并行技術。接下來，基于我們自身的實際操作和評估，我們將向您推薦一些當前表現較為出色的推理框架。

九、總結與展望

在本文中，我們深入探討了一系列旨在提升大模型推理速度的技術和方法，包括但不限于Flash Attention、Page Attention、MOE以及張量并行技術。通過在生產環境中批量部署專用大模型推理集群，我們成功地將包括70B規模模型在內的推理速度降低了50%，穩定地應用這些技術于生產環境，從而證明了這些優化方法的有效性和實用性。

隨著大型模型在各個領域的應用越來越廣泛，如何有效地提升推理速度、降低推理成本成為了一項挑戰。我們的實踐不僅展示了目前可用的一些加速技術，還基于我們的經驗，推薦了幾款經過評測表現優秀的大模型推理框架。這些建議旨在幫助讀者在面對眾多選擇時，能夠挑選出最適合自己需求的推理框架。

展望未來，隨著技術的不斷進步和新算法的不斷涌現，我們相信還會有更多的加速優化技術被開發出來，進一步推動大模型推理效率的提升。最后，我們也期待未來有機會深入探討和介紹更多提升大模型推理速度的新技術和方法。

參考資料

[1] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness(https://arxiv.org/abs/2205.14135)

[2] FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning(https://arxiv.org/abs/2307.08691)

[3] Efficient Memory Management for Large Language Model Serving with PagedAttention(https://arxiv.org/abs/2309.06180)

[4] mixtral-of-experts(https://mistral.ai/news/mixtral-of-experts/)

[5] Mixtral of Experts(https://arxiv.org/abs/2401.04088)

[6] MEDUSA: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads(https://arxiv.org/pdf/2401.10774.pdf)

[7] Scaling Laws for Neural Language Models(https://arxiv.org/pdf/2001.08361.pdf)

[8] Hyung Won Chung(OpenAI), Large Language Models (in 2023) , talked at Seoul National University

[9] Llama 2: Open Foundation and Fine-Tuned Chat Models(https://arxiv.org/abs/2307.09288)

[10] Attention Is All You Need(https://arxiv.org/pdf/1706.03762.pdf)

[11] https://blog.vllm.ai/2023/06/20/vllm.html

[12] https://arxiv.org/pdf/2401.04088.pdf

[13] https://huggingface.co/docs/text-generation-inference/en/conceptual/tensor_parallelism

[14] https://huggingface.co/docs/text-generation-inference/en/conceptual/flash_attention

[15] https://blog.vllm.ai/2023/06/20/vllm.html

本文轉載自??得物技術??，作者： linggong????