一、背景

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

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

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

以下為按時間順序業(yè)界的一些經(jīng)典大模型推理加速技術(shù)，本文試圖為讀者提供一個按時間發(fā)展順序的大模型加速方法綜述。

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

二、大模型發(fā)展面臨的挑戰(zhàn)

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

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

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

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

三、大模型推理加速的優(yōu)化方向

Llama2的模型結(jié)構(gòu)

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

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

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

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

從Llama 2模型的結(jié)構(gòu)分析中，我們可以總結(jié)出大模型在推理計算過程中表現(xiàn)出以下特點：

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

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

1. FlashAttention-Attention計算速度優(yōu)化

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

2. PageAttention-KV Cache顯存管理優(yōu)化

PageAttention[3]的目標是減少顯存碎片，基于PageAttention的VLLM系統(tǒng)能夠?qū)⒘餍械拇笮驼Z言模型（LLM）的吞吐量提高到10倍以上，同時保持耗時分布平穩(wěn)。

3. MOE-縮減推理時模型參數(shù)

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

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

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

四、FlashAttention-Attention算子計算優(yōu)化

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

我們先了解下GPU的內(nèi)存分層結(jié)構(gòu)，參考下圖，圖片來自論文FlashAttention-1[1]。

GPU的內(nèi)存層次結(jié)構(gòu)由三個主要部分組成：SRAM、HBM和DRAM，下面為A100GPU的參考配置。

SRAM（靜態(tài)隨機訪問存儲器）具有最快的訪問速度（19TB/s），但其容量相對較?。▋H20MB）。

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

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

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

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

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

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

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

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

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

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

五、PageAttention-顯存管理優(yōu)化

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

我們先了解下大模型在推理的情況下的內(nèi)存結(jié)構(gòu)分布，下圖來自論文[3]。

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

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

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

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

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

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

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

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

六、MOE-縮減推理時模型參數(shù)

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

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

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

我們來看下Mistral MOE的實現(xiàn)，Mistral MOE是由mistral.ai發(fā)布的8*7B大模型[12]，下圖來自論文[12]，是其8*7B大模型的專家層的結(jié)構(gòu)。

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

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

MOE大模型實現(xiàn)了增加參數(shù)量的同時，推理速度并不降低，是未來大模型的發(fā)展趨勢。

七、Tensor parallelize-張量并行

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

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

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

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

圖片來自文章[13]

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

八、推理框架推薦

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

九、總結(jié)與展望

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

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

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

參考資料

[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

本文轉(zhuǎn)載自??得物技術(shù)??，作者： linggong????