面向生產(chǎn)的 LLM 優(yōu)化

2023-10-23 01:14 作者:HuggingFace 0人讀過 | 我要投稿

注意?:?本文同時也是 Transformers 的文檔。

以 GPT3/4、Falcon 以及 LLama 為代表的大語言模型 (Large Language Model，LLM) 在處理以人為中心的任務(wù)上能力突飛猛進，儼然已成為現(xiàn)代知識型行業(yè)的重要工具。

然而，在實際部署這些模型時，我們?nèi)悦媾R不少挑戰(zhàn):

為了展現(xiàn)可媲美人類的文本理解和生成能力，LLM 的參數(shù)量一般需要達到數(shù)十億 (參見 Kaplan 等人、Wei 等人的論述)，隨之而來的是對推理內(nèi)存的巨大需求。
在許多實際任務(wù)中，LLM 需要廣泛的上下文信息，這就要求模型在推理過程中能夠處理很長的輸入序列。

這些挑戰(zhàn)的關(guān)鍵在于增強 LLM 的計算和存儲效能，特別是如何增強長輸入序列的計算和存儲效能。

本文，我們將回顧迄今為止那些最有效的技術(shù)，以應(yīng)對高效 LLM 部署的挑戰(zhàn):

低精度: 研究表明，低精度 (即 8 比特和 4 比特) 推理可提高計算效率，且對模型性能沒有顯著影響。
Flash 注意力: Flash 注意力是注意力算法的一個變種，它不僅更節(jié)省內(nèi)存，而且通過優(yōu)化 GPU 內(nèi)存利用率從而提升了計算效率。
架構(gòu)創(chuàng)新: 考慮到 LLM 推理的部署方式始終為: 輸入序列為長文本的自回歸文本生成，因此業(yè)界提出了專門的模型架構(gòu)，以實現(xiàn)更高效的推理。這方面最重要的進展有 Alibi、旋轉(zhuǎn)式嵌入 (rotary embeddings) 、多查詢注意力 (Multi-Query Attention，MQA) ?以及分組查詢注意 (Grouped Query Attention，GQA) 。

本文，我們將從張量的角度對自回歸生成進行分析。我們深入研究了低精度的利弊，對最新的注意力算法進行了全面的探索，并討論了改進的 LLM 架構(gòu)。在此過程中，我們用實際的例子來展示每項技術(shù)所帶來的改進。

1. 充分利用低精度的力量

通過將 LLM 視為一組權(quán)重矩陣及權(quán)重向量，并將文本輸入視為向量序列，可以更好地理解 LLM 的內(nèi)存需求。下面，?權(quán)重?表示模型的所有權(quán)重矩陣及向量。

迄今為止，一個 LLM 至少有數(shù)十億參數(shù)。每個參數(shù)均為十進制數(shù)，例如?4.5689?通常存儲成 float32、bfloat16 或 float16 格式。因此，我們能夠輕松算出加載 LLM 所需的內(nèi)存:

加載? $X$ ?B 參數(shù)的 FP32 模型權(quán)重需要大約 4 *? $X$ ?GB 顯存

現(xiàn)如今，很少有模型以 float32 精度進行訓練，通常都是以 bfloat16 精度訓練的，在很少情況下還會以 float16 精度訓練。因此速算公式就變成了:

加載有 ? $X$ ?B 參數(shù)的 BF16/FP16 模型權(quán)重需要大約 2 *?X?GB 顯存

對于較短的文本輸入 (詞元數(shù)小于 1024)，推理的內(nèi)存需求很大程度上取決于模型權(quán)重的大小。因此，現(xiàn)在我們假設(shè)推理的內(nèi)存需求等于將模型加載到 GPU 中所需的顯存量。

我們舉幾個例子來說明用 bfloat16 加載模型大約需要多少顯存:

GPT3?需要 2 * 175 GB =?350 GB?顯存
Bloom?需要 2 * 176 GB =?352 GB?顯存
Llama-2-70b?需要 2 * 70 GB =?140 GB?顯存
Falcon-40b?需要 2 * 40 GB =?80 GB?顯存
MPT-30b?需要 2 * 30 GB =?60 GB?顯存
bigcode/starcoder：https://huggingface.co/bigcode/starcoder 需要 2 * 15.5 =?31 GB?顯存

迄今為止，市面上顯存最大的 GPU 芯片是 80GB 顯存的 A100。前面列出的大多數(shù)模型需要超過 80GB 才能加載，因此必然需要張量并行和/或流水線并行。

?? Transformers 不支持開箱即用的張量并行，因為它需要特定的模型架構(gòu)編寫方式。如果你對以張量并行友好的方式編寫模型感興趣，可隨時查看 TGI(text generation inference) 庫。

?? Transformers 開箱即用地支持簡單的流水線并行。為此，只需使用?device="auto"?加載模型，它會自動將不同層放到相應(yīng)的 GPU 上，詳見此處。但請注意，雖然非常有效，但這種簡單的流水線并行并不能解決 GPU 空閑的問題。可參考此處了解更高級的流水線并行技術(shù)。

如果你能訪問 8 x 80GB A100 節(jié)點，你可以按如下方式加載 BLOOM:

通過使用?device_map="auto"?，注意力層將均勻分布在所有可用的 GPU 上。

本文，我們選用 bigcode/octocoder 模型，因為它可以在單個 40GB A100 GPU 上運行。請注意，下文所有的內(nèi)存和速度優(yōu)化同樣適用于需要模型或張量并行的模型。

由于我們以 bfloat16 精度加載模型，根據(jù)上面的速算公式，預(yù)計使用?“bigcode/octocoder”?運行推理所需的顯存約為 31 GB。我們試試吧！

首先加載模型和分詞器，并將兩者傳遞給?Transformers?的 pipeline。

輸出:

好，現(xiàn)在我們可以把生成的函數(shù)直接用于將字節(jié)數(shù)轉(zhuǎn)換為千兆字節(jié)數(shù)。

我們直接調(diào)用?torch.cuda.max_memory_allocated?來測量 GPU 顯存的峰值占用。

輸出:

相當接近我們的速算結(jié)果！我們可以看到這個數(shù)字并不完全準確，因為從字節(jié)到千字節(jié)需要乘以 1024 而不是 1000。因此，速算公式也可以理解為“最多??GB”。

請注意，如果我們嘗試以全 float32 精度運行模型，則需要高達 64GB 的顯存。

現(xiàn)在幾乎所有模型都是用 bfloat16 中訓練的，如果你的 GPU 支持 bfloat16 的話，你就不應(yīng)該以 float32 來運行推理。float32 并不會提供比訓練精度更好的推理結(jié)果。

如果你不確定 Hub 上的模型權(quán)重的精度如何，可隨時查看模型配置文件內(nèi)的?torch_dtype?項，?如?此處。建議在使用?from_pretrained(..., torch_dtype=...)?加載模型時將精度設(shè)置為與配置文件中的精度相同，該接口的默認精度為 float32。這樣的話，你就可以使用?float16?或?bfloat16?來推理了。

我們再定義一個?flush(...)?函數(shù)來釋放所有已分配的顯存，以便我們可以準確測量分配的 GPU 顯存的峰值。

下一個實驗我們就可以調(diào)用它了。

在最新的 accelerate 庫中，你還可以使用名為?release_memory()?的方法。

那如果你的 GPU 沒有 32GB 顯存怎么辦？研究發(fā)現(xiàn)，模型權(quán)重可以量化為 8 比特或 4 比特，而對模型輸出沒有明顯影響 (參見 Dettmers 等人的論文)。

甚至可以將模型量化為 3 或 2 比特，對輸出的影響仍可接受，如最近的 GPTQ 論文 ?? 所示。

總的來講，量化方案旨在降低權(quán)重的精度，同時盡量保持模型的推理結(jié)果盡可能準確 (?即?盡可能接近 bfloat16)。

請注意，量化對于文本生成特別有效，因為我們關(guān)心的是選擇?最可能的下一個詞元的分布?，而不真正關(guān)心下一個詞元的確切?logit?值。所以，只要下一個詞元?logit?大小順序保持相同，?argmax?或?topk?操作的結(jié)果就會相同。

量化技術(shù)有很多，我們在這里不作詳細討論，但一般來說，所有量化技術(shù)的工作原理如下:

將所有權(quán)重量化至目標精度
加載量化權(quán)重，并把?bfloat16?精度的輸入向量序列傳給模型
將權(quán)重動態(tài)反量化為?bfloat16?，并基于?bfloat16?精度與輸入進行計算
計算后，將權(quán)重再次量化回目標精度。[譯者注: 這一步一般不需要做]

簡而言之，這意味著原來的每個?輸入數(shù)據(jù) - 權(quán)重矩陣乘?，其中為? $X$ ?輸入?， $W$ 權(quán)重矩陣，為輸出:?

$Y%20%3D%20X%20%5Ctimes%20W$

都變成了:

$Y%20%3D%20X%20%5Ctimes%20%5Ctext%7Bdequantize%7D(W)%3B%20%5Ctext%7Bquantize%7D(W)$

當輸入向量走過模型計算圖時，所有權(quán)重矩陣都會依次執(zhí)行反量化和重量化操作。

因此，使用權(quán)重量化時，推理時間通常?不會?減少，反而會增加。

到此為止理論講完了，我們可以開始試試了！要使用 Transformer 權(quán)重量化方案，請確保bitsandbytes?庫已安裝。

然后，只需在?from_pretrained?中添加?load_in_8bit=True?參數(shù)，即可用 8 比特量化加載模型。

現(xiàn)在，再次運行我們的示例，并測量其顯存使用情況。

輸出:

很好，我們得到了與之前一樣的結(jié)果，這就說明準確性沒有損失！我們看一下這次用了多少顯存。

輸出:?

顯存明顯減少！降至 15GB 多一點，這樣就可以在 4090 這樣的消費級 GPU 上運行該模型了。

我們看到內(nèi)存效率有了很大的提高，且模型的輸出沒啥退化。同時，我們也注意到推理速度出現(xiàn)了輕微的減慢。

刪除模型并再次刷一下顯存。

然后，我們看下 4 比特量化的 GPU 顯存消耗峰值是多少?？梢杂门c之前相同的 API 將模型量化為 4 比特 - 這次參數(shù)設(shè)置為?load_in_4bit=True?而不是?load_in_8bit=True?。

輸出:

輸出幾乎與以前相同 - 只是在代碼片段之前缺了?python?這個詞。我們看下需要多少顯存。

輸出:?

僅需 9.5GB！對于參數(shù)量大于 150 億的模型來說，確實不算多。

雖然我們這里看到模型的準確性幾乎沒有下降，但與 8 比特量化或完整的?bfloat16?推理相比，4 比特量化實際上通常會導致不同的結(jié)果。到底用不用它，就看用戶自己抉擇了。

另請注意，與 8 比特量化相比，其推理速度會更慢一些，這是由于 4 比特量化使用了更激進的量化方法，導致??和 ??在推理過程中花的時間更長。

總的來說，我們發(fā)現(xiàn)以 8 比特精度運行?OctoCoder?將所需的 GPU 顯存從 32GB 減少到僅 15GB，而以 4 比特精度運行模型則進一步將所需的 GPU 顯存減少到 9GB 多一點。

4 比特量化讓模型可以在 RTX3090、V100 和 T4 等大多數(shù)人都可以輕松獲取的 GPU 上運行。

更多有關(guān)量化的信息以及有關(guān)如何量化模型以使其顯存占用比 4 比特更少，我們建議大家查看?AutoGPTQ?的實現(xiàn)。

總結(jié)一下，重要的是要記住，模型量化會提高內(nèi)存效率，但會犧牲準確性，在某些情況下還會犧牲推理時間。

如果 GPU 顯存對你而言不是問題，通常不需要考慮量化。然而，如果不量化，許多 GPU 根本無法運行 LLM，在這種情況下，4 比特和 8 比特量化方案是非常有用的工具。

更詳細的使用信息，我們強烈建議你查看 Transformers 的量化文檔。

接下來，我們看看如何用更好的算法和改進的模型架構(gòu)來提高計算和內(nèi)存效率。

2. Flash 注意力: 速度飛躍

當今表現(xiàn)最好的 LLM 其基本架構(gòu)大體相似，包括前饋層、激活層、層歸一化層以及最重要的自注意力層。

自注意力層是大語言模型 (LLM) 的核心，因為其使模型能夠理解輸入詞元之間的上下文關(guān)系。然而，自注意力層在計算以及峰值顯存這兩個方面都隨著輸入詞元的數(shù)目 (也稱為?序列長度?，下文用? $N$ ?表示) 呈?二次方?增長。

雖然這對于較短的輸入序列 (輸入詞元數(shù)小于 1000) 來說并不明顯，但對于較長的輸入序列 (如: 約 16000 個輸入詞元) 來說，就會成為一個嚴重的問題。

我們仔細分析一下。計算長度為 $N$ 的輸入序列 $%5Cmathbf%7BX%7D$ 的自注意力層的輸出 $%5Cmathbf%7BO%7D$ ，其公式為:

$%5Ctextbf%7BO%7D%20%3D%20%5Ctext%7BAttn%7D(%5Cmathbf%7BX%7D)%20%3D%20%5Cmathbf%7BV%7D%20%5Ctimes%20%5Ctext%7BSoftmax%7D(%5Cmathbf%7BQK%7D%5ET)%20%5Ctext%7B%20%EF%BC%8C%E5%85%B6%E4%B8%AD%20%7D%20%5Cmathbf%7BQ%7D%20%3D%20%5Cmathbf%7BW%7D_q%20%5Cmathbf%7BX%7D%2C%20%5Cmathbf%7BV%7D%20%3D%20%5Cmathbf%7BW%7D_v%20%5Cmathbf%7BX%7D%2C%20%5Cmathbf%7BK%7D%20%3D%20%5Cmathbf%7BW%7D_k%20%5Cmathbf%7BX%7D$

$%5Cmathbf%7BX%7D%20%3D%20(%5Cmathbf%7Bx%7D%20_1%2C%20%E2%80%A6%20%5Cmathbf%7Bx%7D_%20%7BN%7D)$ 是注意力層的輸入序列。投影 $%5Cmathbf%7BQ%7D$ 和 $%5Cmathbf%7BK%7D$ 也是 $N$ 個向量組成的序列，其乘積 $%5Cmathbf%7BQK%7D%5ET$ 的大小為 $N%5E2$ 。

LLM 通常有多個注意力頭，因此可以并行進行多個自注意力計算。假設(shè) LLM 有 40 個注意力頭并以 bfloat16 精度運行，我們可以計算出存儲 $%5Cmathbf%7BQK%5ET%7D$ 矩陣的內(nèi)存需求為 $40%20%5Ctimes%202%20%5Ctimes%20N%5E2$ 字節(jié)。當 $N%3D1000$ 時僅需要大約 50MB 的顯存，但當 $N%3D16000$ 時，我們需要 19GB 的顯存，當 $N%3D100%2C000$ 時，僅存儲矩陣就需要近 1TB。

總之，隨著輸入上下文越來越長，默認的自注意力算法所需的內(nèi)存很快就會變得非常昂貴。

伴隨著 LLM 在文本理解和生成方面的進展，它們正被應(yīng)用于日益復雜的任務(wù)。之前，我們主要用模型來對幾個句子進行翻譯或摘要，但現(xiàn)在我們會用這些模型來管理整頁的文本，這就要求它們具有處理長輸入文本的能力。

我們?nèi)绾螖[脫長輸入文本對內(nèi)存的過高要求？我們需要一種新的方法讓我們在計算自注意力機制時能夠擺脫矩陣。 Tri Dao 等人開發(fā)了這樣一種新算法，并將其稱為?Flash 注意力。

簡而言之，F(xiàn)lash 注意力將 $%5Cmathbf%7BV%7D%20%5Ctimes%20%5Ctext%7BSoftmax%7D(%5Cmathbf%7BQK%7D%5ET)$ 的計算分解成若干步驟，通過迭代多個 softmax 計算步來將輸出分成多個較小的塊進行計算:

$%5Ctextbf%7BO%7D%20_i%20%5Cleftarrow%20s%5Ea_%20%7Bij%7D%20%5Ctimes%20%5Ctextbf%7BO%7D%20_i%20%2B%20s%5Eb_%20%7Bij%7D%20%5Ctimes%20%5Cmathbf%7BV%7D%20_%7Bj%7D%20%5Ctimes%20%5Ctext%7BSoftmax%7D(%5Cmathbf%7BQK%7D%5ET_%20%7Bi%2Cj%7D)%20%5Ctext%7B%EF%BC%8C%E5%9C%A8%20%7D%20i%2C%20j%20%5Ctext%7B%20%E4%B8%8A%E8%BF%AD%E4%BB%A3%7D$

其中 $s%5Ea_%7Bij%7D$ 和 $s%5Eb_%7Bij%7D$ 是隨著每個? $i$ 和? $j$ 迭代更新的 softmax 統(tǒng)計歸一化值。

請注意，整個 Flash 注意力有點復雜，這里已經(jīng)大大簡化了。如果想要深入理解，可以閱讀 Flash Attention 的論文。

要點如下:

通過跟蹤 softmax 統(tǒng)計歸一化值再加上一些聰明的數(shù)學技巧，與默認的自注意力層相比，F(xiàn)lash 注意力的計算結(jié)果?完全相同，而內(nèi)存成本僅隨著 $N$ ?線性增加。

僅看這個公式，直覺上來講，F(xiàn)lash 注意力肯定比默認的自注意力公式要慢很多，因為需要進行更多的計算。確實，與普通注意力相比，F(xiàn)lash 注意力需要更多的 FLOP，因為需要不斷重新計算 softmax 統(tǒng)計歸一化值 (如果感興趣，請參閱論文以了解更多詳細信息)。

然而，與默認注意力相比，F(xiàn)lash 注意力的推理速度要快得多，這是因為它能夠顯著減少對較慢的高帶寬顯存的需求，而更多使用了更快的片上內(nèi)存 (SRAM)。

從本質(zhì)上講，F(xiàn)lash 注意力確保所有中間寫入和讀取操作都可以使用快速?片上?SRAM 來完成，而不必訪問較慢的顯存來計算輸出向量 $%5Cmathbf%7BO%7D$ 。

實際上，如果能用的話，我們沒有理由不用 Flash 注意力。該算法在數(shù)學上給出相同的輸出，但速度更快且內(nèi)存效率更高。

我們看一個實際的例子。

我們的?OctoCoder?模型現(xiàn)在被輸入了長得多的提示，其中包括所謂的“系統(tǒng)提示”。系統(tǒng)提示用于引導 LLM 去適應(yīng)特定的用戶任務(wù)。

接下來，我們使用系統(tǒng)提示，引導?OctoCoder?成為更好的編程助手。

為了演示需要，我們將系統(tǒng)提示復制十倍，以便輸入長度足夠長以觀察 Flash 注意力帶來的內(nèi)存節(jié)省。然后在其后加上原始提示?"Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here"?:?

以 bfloat16 精度再次初始化模型。

現(xiàn)在，我們可以像以前一樣運行模型，同時測量其峰值 GPU 顯存需求及推理時間。

輸出:

輸出與之前一樣，但是這一次，模型會多次重復答案，直到達到 60 個詞元為止。這并不奇怪，因為出于演示目的，我們將系統(tǒng)提示重復了十次，從而提示模型重復自身。

注意，在實際應(yīng)用中，系統(tǒng)提示不應(yīng)重復十次 —— 一次就夠了！

我們測量一下峰值 GPU 顯存需求。

輸出:

正如我們所看到的，峰值 GPU 顯存需求現(xiàn)在明顯高于以前，這主要是因為輸入序列變長了。整個生成過程也需要一分多鐘的時間。

我們調(diào)用?flush()?來釋放 GPU 內(nèi)存以供下一個實驗使用。

為便于比較，我們運行相同的函數(shù)，但啟用 Flash 注意力。為此，我們將模型轉(zhuǎn)換為 BetterTransformers，這會因此而啟用 PyTorch 的 SDPA 自注意力，其實現(xiàn)是基于 Flash 注意力的。

現(xiàn)在我們運行與之前完全相同的代碼片段，但此時 Transformers 在底層將使用 Flash 注意力。

輸出:?

結(jié)果與之前完全相同，但由于 Flash 注意力，我們可以觀察到非常顯著的加速。

我們最后一次測量一下內(nèi)存消耗。

輸出:?

我們幾乎一下就回到了原來的 29GB 峰值 GPU 顯存。

我們可以觀察到，與剛開始的短輸入序列相比，使用 Flash 注意力且輸入長序列時，我們只多用了大約 100MB 的 GPU 顯存。

3. 架構(gòu)背后的科學: 長文本輸入和聊天式 LLM 的策略選擇

到目前為止，我們已經(jīng)研究了通過以下方式提高計算和內(nèi)存效率:

將權(quán)重轉(zhuǎn)換為較低精度的格式
用內(nèi)存和計算效率更高的版本替換自注意力算法

現(xiàn)在讓我們看看如何改變 LLM 的架構(gòu)，使其對于需要長文本輸入的任務(wù)更高效，?例如?:

檢索增強問答
總結(jié)
聊天

請注意，?聊天?應(yīng)用不僅需要 LLM 處理長文本輸入，還需要 LLM 能夠有效地處理用戶和助手之間的多輪對話 (例如 ChatGPT)。

一旦經(jīng)過訓練，LLM 的基本架構(gòu)就很難改變，因此提前考慮 LLM 的任務(wù)特征并相應(yīng)地優(yōu)化模型架構(gòu)非常重要。模型架構(gòu)中有兩個重要組件很快就會成為長輸入序列的內(nèi)存和/或性能瓶頸。

位置嵌入 (positional embeddings)
鍵值緩存 (key-value cache)

我們來一一詳細探討:

3.1 改進 LLM 的位置嵌入

自注意力機制計算每個詞元間的相關(guān)系數(shù)。例如，文本輸入序列?“Hello”, “I”, “l(fā)ove”, “you”?的 $%5Ctext%7BSoftmax%7D(%5Cmathbf%7BQK%7D%5ET)$ 矩陣看起來如下:

每個詞元都會被賦予一個概率值，表示其對另一個詞元的關(guān)注度。例如，?“l(fā)ove”?這個詞關(guān)注?“Hello”?這個詞的概率為 0.05%，關(guān)注?“I”?的概率為 0.3%，而對自己的關(guān)注概率則為 0.65%。

基于自注意力但沒有位置嵌入的 LLM 在理解輸入文本彼此的相對位置上會遇到很大困難。這是因為在經(jīng)由 $%5Cmathbf%7BQK%7D%5ET$ 來計算相關(guān)概率時，其計算是與詞元間的相對距離無關(guān)的，即該計算與詞元間的相對距離的關(guān)系為 $O(1)$ 。因此，對于沒有位置嵌入的 LLM，每個詞元似乎與所有其他詞元等距。?此時?，區(qū)分?“Hello I love you”?和?“You love I hello”?會比較困難。

為了讓能夠 LLM 理解語序，需要額外的?提示?，通常我們用?位置編碼?(也稱為?位置嵌入?) 來注入這種提示。位置編碼將每個詞元的位置編碼為數(shù)字，LLM 可以利用這些數(shù)字更好地理解語序。

Attention Is All You Need??論文引入了正弦位置嵌入 $%5Cmathbf%7BP%7D%20%3D%20%5Cmathbf%7Bp%7D_1%2C%20%5Cldots%2C%20%5Cmathbf%7Bp%7D_N$ 。其中每個向量? $%5Cmathbf%7Bp%7D_i$ 為其位置 $i$ 的正弦函數(shù)。然后將位置編碼與輸入序列向量簡單相加 $%5Cmathbf%7B%5Chat%7BX%7D%7D%20%3D%20%5Cmathbf%7B%5Chat%7Bx%7D%7D_1%2C%20%5Cldots%2C%20%5Cmathbf%7B%5Chat%7Bx%7D%7D_N%24%20%3D%20%24%5Cmathbf%7Bx%7D_1%20%2B%20%5Cmathbf%7Bp%7D_1%2C%20%5Cldots%2C%20%5Cmathbf%7Bx%7D_N%20%2B%20%5Cmathbf%7Bp%7D_N$ 從而提示模型更好地學習語序。

其他工作 (如 Devlin 等人的工作) 沒有使用固定位置嵌入，而是使用可訓練的位置編碼，在訓練期間學習位置嵌入 $%5Cmathbf%7BP%7D$ 。

曾經(jīng)，正弦位置嵌入以及可訓練位置嵌入是將語序編碼進 LLM 的主要方法，但這兩個方法會有一些問題:

正弦位置嵌入以及可訓練位置嵌入都是絕對位置嵌入，?即?為每個位置 id ( $0%2C%20%5Cldots%2C%20N$ )生成一個唯一的嵌入。正如 Huang et al. 和 Su et al. 的工作所示，絕對位置嵌入會導致 LLM 在處理長文本輸入時性能較差。對長文本輸入而言，如果模型能夠?qū)W習輸入詞元間的相對距離而不是它們的絕對位置，會比較好。
當使用訓練位置嵌入時，LLM 必須在固定的輸入長度 $N$ ?上進行訓練，因此如果推理時的輸入長度比訓練長度更長，外插會比較麻煩。

最近，可以解決上述問題的相對位置嵌入變得越來越流行，其中應(yīng)用最多的有兩個:

旋轉(zhuǎn)位置嵌入 (Rotary Position Embedding， RoPE)
ALiBi

RoPE?和?ALiBi?都認為，最好直接在自注意力算法中向 LLM 提示語序，因為詞元是通過自注意力機制互相關(guān)聯(lián)的。更具體地說，應(yīng)該通過 $%5Cmathbf%7BQK%7D%5ET$ 修改的計算來提示語序。

簡而言之，?RoPE?指出位置信息可以編碼為?查詢 - 鍵值對?，?如? $i%2C%20j$ 通過分別將每個向量根據(jù)其在句子中的位置旋轉(zhuǎn)角度 $%5Ctheta%20%5Ctimes%20i%20%E5%92%8C%20%5Ctheta%20%5Ctimes%20j$ :

$%5Cmathbf%7B%5Chat%7Bq%7D%7D_i%5ET%20%5Cmathbf%7B%5Chat%7Bx%7D%7D_j%20%3D%20%5Cmathbf%7B%7Bq%7D%7D%20_i%5ET%20%5Cmathbf%7BR%7D_%20%7B%5Ctheta%2C%20i%20-j%7D%20%5Cmathbf%7B%7Bx%7D%7D_j.$

$%5Cmathbf%7BR%7D_%7B%5Ctheta%2C%20i%20-%20j%7D$ 表示旋轉(zhuǎn)矩陣。 $%5Ctheta$ ?在不可訓練的預(yù)定義值，其值取決于訓練期間最大輸入序列長度。

通過這樣做， $%5Cmathbf%7Bq%7D_i$ ?和? $%5Cmathbf%7Bq%7D_j$ ?之間的概率得分僅受 $i%20%5Cne%20j$ ?是否成立這一條件影響，且其值僅取決于相對距離 $i%20-%20j$ , 而與每個向量的具體位置? $i$ 和 $j$ ?無關(guān)。

如今，多個最重要的 LLM 使用了?RoPE?，例如:

Falcon
Llama
PaLM

另一個方案是?ALiBi?，它提出了一種更簡單的相對位置編碼方案。在計算 softmax 之前， $%5Cmathbf%7BQK%7D%5ET$ 矩陣的每個元素會減去被一個預(yù)定義系數(shù)?m?縮放后的對應(yīng)兩個向量間的相對距離。

編輯搜圖

如 ALiBi 論文所示，這種簡單的相對位置編碼使得模型即使在很長的文本輸入序列中也能保持高性能。

當前也有多個最重要的 LLM 使用了?ALiBi?，如:

MPT
BLOOM

RoPE?和?ALiBi?位置編碼都可以外推到訓練期間未見的輸入長度，而事實證明，與?RoPE?相比，?ALiBi?的外推效果要好得多。對于 ALiBi，只需簡單地增加下三角位置矩陣的值以匹配輸入序列的長度即可。而對于?RoPE?，如果輸入長度比訓練期間的輸入長得多，使用訓練期間? $%5Ctheta$ ?值的生成效果不好，?參見?Press et al.。然而，社區(qū)已經(jīng)找到了一些調(diào)整? $%5Ctheta$ ?的有效技巧。從而允許?RoPE?位置嵌入能夠很好地應(yīng)對輸入序列外插的狀況 (請參閱此處)。

RoPE 和 ALiBi 都是相對位置嵌入，其嵌入?yún)?shù)是?不可?訓練的，而是基于以下直覺:

有關(guān)輸入文本的位置提示應(yīng)直接提供給自注意力層的 QKt 矩陣
應(yīng)該激勵 LLM 學習基于恒定?相對?距離的位置編碼
輸入詞元間彼此距離越遠，它們的?查詢 - 鍵?概率越低。 RoPE 和 ALiBi 都降低了距離較遠詞元間的?查詢 - 鍵?概率。RoPE 通過增加?查詢 - 鍵?向量之間的夾角來減少它們的向量積。而 ALiBi 通過從向量積中減去一個更大的數(shù)來達成這個目的。

總之，打算部署在需要處理長文本輸入的任務(wù)中的 LLM 可以通過相對位置嵌入 (例如 RoPE 和 ALiBi) 來進行更好的訓練。另請注意，使用了 RoPE 和 ALiBi 的 LLM 即使是僅在固定長度 (例如 $N_1%20%3D%202048$ ) 上訓練的，其仍然可以在推理時通過位置嵌入外插來處理比 $N_1$ 長得多的文本輸入 (如 $N_2%20%3D%208192%20%3E%20N_1$ )。

3.2 鍵值緩存

使用 LLM 進行自回歸文本生成的工作原理是把輸入序列輸入給模型，并采樣獲得下一個詞元，再將獲得的詞元添加到輸入序列后面，如此往復，直到 LLM 生成一個表示結(jié)束的詞元。

請查閱 Transformer 的文本生成教程以更直觀地了解自回歸生成的工作原理。

下面，我們快速運行一個代碼段來展示自回歸是如何工作的。我們簡單地使用?torch.argmax?獲取最有可能的下一個詞元。

輸出:?

正如我們所看到的，每次我們都把剛剛采樣出的詞元添加到輸入文本中。

除了極少數(shù)例外，LLM 都是基于因果語言模型的目標函數(shù)進行訓練的，因此我們不需要注意力矩陣的上三角部分 - 這就是為什么在上面的兩個圖中，上三角的注意力分數(shù)是空的 (?也即?概率為 0)。想要快速入門因果語言模型，你可以參考這篇?圖解自注意力??博文。

因此，當前詞元?永遠僅?依賴于其前面的詞元，更具體地說， $%5Cmathbf%7Bq%7D%20_i$ 向量永遠與任何的 $j%20%3E%20i$ 鍵、值向量無關(guān)聯(lián)。相反 $%5Cmathbf%7Bq%7D%20_i$ 僅關(guān)注其之前的鍵、值向量 $%5Cmathbf%7Bk%7D_%20%7Bm%20%3C%20i%7D%2C%20%5Cmathbf%7Bv%7D_%20%7Bm%20%3C%20i%7D%20%5Ctext%7B%EF%BC%8C%7D%20m%20%5Cin%20%7B0%2C%20%5Cldots%20i%20-%201%7D$ 。為了減少不必要的計算，因此可以把先前所有步的每一層的鍵、值向量緩存下來。

接下來，我們將告訴 LLM 在每次前向傳播中都利用鍵值緩存來減少計算量。在 Transformers 中，我們可以通過將?use_cache?參數(shù)傳給?forward?來利用鍵值緩存，這樣的話，每次推理僅需傳當前詞元給?forward?就可以。

輸出:?

正如我們所看到的，當使用鍵值緩存時，輸入文本的長度?沒有?增加，每次都只有一個向量。另一方面，鍵值緩存的長度每解碼步都增加了一。

利用鍵值緩存意味著 $%5Cmathbf%7BQK%7D%5ET$ ?本質(zhì)上減少為 $%5Cmathbf%7Bq%7D_c%5Cmathbf%7BK%7D%5ET$ ，其中 $%5Cmathbf%7Bq%7D_c$ ?是當前輸入詞元的查詢投影，它?始終?只是單個向量。

使用鍵值緩存有兩個優(yōu)點:

與計算完整的 $%5Cmathbf%7BQK%7D%5ET$ ?矩陣相比，計算量更小，計算效率顯著提高，因此推理速度也隨之提高。
所需的最大內(nèi)存不隨生成的詞元數(shù)量呈二次方增加，而僅呈線性增加。

用戶應(yīng)該?始終?使用鍵值緩存，因為它的生成結(jié)果相同且能顯著加快長輸入序列的生成速度。當使用文本 pipeline 或?generate?方法時，Transformers 默認啟用鍵值緩存。

請注意，鍵值緩存對于聊天等需要多輪自回歸解碼的應(yīng)用程序特別有用。我們看一個例子。

在這個聊天示例中，LLM 需自回歸解碼兩次:

第一次，鍵值緩存為空，輸入提示為?"User: How many people live in France?"?，模型自回歸生成文本?"Roughly 75 million people live in France"?，同時在每個解碼步添加鍵值緩存。
第二次輸入提示為?"User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many in Germany?"?。由于緩存，前兩個句子的所有鍵值向量都已經(jīng)計算出來。因此輸入提示僅包含?"User: And how many in Germany?"?。在處理縮短的輸入提示時，計算出的鍵值向量將添加到第一次解碼的鍵值緩存后面。然后，助手使用鍵值緩存自回歸地生成第二個問題的答案?"Germany has ca. 81 million inhabitants"?，該鍵值緩存是?"User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many are in Germany?"?的編碼向量序列。

這里需要注意兩件事:

保留所有上下文對于在聊天場景中部署的 LLM 至關(guān)重要，以便 LLM 理解對話的所有上文。例如，上面的示例中，LLM 需要了解用戶在詢問?"And how many are in Germany"?時指的是人口。
鍵值緩存對于聊天非常有用，因為它允許我們不斷增長聊天歷史記錄的編碼緩存，而不必對聊天歷史記錄從頭開始重新編碼 (當使用編碼器 - 解碼器時架構(gòu)時我們就不得不這么做)。

然而，還有一個問題。雖然 $%5Cmathbf%7BQK%7D%5ET$ 矩陣所需的峰值內(nèi)存顯著減少，但對于長輸入序列或多輪聊天，將鍵值緩存保留在內(nèi)存中還是會非常昂貴。請記住，鍵值緩存需要存儲先前所有輸入向量? $%5Cmathbf%7Bx%7D_i%20%5Ctext%7B%2C%20for%20%7D%20i%20%5Cin%20%7B1%2C%20%5Cldots%2C%20c%20-%201%7D$ ?的所有層、所有注意力頭的鍵值向量。

我們計算一下我們之前使用的 LLM?bigcode/octocoder?需要存儲在鍵值緩存中的浮點數(shù)的個數(shù)。浮點數(shù)的個數(shù)等于序列長度的兩倍乘以注意力頭的個數(shù)乘以注意力頭的維度再乘以層數(shù)。假設(shè)輸入序列長度為 16000，我們計算得出:?

輸出:?

大約 80 億個浮點數(shù)！以?float16?精度存儲 80 億個浮點值需要大約 15 GB 的顯存，大約是模型本身權(quán)重的一半！

研究人員提出了兩種方法，用于顯著降低鍵值緩存的內(nèi)存成本:

多查詢注意力 (Multi-Query-Attention，MQA)

多查詢注意力機制是 Noam Shazeer 在?Fast Transformer Decoding: One Write-Head is All You Need?論文中提出的。正如標題所示，Noam 發(fā)現(xiàn)，可以在所有注意力頭之間共享同一對鍵、值投影權(quán)重，而不是使用?n_head?對鍵值投影權(quán)重，這并不會顯著降低模型的性能。

通過共享同一對鍵、值投影權(quán)重，鍵值向量? $%5Cmathbf%7Bk%7D_i%2C%20%5Cmathbf%7Bv%7D_i$ ?在所有注意力頭上相同，這意味著我們只需要緩存 1 個鍵值投影對，而不需要 n_head 對。

除了節(jié)省內(nèi)存之外，MQA 還可以提高計算效率。在自回歸解碼中，需要重新加載大的鍵值向量，與當前的鍵值向量對相串接，然后將其輸入到每一步 $%5Cmathbf%7Bq%7D_c%5Cmathbf%7BK%7D%5ET$ 的計算中。對于自回歸解碼，不斷重新加載所需的內(nèi)存帶寬可能成為嚴重的性能瓶頸。通過減少鍵值向量的大小，需要訪問的內(nèi)存更少，從而減少內(nèi)存帶寬瓶頸。欲了解更多詳細信息，請查看 Noam 的論文。

這里的重點是，只有使用鍵值緩存時，將鍵值注意力頭的數(shù)量減少到 1 才有意義。沒有鍵值緩存時，模型單次前向傳播的峰值內(nèi)存消耗保持不變，因為每個注意力頭 $%5Cmathbf%7BQK%7D%5ET$ 查詢向量不同，因此每個注意力頭的矩陣也不相同。

MQA 已被社區(qū)廣泛采用，現(xiàn)已被許多流行的 LLM 所采用:

此外，本文所使用的檢查點 -?bigcode/octocoder?- 也使用了 MQA。

Falcon
PaLM
MPT
BLOOM
分組查詢注意力 (Grouped-Query-Attention，GQA)

分組查詢注意力由來自 Google 的 Ainslie 等人提出，它們發(fā)現(xiàn)，與原始的多頭鍵值投影相比，使用 MQA 通常會導致生成質(zhì)量下降。該論文認為，通過不太大幅度地減少查詢頭投影權(quán)重的數(shù)量可以獲得更高的模型性能。不應(yīng)僅使用單個鍵值投影權(quán)重，而應(yīng)使用?n < n_head?個鍵值投影權(quán)重。通過將?n?設(shè)為比?n_head?小得多的值 (例如 2，4 或 8)，幾乎可以保留 MQA 帶來的所有內(nèi)存和速度增益，同時更少地犧牲模型能力，或者說說僅略微犧牲模型性能。

此外，GQA 的作者發(fā)現(xiàn)，現(xiàn)有的模型檢查點可以通過?升級訓練?，變成 GQA 架構(gòu)，而其所需的計算量僅為原始預(yù)訓練計算的 5%。雖然 5% 的原始預(yù)訓練計算量仍然很大，但 GQA?升級訓練?允許現(xiàn)有 checkpoint 通過這個機制，升級成能處理長輸入序列的 checkpoint，這點還是挺誘人的。

GQA 最近才被提出，這就是為什么截至本文撰寫時其被采用得較少。GQA 最著名的應(yīng)用是 Llama-v2。

總之，如果部署自回歸解碼的 LLM 并且需要處理長輸入序列 (例如聊天)，我們強烈建議使用 GQA 或 MQA。

總結(jié)

研究界不斷提出新的、巧妙的方法來加速更大的 LLM 的推理。舉個例子，一個頗有前景的研究方向是投機解碼，其中“簡單詞元”是由更小、更快的語言模型生成的，而只有“難詞元”是由 LLM 本身生成的。詳細介紹超出了本文的范圍，但可以閱讀這篇不錯的博文。

GPT3/4、Llama-2-70b、Claude、PaLM 等海量 LLM 能夠在 Hugging Face Chat 或 ChatGPT 等聊天應(yīng)用中快速運行的原因是很大一部分歸功于上述精度、算法和架構(gòu)方面的改進。展望未來，GPU、TPU 等加速器只會變得更快且內(nèi)存更大，但人們?nèi)匀粦?yīng)該始終確保使用最好的可用算法和架構(gòu)來獲得最大的收益 ??。

英文原文:?https://hf.co/blog/optimize-llm

原文作者: Patrick von Platen

譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態(tài)數(shù)據(jù)上的應(yīng)用及大規(guī)模模型的訓練推理。

審校/排版: zhongdongy (阿東)

標簽：

最美情侣中文字幕电影,在线麻豆精品传媒,在线网站高清黄,久久黄色视频

面向生產(chǎn)的 LLM 優(yōu)化

1. 充分利用低精度的力量

2. Flash 注意力: 速度飛躍

3.1 改進 LLM 的位置嵌入

3.2 鍵值緩存

總結(jié)