LLM 推理最佳化探微 (1) :Transformer 解碼器的推理過程詳解

Baihai_IDP發表於2024-01-31

編者按:隨著 LLM 賦能越來越多需要實時決策和響應的應用場景,以及使用者體驗不佳、成本過高、資源受限等問題的出現,大模型高效推理已成為一個重要的研究課題。為此,Baihai IDP 推出 Pierre Lienhart 的系列文章,從多個維度全面剖析 Transformer 大語言模型的推理過程,以期幫助讀者對這個技術難點建立系統的理解,並在實踐中做出正確的模型服務部署決策。

本文是該系列文章的第一篇,作者的核心觀點是:透徹理解 Transformer 解碼器的推理過程,是實現大語言模型高效能服務的基礎。

作者透過解析文字生成的流程,明確了啟動階段和生成階段的概念,並指出了鍵值快取在其中起到的關鍵作用,為後續最佳化方法做好了理論鋪墊。

作者 | Pierre Lienhart

編譯 | 嶽揚

01 Introduction

在本系列博文中,我將帶領各位讀者深入探討 LLM 推理的各方面知識以及實踐中可能面臨的挑戰。 本系列博文所說的 LLM 推理是指使用僅包含解碼器的 Transformer 模型生成詞元(tokens),因為大多數挑戰及其相關的補救方法都來自於該特定架構和使用場景。 雖然重點在解碼器模型,但我認為該系列部落格中的一些較好的見解也可以用於理解和最佳化編碼器模型的推理過程。

我假定各位讀者已經對 Transformer 架構和著 名的《Attention Is All You Need》論文[1]中介紹的縮放點積注意力(SDPA)機制有了基本的瞭解。不過,閱讀本文無需深入理解注意力機制背後的原理。

希望在本系列部落格結束時,各位讀者能夠理解與 LLM 推理相關的術語,如鍵-值 (KV) 快取、記憶體頻寬約束(memory-bandwidth bound)等,能夠理解各種推理最佳化(模型量化(quantization)、核心融合(fused kernels)、模型架構修改(model architecture modifications)等)和配置(batch size、使用哪種 GPU 等)技術,最後能夠將它們與延遲、吞吐量和成本等關鍵效能指標聯絡起來。

期望各位讀者透過理解文章中的內容,形成一種具備洞察力的思維模式,以便在處理 LLM 服務的配置和最佳化時,能夠做出基於充分資訊的、迅速而明智的決策。和其他系列文章一樣,我希望本系列文章為初次部署 LLM 到服務端的讀者,提供他們希望獲得的有關資訊和指導。

現在,讓我來介紹一下本系列的大綱。

首先,需要理解使用 Transformer 解碼器(Transformer decoder)生成 tokens 需要以下兩個步驟。這兩個步驟分別是處理提示語步驟和多個自迴歸步驟。兩個步驟在硬體利用上有著截然不同的特徵,我們將在整個系列中詳細探討這種區別。

然後,我們將介紹對自迴歸步驟進行的第一個非常常見的最佳化,即 KV 快取(KV caching)。重點介紹 KV 快取是不可避免的,因為它作為一種關鍵的輸入影響著整個自迴歸階段。正如我們將看到的那樣, KV 快取並非免費的午餐,它也會引發一系列問題。 在接下來的文章中,我們將更深入地探討這些挑戰及其緩解方法,並專門針對這一主題進行深入討論。

當對 KV 快取瞭如指掌之後,我們可以更深入地瞭解執行 Transformers 進行推理時,Transformer 模型如何(未)充分利用硬體資源。在這一階段,我們必須引入算術強度(arithmetic intensity)這一關鍵概念,這是一個名為 "屋頂線模型"(roofline model)的 mental model (譯者注:"mental model" 意指人們在思考和理解複雜概念時構建的一種抽象思維模型。"屋頂線模型" 是一種特定的mental model,用於評估計算密集型任務的效能極限。這種模型通常以圖形的形式呈現,幫助人們直觀地理解任務在硬體效能方面的限制,並在最佳化過程中提供指導。),並將其與峰值 FLOPS、記憶體頻寬等關鍵硬體特徵以及延遲、吞吐量和成本等關鍵效能指標聯絡起來。然後,我們將這些知識應用於 Transformers 的推理,彙總關鍵見解,以瞭解如何更好地利用硬體並改善效能指標。透過對這一階段基礎知識的理解,我們將更好地把握效能最佳化過程中各種策略的真實動機,為實際應用提供更有深度的指導。

模型量化(Quantization)一直是去年最熱門的最佳化策略之一,能夠為效能帶來重大提升。雖然模型量化本身就值得寫一系列文章,但我將僅用一篇文章介紹,以期為讀者打牢堅實的基礎,明確模型量化演算法在哪些方面能夠有幫助,哪些方面沒有幫助。

最後,我們需要介紹目前的這些 LLM 服務框架是如何工作的。 遺憾的是,要想在 LLM 推理過程獲得最 佳效能,僅最佳化模型本身是不夠的。 事實上,模型伺服器透過有效地管理傳入的請求和硬體資源,在確保較好的 end-to-end 效能方面確實發揮著關鍵作用。我希望最後這篇文章能為讀者提供有用的見解,幫助讀者更好地部署LLM。

以下是本系列部落格內容計劃:

  • Introduction

  • The two-step process behind LLMs’ responses

  • KV caching unveiled

  • KV caching: A deeper look

  • Arithmetic intensity: Are you compute or memory-bandwidth bound?

  • Arithmetic intensity (and memory) is all you need

  • Shrink all the things! A guided tour of LLM quantization

  • Why you can’t just serve LLMs using a good old model server?

閒話少說,讓我們直接開始吧!

02 LLM 做出回答前歷經的兩個步驟

此小節將為各位讀者提供一個熱身的機會,我們先來回顧 Transformer 架構和使用基於 Transformer 的解碼器生成文字的基礎知識。對於本系列部落格的理解和學習,建立一致的專業用詞體系是非常重要的。我將使用粗體標出我個人喜歡的專業術語。然後我們來一起了解文字生成的兩個階段:啟動階段和生成(或解碼)階段。

首先,我們來複習一下 Transformer。為簡單起見,我們假設每次只處理單個序列(sequence)(即 batch size 為 1)。在下圖中,我描繪了基於原始 Transformer 架構的解碼器(圖 1)的主要層次結構,該解碼器(decoder)用於從輸入的 tokens 序列中生成輸出tokens。

LLM 推理最佳化探微 (1) :Transformer 解碼器的推理過程詳解

圖 1 -  Transformer 解碼器的模型輪廓圖

請注意,解碼器本身並不輸出 tokens,而是輸出 logits(數量與詞彙表的大小相同)(譯者注:logits 是一個數值向量,其維度等於詞彙表的大小,表示每個 token 的可能性分數。)。順便說一下,輸出 logits 的最後一層通常被稱為 language model head 或 LM head 。在生成文字時,透過 logits 提取 tokens 的過程是透過一種被稱為搜尋策略(search strategy)、生成策略(generation strategy)或解碼策略(decoding strategy)的啟發式方法完成的。常見的解碼策略包括:

  • 貪婪解碼(Greedy decoding) ,簡單來說,就是選取具有最大 logit 的token,在選擇最終的 token 之前,也可以對 logits 進行一些變換(如重複懲罰(repetition penalty))。

  • 抽樣解碼(Sampling decoding) ,模型的 logits 被看作是一個多項分佈,採用該分佈進行抽樣。換句話說,就是透過抽樣從詞彙表(vocabulary)中選擇一個token。我們從中進行抽樣的分佈可以首先透過簡單的一些轉換操作(如temperature scaling、top-k和top-p)來調整,這是最為人熟知的。

  • 還有更復雜的啟發式方法,比如 束搜尋(beam search)、對比解碼(contrastive decoding) [2]等。

為了簡單起見,我們將假定解碼策略是模型的一部分(圖2)。這種 mental model 實際上在部署、執行和提供大語言模型服務的整體解決方案中非常有用,這種將 tokens 序列作為輸入並返回相應輸出 token 的實體通常被稱為執行引擎(execution engine)或推理引擎(inference engine)。

LLM 推理最佳化探微 (1) :Transformer 解碼器的推理過程詳解

圖 2 - 一個簡化的 Transformer 解碼器模型示意圖

那麼如何生成多個 tokens 呢?使用基於 Transformer 的解碼器從輸入文字序列(通常稱為提示語(prompt))生成文字(通常也被稱為對輸入文字的擴充套件或補充)基本上包括以下步驟:

  1. 將模型權重載入到 GPU

  2. 在 CPU 上對提示語(prompt)進行分詞(tokenizing),並將 token 張量傳輸到 GPU (圖3)

LLM 推理最佳化探微 (1) :Transformer 解碼器的推理過程詳解

圖 3 - 分詞步驟示意圖

    3.  將分詞完成後的提示語輸入神經網路,生成擴充套件的第一個token

這一階段通常被稱為啟動階段(initiation phase)。 在下一篇文章中,我們將看到它也經常被稱為預填充階段(pre-fill phase)。

    4. 將生成的 token 附加到輸入的 token 序列中,並將其用作生成擴充套件文字中第二個 token 的新輸入。然後,重複此過程,直到生成了停止序列(stop sequence)(例如,單個 end-of-sequence(EOS) token)或達到所配置的最大序列長度(圖4))。

這個由多個步驟組成的階段通常被稱為生成階段(generation phase)、解碼階段(decoding phase)、自迴歸階段(auto-regressive phase),甚至是增量階段(incremental phase)。

步驟3和步驟4都在下面的圖中(圖4)有所說明。

LLM 推理最佳化探微 (1) :Transformer 解碼器的推理過程詳解

圖4 - tokens 生成過程的啟動和解碼階段

    5. 將完成的 tokens 從 GPU 獲取到 CPU ,並對它們進行 detokenize(譯者注:”detokenize“指的是將模型生成的 tokens 序列轉換回原始文字或句子的過程。可能包括去除 tokens 之間的空格、新增標點符號、還原縮寫等操作,以還原生成文字的自然語言形式。),以獲取生成的文字(圖5)。

LLM 推理最佳化探微 (1) :Transformer 解碼器的推理過程詳解

圖 5 - detokenize步驟

注意:最近出現的旨在降低推理延遲的更先進技術,如投機抽樣(speculative sampling[3] )或前向解碼(lookahead decoding[4]),並不完全遵循上述簡單演算法。

說到這裡,你應該會感到失望、困惑或兩者兼而有之。您可能會問我:那麼啟動階段和解碼階段有什麼區別?現在看來,這充其量只是人為刻意區分的。啟動階段的感覺確實像是 while 迴圈的初始化步驟,而且我們在這兩個階段所做的事情本質上是一樣的:在每次迭代時,我們都對 tokens 序列進行前向傳遞,每次傳遞序列都會增加一個 token。

你說得沒錯。在這一點上,無論在硬體上如何進行計算,兩個階段之間確實沒有區別,因此兩個階段在這方面都沒有什麼特別之處。不過,正如我們將在下一篇文章中看到的那樣,這種設定涉及大量冗餘計算,因此在許多情況下效率低下。緩解這種情況的一種重要方式是快取我們不想重新計算的內容。這種最佳化被稱為 KV 快取,並引入了我一直在暗示的這兩個階段之間的關鍵差異。下一篇文章見!

Thanks for reading!

END

參考資料


來自 “ ITPUB部落格 ” ,連結:https://blog.itpub.net/70018536/viewspace-3005859/,如需轉載,請註明出處,否則將追究法律責任。

相關文章