通過堆疊 Transformer 層來創建大型模型,可以提高準確性、實現少量學習功能,并且在各種語言任務中實現近乎人類的性能。這些基礎模型的訓練成本高昂,而且在推理過程中可能會占用大量內存和計算資源(這是一種重復性成本)。目前最熱門的大型語言模型 (LLM)可以達到數百億到數千億的參數規模,并且根據用例,可能需要處理長輸入(或上下文),這也會增加費用。
本文討論了大型語言模型(LLM)推理中最緊迫的挑戰及其實用解決方案。建議讀者對 Transformer 架構 和通用注意力機制有基本的了解。我們將在下一節重點討論掌握 LLM 推理復雜性的重要性。
了解 LLM 推理
大多數熱門的僅使用解碼器的 LLM (例如 GPT-3)都基于因果關系建模目標進行了預訓練,本質上是作為次詞預測器。這些 LLM 接受一系列標記作為輸入,并以自回歸方式生成后續標記,直到它們滿足停止標準(例如,要生成的標記數量的限制或停止詞列表),或直到它生成特殊的<end>
標記標志著生成過程的結束。此過程涉及兩個階段:預填充階段和解碼階段。
請注意,令牌是模型處理的語言的原子部分。一個令牌大約是四個英語字符。自然語言中的所有輸入在輸入模型之前都會轉換為令牌。
預填充階段或處理輸入
在預填充階段,LLM 處理輸入令牌以計算中間狀態(鍵和值),這些狀態用于生成“第一個”新令牌。每個新令牌都依賴于之前的所有令牌,但由于輸入的完整范圍已知,因此在高層次上,這是一個高度并行的矩陣矩陣運算。它實際上飽和了 GPU 利用率。
解碼相位或生成輸出
在解碼階段,LLM 一次自動回歸生成一個輸出令牌,直到滿足停止條件為止。每個連續輸出令牌都需要知道先前迭代的所有輸出狀態(鍵和值)。這就像矩陣向量運算,與預填充階段相比,它未充分利用 GPU 計算能力。數據(權重、鍵、值、激活)從內存傳輸到 GPU 的速度主宰延遲,而不是計算實際發生的速度。換言之,這是一個受內存限制的運算。
本文介紹的許多推理挑戰和相應解決方案都與此解碼階段的優化相關:高效注意力模塊、有效管理鍵和值等。
不同的 LLM 可能會使用不同的 tokenizer,因此,比較它們之間的輸出令牌可能并不簡單。在比較推理吞吐量時,即使兩個 LLM 的每秒輸出令牌相似,但如果它們使用不同的 tokenizer,它們可能不相同。這是因為相應的令牌可能表示不同數量的字符。
批處理
提高 GPU 利用率和有效吞吐量的最簡單方法是通過批處理。由于多個請求使用相同的模型,因此權重的內存成本分散。將更多的批量轉移到 GPU 以一次性處理,將利用更多的可用計算。
但是,批量大小只能增加到一定的限制,此時它們可能會導致內存溢出。為了更好地理解發生這種情況的原因,需要查看鍵值 (KV) 緩存和 LLM 內存要求。
傳統的批處理(也稱為靜態批處理)次優。這是因為,對于批量中的每個請求,LLM 可能會生成不同數量的完成令牌,隨后它們的執行時間也不同。因此,批量中的所有請求必須等待最長的請求完成,而生成長度的巨大差異可能會使這種情況更加嚴重。有一些方法可以緩解這種情況,例如動態批處理,稍后將討論這一點。
鍵值緩存
解碼階段的一個常見優化是 KV 緩存。解碼階段在每個時間步長生成單個令牌,但每個令牌都取決于之前所有令牌的鍵值張量(包括在預填充時計算的輸入令牌 KV 張量,以及在當前時間步長之前計算的任何新 KV 張量)。
為避免在每個時間步長為所有令牌重新計算所有這些張量,我們可以將其緩存在 GPU 顯存中。每次迭代計算新元時,只需將其添加到正在運行的緩存中,以便在下一次迭代中使用。在某些實現中,模型的每一層都有一個 KV 緩存。

LLM 內存需求
實際上,GPU LLM 內存需求的兩個主要貢獻者是模型權重和 KV 緩存。
- 模型權重:模型參數占用內存。例如,具有 70 億個參數(例如 Llama-2-7b),以 16 位精度(FP16 或 BF16)加載時,顯存大小約為 70 億 * 2 字節(FP16)= 14 GB。
- KV 緩存:內存中的自注意力張量被用作緩存,以避免重復計算。
在批處理時,批量中每個請求的 KV 緩存仍必須單獨分配,并且可能會占用大量內存。以下公式描述了適用于當今大多數常見 LLM 架構的 KV 緩存的大小。
每個令牌的 KV 緩存大小(字節)= 2 * (num_layers) * (num_heads * dim_head) * precision_in_bytes
第一個系數為 2 解釋了 K 和 V 矩陣。通常,(num_heads*dim_head) 的值與 Transformer 的 hidden_size (或模型維度 d_model)相同。這些模型屬性通常可在模型卡或相關配置文件中找到。
輸入序列中的每批輸入中的每個令牌都需要此內存大小。假設半精度,KV 緩存的總大小由以下公式給出。
以字節為單位的 KV 緩存總大小 = (batch_size) + (sequence_length) * 2 * (num_layers) * (hidden_size) * sizeof(FP16)
例如,對于 16 位精度的 Lama 2 7B 模型,批量大小為 1,則 KV 緩存的大小為 1*4096*2*32*4096*2 字節,即小于 2 GB.
高效管理此 KV 緩存是一項艱巨的任務。隨著批量大小和序列長度的線性增長,顯存需求可以快速擴展。因此,它限制了可以提供的吞吐量,并對長上下文輸入構成了挑戰。這正是本文介紹的幾項優化背后的原因。
使用模型并行擴展 LLM
減少模型權重在每個設備上的顯存占用的一種方法是將模型分配到多個 GPU 上。通過擴展顯存和計算占用,可以運行更大的模型或更大批量的輸入。模型并行化是在需要比單個設備上可用的內存更多的模型上進行訓練或推理的必要條件,并使訓練時間和推理措施(延遲或吞吐量)適合某些用例。有幾種方法可以根據模型權重的分割方式并行化模型。
請注意,數據并行性也是一種技術,通常與下面列出的其他技術在同一上下文中提及。在這種情況下,模型的權重在多個設備上復制,輸入的(全局)批量大小在每個設備上被分解為微批量。它通過處理更多的批量來減少總體執行時間。但是,它是一種訓練時間優化,在推理期間相關性較小。
管道并行
工作流并行涉及將模型(縱向)分片為多個數據塊,其中每個數據塊由在單獨設備上執行的層子集組成。圖 2a 展示了四路工作流并行,其中模型按順序進行分區,并在每臺設備上執行所有層的四分之一子集。一臺設備上的一組操作的輸出將傳遞給下一個設備,該設備將繼續執行后續的數據塊。和
分別指示設備 n 上的正向和反向傳遞。有效地將在每個設備上存儲模型權重的內存需求調至四分之一。
此方法的主要局限性在于,由于處理的順序性質,某些設備或層可能會在等待前幾層的輸出(激活、梯度)時保持空閑狀態,從而導致前向和反向通道效率低下或出現“管道氣泡”。在圖 2b 中,白色空白區域是具有樸素管道并行性的大型管道氣泡,其中設備處于空閑狀態且未得到充分利用。
如圖 2c 所示,微批處理可以在一定程度上減輕這種情況。輸入的全局批量大小被拆分成子批量,逐個進行處理,并在結束時累加梯度。請注意,和
分別指示設備上的正向和反向傳遞
使用微批
.此方法會縮小管道氣泡的大小,但并不能完全消除它們。

張量并行度
張量并行性涉及將模型的各個層(水平)分片為較小的獨立計算塊,這些計算塊可在不同的設備上執行。注意力塊和多層感知器 (MLP) 層是可以利用張量并行性的 Transformer 的主要組件。在多頭注意力塊中,可以將每個頭部或頭部組分配給不同的設備,以便可以獨立和并行計算。

圖 3a 展示了兩層 MLP 上的雙向張量并行性示例,其中每層均由一個圓框表示。在第一層中,權重矩陣分為
和
.計算
和
可以在同一批次上獨立執行(
是輸入的身份運算)
在兩臺不同的設備上運行。這實際上將在每臺設備上存儲權重所需的內存需求減半。
組合第二層中的輸出。
圖 3b 是自注意力層中雙向張量并行性的示例。多個注意力頭本質上是并行的,可以在設備之間分割。
序列并行度
Tensor 并行性存在局限性,因為它需要將層劃分為獨立的、可管理的塊。它不適用于 LayerNorm 和 Dropout 等操作,而是在張量并行組中進行復制。盡管 LayerNorm 和 Dropout 的計算成本較低,但它們確實需要大量內存來存儲(冗余的)激活。
如《減少大型 Transformer 模型中的激活重新計算》所示,這些操作獨立于輸入序列,并且可以沿著“序列維度”進行劃分,從而提高內存效率。這稱為序列并行性。

模型并行技術并非獨有,可以結合使用。它們可以幫助擴展和減少 LLM 的每個 GPU 顯存占用,但也有專門針對注意力模塊的優化技術。
優化注意力機制
擴展的點積注意力 (SDPA) 運算將查詢和鍵值對映射到輸出,詳情請參閱《您只需集中注意力》。
多頭注意力
作為對 SDPA 的增強,通過多次并行執行 Q、K 和 V 矩陣的不同學習投影來執行注意力層,使模型能夠共同處理來自不同位置的不同表征子空間的信息。這些子空間是獨立學習的,使模型能夠更深入地了解輸入中的不同位置。
如圖 5 所示,我們將多個并行注意力運算的輸出連接起來,并對其進行線性投影,以將其組合起來。每個并行注意力層稱為“頭部”,這種方法稱為多頭注意力 (MHA).
在原始作品中,每個注意力頭都在模型的簡化維度上操作(例如)。這使得計算成本與單頭注意力相似。

多查詢注意力
對 MHA 的一種推理優化,稱為多查詢注意力 (MQA),正如在 快速 Transformer 解碼 中所述,它在多個注意力頭之間共享鍵和值。查詢向量仍會像之前一樣進行多次投影。
雖然 MQA 中完成的計算量與 MHA 相同,但從內存讀取的數據量(鍵、值)只是之前的一小部分。受內存帶寬限制時,這可以提高計算利用率。它還減少了內存中的 KV 緩存的大小,為更大的批量大小提供了空間。
減少鍵值頭可能會導致準確性降低。此外,要利用這種優化的模型進行推理,需要進行訓練(或至少經過微調,訓練量約為原始訓練量的 5%),同時需要啟用 MQA。
分組查詢注意力
分組查詢注意力(GQA)通過將關鍵幀和值投射到幾組查詢頭(圖 6),在 MHA(多頭注意力)和 MQA(多查詢注意力)之間取得平衡。在每個組中,其行為類似于多查詢注意力。
圖 6 顯示了多頭注意力具有多個鍵值頭(左)。分組查詢注意力(中心)的鍵值頭超過一個,但少于查詢頭的數量,這是內存需求和模型質量之間的平衡。多查詢注意力(右)具有單個鍵值頭以幫助節省內存。

最初使用 MHA 訓練的模型,可以使用原始訓練計算的一小部分通過 GQA 進行“上訓練”。它們的質量接近 MHA,同時保持接近 MQA 的計算效率。Lama 2 70B 是一個利用 GQA 的模型示例。
MQA 和 GQA 等優化通過減少存儲的鍵和值頭的數量來幫助減少 KV 緩存所需的內存。管理此 KV 緩存的方式可能仍然效率低下。與優化注意力模塊本身不同,下一節將介紹一種更高效的 KV 緩存管理技術。
閃光注意力
優化注意力機制的另一種方法是修改某些計算的順序,以更好地利用 GPU 的內存層次結構。神經網絡通常以層的形式進行描述,大多數實現也以這種方式進行布局,按順序對輸入數據進行一種計算。這并不總是能實現最佳性能,因為對已引入更高、性能更高的內存層次結構的值進行更多計算會很有好處。
在實際計算期間將多個層融合在一起,可以最大限度地減少 GPU 從內存讀取和寫入內存的次數,并將需要相同數據的計算分組,即使它們是神經網絡中不同層的一部分。
一種非常流行的技術融合是 FlashAttention,這是一種 I/O 感知的精確注意力算法,詳見FlashAttention:通過 I/O 感知實現快速且節省內存的精確注意力。所謂精確注意力,意味著它在數學上與標準的多頭注意力(提供用于多查詢和分組查詢注意力的變體)相同,因此可以無縫集成到現有的模型架構中,甚至適用于未經修改的已訓練模型。
I/O 感知表示在將運算融合在一起時,它會考慮到之前討論的一些內存移動成本。特別是,FlashAttention 使用“平鋪”來完全計算并一次性寫出一小部分最終矩陣,而不是按步驟對整個矩陣進行部分計算,寫出中間值。
圖 7 顯示了 40 GB GPU 上平鋪的 FlashAttention 計算模式和內存層次結構。右側圖表顯示了因融合和重新排序 Attention 機制的不同組件而導致的相對加速。

通過分頁有效管理 KV 緩存
由于輸入大小不可預測,因此 KV 緩存有時會靜態地“過度置備”,以考慮到盡可能大的輸入(支持的序列長度)。例如,如果模型的支持最大序列長度為 2048,那么無論請求中的輸入大小和生成的輸出大小如何,都將在內存中保留大小為 2048 的空間。此空間可能會被連續分配,而且其中大部分仍然未使用,從而導致內存浪費或碎片化。在請求的整個生命周期中,此保留空間被占用。

受操作系統分頁機制啟發,PagedAttention 算法允許在內存中非連續空間存儲連續的鍵和值。它將每個請求的 KV 緩存分割成代表固定數量令牌的塊,并且這些令牌可以以非連續的方式進行存儲。
在注意力計算期間,系統會根據需要使用記錄帳戶的塊表來獲取這些塊。生成新令牌后,系統會進行新的塊分配。這些塊的大小是固定的,可以消除因挑戰(例如需要不同分配的不同請求)而產生的低效問題。這顯著限制了內存浪費,從而實現了更大的批量大小(從而提高吞吐量)。
模型優化技術
到目前為止,我們已經討論了 LLM 消耗內存的不同方式、在多個不同 GPU 之間分配內存的一些方式,以及優化注意力機制和 KV 緩存。還有一些模型優化技術,可以通過修改模型權重本身來減少每個 GPU 上的內存使用。GPU 還有專用硬件來加速對這些修改值的操作,從而提高模型的速度。
量化
量化是降低模型權重和激活精度的過程。大多數模型都以 32 或 16 位精度進行訓練,其中每個參數和激活元素占用 32 或 16 位內存(單精度浮點)。但是,大多數深度學習模型可以用每個值 8 位甚至更少的位來有效表示。
圖 9 顯示了一種可能的量化方法前后的值分布。在這種情況下,舍入會損失一些精度,裁剪會損失一些動態范圍,從而以更小的格式表示值。

降低模型的精度可以產生一些好處。如果模型占用的內存空間較少,您可以在相同數量的硬件上擬合更大的模型。量化還意味著您可以在相同數量的帶寬上傳輸更多參數,這有助于加速帶寬受限的模型。
LLM 有許多不同的量化技術,涉及降低激活值、權重或兩者的精度。量化權重更直接,因為它們在訓練后是固定的。但是,由于激活值保持在更高的精度,因此這可能會在表格上保留一些性能。GPU 沒有用于乘以 INT8 和 FP16 數字的專用硬件,因此必須將權重轉換回更高的精度,以用于實際操作。
也可以量化 Transformer 塊和網絡層的激活、輸入,但這也有其自身的挑戰。激活向量通常包含離群值,從而有效地增加其動態范圍,使得以低于權重的精度表示這些值更具挑戰性。
一種方法是,通過在模型中傳遞代表性數據集,并選擇以比其他模型更高的精度表示某些激活 (LLM.int8 ()),找出這些異常值可能出現的位置。另一種方法是借用易于量化的權重動態范圍,并在激活中重復使用該范圍。
稀疏
與量化類似,研究表明,許多深度學習模型在修剪或將接近 0 的某些值替換為 0 本身方面都很可靠。稀疏矩陣是許多元素為 0 的矩陣。它們可以用凝聚形式表示,與完整的密集矩陣相比,它占用的空間更小。

尤其是 GPU,對于某些類型的結構化稀疏每 4 個值中就有 2 個由 0 表示。稀疏表示還可以與量化相結合,以實現更高的執行速度。尋找以稀疏格式表示大型語言模型的最佳方法仍然是一個活躍的研究領域,并為未來提高推理速度提供了一個大有希望的方向。
蒸
縮小模型大小的另一種方法是通過名為蒸此過程涉及訓練一個較小的模型(稱為學生),以模擬較大模型(教師)的行為。
成功的模型壓縮示例包括DistilBERT,它將 BERT 模型的大小壓縮了 40%,同時保留了 97% 的語言理解能力,并且速度提升了 60%。
盡管在大型語言模型(LLM)中的知識蒸餾是一個活躍的研究領域,但更多相關信息可以在提取神經網絡中的知識中找到:
- 學生網絡經過訓練,可以反映更大的教師網絡的性能,使用測量其輸出之間差異的損失函數。除了可能包括將學生的輸出與真值標簽相匹配的原始損失函數之外,還可以實現這一目標。
- 匹配的教師輸出可以是最后一層(稱為邏輯)或中間層激活。
圖 11 展示了知識蒸的通用框架。教師的邏輯是學生使用蒸損失進行優化的軟目標。其他蒸方法可能會使用其他損失度量來“提煉”教師的知識。

另一種訓練方法是使用由教師模型合成的數據對學生大語言模型 (LLM) 進行監督式訓練,這在人類注釋稀缺或不可用的情況下特別有用。逐步提煉!除了作為真實標簽之外,還可以從教師 LLM 中提取推理過程。這些推理過程作為中間步驟,可以以數據高效的方式訓練較小的學生 LLM。
需要注意的是,目前許多先進的 LLM 都有嚴格的許可證,禁止使用其輸出來訓練其他 LLM,因此很難找到合適的教師模型。
模型服務技術
模型執行通常受內存帶寬限制,尤其是權重中的帶寬限制。即使應用了前面介紹的所有模型優化,它仍然很可能受內存限制。因此,您希望在加載模型權重時盡可能多地處理這些權重。換言之,嘗試并行執行操作。可以采取兩種方法:
- 動態批處理:涉及同時執行多個不同的請求。
- 推測推理:涉及并行執行序列中的多個異步步驟以節省時間。
動態批處理
LLM 具有一些獨特的執行特性,在實踐中可能難以有效地批處理請求。單個模型可同時用于各種看起來非常不同的任務。從聊天機器人中的簡單問答響應到文檔摘要或生成長代碼塊,工作負載具有高度動態化的特點,輸出大小變化了幾個數量級。
這種通用性會使批處理請求并有效地并行執行這些請求(這是服務神經網絡的常見優化)具有挑戰性。這可能會導致一些請求比其他請求更早完成。
為了管理這些動態負載,許多 LLM 服務解決方案包括一種名為連續或動態批處理的優化調度技術。這利用了一個事實,即 LLM 的整個文本生成過程可以分解為模型上的多次執行迭代。
在運行中批處理時,服務器運行時無需等待整個批處理完成再轉到下一組請求,而是會立即從批量中移除已完成的序列。然后,當其他請求仍處于運行中時,服務器會開始執行新請求。因此,運行中批處理可以在實際用例中大幅提高 GPU 的整體利用率。
推測推理
推測推理也稱為推測采樣、輔助生成或塊級并行解碼,是并行執行 LLM 的另一種方式。通常,GPT 風格的大型語言模型是按令牌生成文本標記的自回歸模型。
生成的每個令牌都依賴于之前提供上下文的所有令牌。這意味著,在常規執行中,不可能并行從同一序列生成多個令牌 – 您必須等待生成第 n 個令牌,然后才能生成 n+1.
圖 12 展示了一個推理示例,其中草稿模型暫時預測多個并行驗證或拒絕的未來步驟。在這種情況下,草稿中的前兩個預測標記被接受,而最后一個標記被拒絕并刪除,然后再繼續生成。

預測采樣提供了一種變通方法。此方法的基本思路是使用一些“更便宜”的過程來生成包含多個令牌的草稿延續。然后,并行執行多個步驟的主“驗證”模型,在需要執行步驟時將廉價草稿用作“預測”上下文。
如果驗證模型生成的令牌與草稿相同,則您知道接受這些令牌用于輸出。否則,您可以在第一個不匹配令牌之后丟棄所有內容,然后使用新草稿重復此過程。
如何生成草稿令牌有許多不同的選項,每種選項都有不同的權衡。您可以訓練多個模型,或在單個預訓練模型上微調多個頭部,以預測未來的多個步驟的令牌。或者,您可以使用小型模型作為草稿模型,使用更大、功能更強大的模型作為驗證器。
結束語
本文將概述許多最熱門的解決方案,以幫助高效優化和服務 LLM,無論是在數據中心還是在 PC 的邊緣。其中許多技術都經過優化,并通過 NVIDIA TensorRT-LLM 這是一個開源庫,由 TensorRT 深度學習編譯器、優化的內核、預處理和后處理步驟以及多 GPU/多節點通信基元組成,可在 NVIDIA GPU 上實現突破性性能。如需了解詳情,請參閱 借助 NVIDIA TensorRT-LLM 優化大型語言模型的推理(現已公開發布)。
NVIDIA TensorRT-LLM 現在由 NVIDIA Triton 推理服務器 支持,使企業能夠跨不同的 AI 框架、硬件加速器和部署模式同時為多個 AI 模型提供服務,并實現峰值吞吐量和更低延遲。
TensorRT-LLM 還支持 NVIDIA NeMo,它提供了一個端到端的云原生企業框架,供開發者構建、自定義和部署具有數十億參數的生成式 AI 模型。欲了解更多信息,請開始使用 NeMo。
?