語言模型推論的光速極限

在開發 calm 的過程中,這是一個從頭開始打造、精簡且快速的transformer語言模型 CUDA 推論實作,一個關鍵的考量是確立推論過程的光速極限,並衡量相對於該極限的進展。在本文中,我們將討論這個理論上限及其意涵。              --Arseny Kapoulkine--

如果您有興趣瞭解更多推導過程和一些圖表,這個 Python notebook 做了相同的建模。
https://github.com/zeux/calm

推論機制

當語言模型 (特指 decoder-only 的文字 transformer 模型,但本文其餘部分將統稱為語言模型) 生成字詞時,它是一次一個字詞地生成;語言模型可以被理解為一個函數,它接收一個字詞作為輸入,並產生一個包含詞彙表中所有字詞 (通常有 5 萬到 25 萬個字詞,每個字詞由幾個字母組成) 機率的陣列。接著,程式根據這些機率從所有字詞中取樣,產生下一個字詞,然後重複這個過程。這意味著在生成一段文字序列時,沒有辦法平行化 ─ 生成過程可以被視為一次一個字詞。

語言模型在處理一個字詞時,大致上會進行兩種運算:矩陣向量乘法,其中一個大矩陣 (如 8192x8192) 乘以一個向量產生另一個向量;以及注意力 (attention) 的計算。在生成過程中,模型不僅可以看到當前字詞的狀態,還可以看到序列中所有先前字詞的內部狀態 ─ 包括使用者在提示中寫的字詞,以及模型本身生成的字詞。這些狀態儲存在KV 快取(key-value cache) 的結構中,本質上是每個先前位置的一組 key 和 value 向量。注意力機制為當前字詞生成一個 query 向量,計算它與所有先前位置的所有 key 向量之間的內積,然後將結果標準化,並透過對所有先前位置的所有 value 向量進行加權求和來計算一個 value 向量,其中內積作為權重。

現在,矩陣向量乘法和注意力計算有一個重要的共通點:對於從矩陣或 KV 快取讀取的每個元素,我們只需要做非常少量的浮點運算。矩陣向量乘法對每個矩陣元素做一次乘法和加法 (2 個 FLOP);注意力計算對每個 key 元素做一次乘法和加法來計算內積,對每個 value 元素做一次乘法和加法來計算加權和。

現代 CPU 和 GPU 的 ALU 運算(乘法、加法)速度遠高於從記憶體讀取輸入的速度。例如:

- AMD Ryzen 7950X 有 67 GB/s 的記憶體頻寬和 2735 GFLOPS,FLOP與byte之比為 40:1。

- NVIDIA GeForce RTX 4090 有 1008 GB/s 的記憶體頻寬和 83 TFLOPS,FLOP與byte之比為 82:1。 

- NVIDIA H100 SXM(一張資料中心用的卡)有 3350 GB/s 的記憶體頻寬和 67 TFLOPS,FLOP與byte之比似乎只有 20:1;然而,對於類似矩陣乘法的問題,tensor core 在沒有稀疏性的情況下可提供約 494 TFLOPS,使得 FLOP 與 byte 之比達到 147:1。

對於 FP16 或 FP8 等更小的浮點數,情況會變得更糟:H100 的 tensor core 在密集 FP8 矩陣上的理論吞吐量為 1979 TFLOPS,將 FLOP 與 byte 之比提高到 590:1。不用說,在任何這些配置下,無論是否使用 tensor core,以及無論使用何種浮點格式,ALU 都是非常充足的。

因此,任何只需要對每個元素做兩次運算的問題,都一定會受到頻寬的限制,我們應該能夠根據模型配置、KV快取的大小和可用頻寬來估計推論過程所需的最小時間。

Mistral 的光速極限

在不過於深入討論確切公式和矩陣的情況下,讓我們看看像 Mistral 7B 這樣的模型,它有 72 億個參數 (因此所有矩陣元素的總數是 72 億)。

參數的組成如下:

- 4096 * 32000 = 1.31 億個嵌入矩陣的參數;這個矩陣不會用在矩陣向量乘法中,因為對每個字詞只會讀取矩陣的一行,所以我們不會將其納入頻寬計算。

- 32 * (4096 * (128 * 32 + 128 * 8 * 2) + 4096 * 128 * 32) = 13.42 億個用於計算注意力相關向量的參數。 

- 32 * (4096 * 14336 * 3) = 56.37 億個透過前饋網路轉換隱藏狀態的參數。

- 4096 * 32000 = 1.31 億個將隱藏狀態轉換為字詞機率的參數;這個矩陣會用在矩陣乘法中,與嵌入矩陣不同。

這加起來大約有 71.11 億個在矩陣乘法中使用的「活躍」參數。如果模型對矩陣元素使用 FP16,我們最終必須為每個字詞讀取約 14.2 GB 的資料。此外,雖然每個矩陣在為下一個字詞執行推論時都會再次使用,但快取的大小通常在幾十 MB 左右,因此我們可以假設這個過程無法比記憶體頻寬更快,因為權重在推論之間不會留在快取中。

這涵蓋了矩陣運算;注意力計算需要讀取直到當前字詞的 KV 快取,因此讀取的資料量取決於模型在生成新字詞時看到的字詞數量 ─ 包括系統提示 (通常對使用者隱藏)、使用者提示、先前的模型輸出,並且可以包括較長聊天階段的多個使用者提示。

對於 Mistral,KV 快取為每層的每個 key 儲存 8 個 128 維的向量,為每層的每個 value 儲存 8 個 128 維的向量,總共是每個字詞 32 * 128 * 8 * 2 = 65K 個元素;如果 KV 快取使用 FP16 儲存單個元素,那麼對於第 P 個字詞,我們需要讀取 P * 130 KB 的記憶體 ─ 例如,第 1000 個字詞需要從 KV 快取讀取 130 MB 的資料。

有了這些數字,現在很容易計算推論所需的最小時間。例如,在 NVIDIA RTX 4090 (1008 GB/s) 上,讀取 14.2 GB 大約需要 14.1 毫秒,因此我們可以預期對位置數較低的字詞,每個字詞大約需要 14.1 毫秒 (KV 快取的影響可以忽略不計)。如果我們使用 8 位元的權重,我們需要讀取 7.1 GB,這大約需要 7.0 毫秒。這些是下限 ─ 它們代表了理論上每個字詞可能的最小時間。

理論上限有用嗎?

我們做了一堆數學計算,得到了一些告訴我們推論速度無法超過給定閾值的數字 ─ 這有用嗎?讓我們看看為什麼它可能有用的幾個原因。

要真正達到該時間,你需要高品質的軟體實作,以及能夠達到理論峰值頻寬的硬體。這意味著,如果給定的實作遠低於最佳數字,就值得進行調查:效能可能遺留在軟體或硬體方面。例如,在 RTX 4090 上,使用 16 位元權重時,calm 可以達到每個字詞約 15.4 毫秒;使用 8 位元權重時,可以達到每個字詞約 7.8 毫秒 ─ 這大約是理論上可能效能的 90%。在使用 CPU 推論的 Apple M2 Air 上,calm 和 llama.cpp 都只達到了理論 100 GB/s 頻寬的約 65%,這表明引用的峰值頻寬只有在 iGPU 的幫助下才能完全利用。

頻寬與每個元素使用的位元組數呈線性關係;這意味著我們既可以從更小的權重格式 (量化) 估計理論效益,也可以透過將實際效能與理論極限進行比較來驗證實作的品質。例如,在 RTX 4090 上,使用 16 位元權重時,llama.cpp 可以達到每個字詞約 17.1 毫秒 (峰值的 82%);使用 8.5 位元權重時,可以達到每個字詞約 10.3 毫秒 (峰值的 71%);使用 4.5 位元權重時,可以達到每個字詞約 6.7 毫秒 (峰值的 58%),這表明對更小的格式還有很大的優化空間。

除了提供解碼時間的下限之外,上面的建模還表明,推論過程嚴重低估了 ALU 單元的利用率。要解決這個問題,需要改變 FLOP 與 byte 的平衡;投機解碼等技術試圖幫助解決這個問題,但對於多使用者的例子,我們可以注意到,當同時處理多個使用者請求時,我們可以在同一個矩陣上同時執行多個矩陣向量乘法(也就是矩陣矩陣乘法!)─ 矩陣矩陣乘法的最佳實作會在矩陣足夠大時變成ALU bound。這就是為什麼這種 ALU 與 byte 的不平衡對於生產環境的推論系統而言不是一個關鍵問題 ─ 當你要求 ChatGPT 幫助一項任務時,你的請求會與同一 GPU 上的許多其他請求同時進行評估,頻寬會被更有效地利用。至關重要的是,請求批次處理通常無法幫助減少 KV 快取的頻寬需求(除非這些請求共享非常大的前綴),因為 KV 快取的大小和頻寬會隨著請求數量增加,而權重矩陣保持不變。

像 Mixtral 這樣的專家混合模型有略微不同的縮放特性:批次處理最初只會增加所需的頻寬,但一旦專家利用率變得顯著,推論就會越來越受 ALU 的限制。

最後,如果無法進行批次處理,頻寬就成為預期推論效能的關鍵估算指標,它在不同的模型變體/設備類型或架構中保持不變,你可以用它來決定需要使用的硬體。例如,NVIDIA RTX 4080 有 716 GB/s 的頻寬,所以你可以預期它執行語言模型推論的速度大約是 RTX 4090 的 0.7 倍 ─ 這可能與其他工作負載(如遊戲、光線追蹤或其他類型神經網路的推論)的相對效能不同!

結論

對於像這樣計算量和記憶體訪問量已知的問題,將理論光速建模作為基準非常重要,因為它有助於驗證實作的品質並預測架構變化的影響。

理想情況下,你的推論實作應該仔細計算實現的有效頻寬,你應該在效能分析時將其作為主要指引 ─ 因為你知道這個值的極限!不過,一定要仔細計算 ─ calm 曾多次出現架構上的特殊情況,導致計算出的頻寬略有偏差:)

附錄:群組查詢注意力 (Group Query Attention)

Mistral-7B 是一個非常平衡的模型;在上面的所有計算中,KV 快取似乎並不是成本結構中不可或缺的部分。其中一個原因是相對較短的上下文(Mistral-7B 使用了視窗注意力,將頻寬消耗限制在 4096 個字詞的視窗中),但另一個可能更重要的原因,是使用了群組查詢注意力(Group Query Attention, GQA)。

在群組查詢注意力中(4倍比率),為了產生4個內積,我們不是使用4個查詢向量並與4個對應的鍵向量計算內積,而是取一個鍵向量但使用4個查詢向量,並執行4次內積運算。這使我們能夠減少KV快取的大小和所需的頻寬 ─ 我們不再是從KV快取中讀取每個元素並只對其進行一次乘加運算,而是現在進行4次運算,這在一定程度上重新平衡了ALU與頻寬的比率,對我們有利。

這對KV快取的記憶體大小也至關重要,但對於如此短的上下文可能並不明顯:4096個字詞的上下文在Mistral中佔用0.5 GiB,但沒有GQA的類似模型(如Llama 7B)「只」需要2 GiB。讓我們看看最近一個沒有使用GQA的模型,Cohere的Command-R。

該模型本身有約350億個參數,因此在每個權重16位元的情況下,我們在推論期間需要為每個字詞讀取70 GB的權重。對於每個字詞,它需要在KV快取中儲存40 * 128 * 64 * 2 = 655K個元素,在每個元素16位元的情況下,每個字詞是1.3 MB。

因此,4096個字詞的上下文將佔用約5.3 GB;與約70 GB的權重相比,這已經相當可觀了。然而,如果考慮到Cohere的模型號稱有20萬個字詞的上下文視窗,情況會變得更加可怕 ─ 要計算20萬個上下文視窗的最後一個字詞,你需要讀取260 GB!(讓我們暫時忽略你還需要260 GB的顯存來儲存它)

在典型的「生產」(仍然是單一使用者)環境中,情況會進一步改變。權重通常會使用4位元量化(通常實作為每個權重約4.5位元),而KV快取可能會使用8位元(FP8)值。如果我們「保守地」假設10萬個字詞的上下文(是宣稱最大值的一半),這將給我們帶來約19.7 GB的模型權重和約65 GB的KV快取,而要計算最後一個字詞,我們需要從記憶體中讀取所有這些資料。突然間,注意力計算從微不足道變成了佔用時間的約75%,假設兩者都以峰值頻寬執行!

儘管10萬個字詞的上下文可能看起來有點極端,但在多使用者的環境中,這也是預期工作負載的公平表示。批次處理允許我們使矩陣乘法變成ALU bound,並為每批值(即每64+個使用者請求)讀取一次模型權重,但每個使用者請求通常都有自己的KV快取,因此注意力計算仍然是頻寬bound ─ 並且需要大量記憶體才能將所有使用者的請求放在單個節點上!

如果這些模型使用了4倍GQA,KV快取所需的大小和頻寬就會小4倍;雖然對於數萬個字詞的上下文來說仍然很可觀,但會更容易管理。對於Cohere的預期使用情境,使用GQA可能會有一些品質下降 ─ 看看技術報告會很有趣,因為它可能包含相關的消融研究,但純粹從成本/效能的角度來看,每個基於transformer的語言模型都需要評估GQA,因為其優勢太顯著了,不容忽視。

個人見解:

在生成文字時,語言模型是逐字逐句依序產生,沒有辦法平行化處理。模型在處理每個字詞時,主要進行矩陣向量乘法和注意力(attention)計算。現代 CPU 和 GPU 的算術運算能力遠高於記憶體讀取速度,以至於推論過程的效能幾乎完全取決於記憶體頻寬。

藉由分析模型參數的組成和注意力機制的快取需求,我們可以預估推論每個字詞所需的最小時間。這個理論下限對於評估軟硬體效能非常有用:

1. 如果實際效能遠低於理論值,代表軟硬體可能有優化空間。
2. 不同精度(如FP16、INT8)的理論效能可以預估量化的潛在效益。
3. 即使硬體架構不同,記憶體頻寬仍是一致的效能指標。

我認為作者提出了一個實用且具啟發性的效能模型。不只可以用來診斷效能瓶頸,也能輔助軟硬體設計的取捨。唯一的缺點是沒有涵蓋多使用者的情境,不過文中也簡單討論了批次處理(batching)的影響。

此外,作者特別強調了群組查詢注意力(Group Query Attention)的重要性。藉由增加每個注意力快取元素的運算量,GQA 顯著降低了快取占用的記憶體容量和頻寬。以 Cohere 的 Command-R 模型為例,如果使用 4x GQA,快取所需的記憶體和頻寬可以減少 4 倍之多。

我認為速度極限分析是效能工程不可或缺的一環。透過將實際效能與理論極限比較,我們可以系統性地找出優化方向,而不是憑直覺或運氣。當然,模型品質的權衡也不能忽視。畢竟再怎麼快,生成愚蠢的回覆也是沒有意義的。在效能與品質之間取得平衡,才是打造實用語言模型的不二法門。

不過,我也想補充幾點:

1. 文中假設讀取權重是效能的主要限制,但對某些硬體(如AI加速晶片)來說,這可能並不完全成立。晶片內部通常有專門的高速記憶體來儲存權重,因此讀取速度可能不是最關鍵的因素。

2. 推論速度只是語言模型實用性的一個面向。即使透過種種優化來逼近理論極限,品質和穩定性的問題仍然存在。生成不連貫或有害的內容,對使用者體驗的傷害可能大過延遲幾毫秒。品質和效能需要一起考量。

3. 文章聚焦於單一使用者的情境,但實際應用常常需要同時服務大量使用者。這時批次處理和記憶體效率的權衡會變得更加複雜。找出最佳的批次大小和排程策略,是一個值得深入研究的課題。

作者提出了一個實用的效能分析框架,揭示了當前語言模型推論的核心挑戰。這些洞見不僅有助於改進現有系統,也為未來的架構設計提供了方向。期待看到更多在效能、品質、應用場景等多方面尋求平衡的創新方案。