面向生產的 LLM 最佳化

HuggingFace發表於2023-10-19

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

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

然而，在實際部署這些模型時，我們仍面臨不少挑戰:

  • 為了展現可媲美人類的文字理解和生成能力，LLM 的引數量一般需要達到數十億 (參見 Kaplan 等人、Wei 等人 的論述)，隨之而來的是對推理記憶體的巨大需求。
  • 在許多實際任務中，LLM 需要廣泛的上下文資訊，這就要求模型在推理過程中能夠處理很長的輸入序列。

這些挑戰的關鍵在於增強 LLM 的計算和儲存效能，特別是如何增強長輸入序列的計算和儲存效能。

本文，我們將回顧迄今為止那些最有效的技術，以應對高效 LLM 部署的挑戰:

  1. 低精度: 研究表明，低精度 (即 8 位元和 4 位元) 推理可提高計算效率，且對模型效能沒有顯著影響。
  2. Flash 注意力: Flash 注意力是注意力演算法的一個變種，它不僅更節省記憶體，而且透過最佳化 GPU 記憶體利用率從而提升了計算效率。
  3. 架構創新: 考慮到 LLM 推理的部署方式始終為: 輸入序列為長文字的自迴歸文字生成，因此業界提出了專門的模型架構，以實現更高效的推理。這方面最重要的進展有 Alibi、旋轉式嵌入 (rotary embeddings) 、多查詢注意力 (Multi-Query Attention，MQA) 以及 分組查詢注意 (Grouped Query Attention，GQA) 。

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

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

透過將 LLM 視為一組權重矩陣及權重向量，並將文字輸入視為向量序列，可以更好地理解 LLM 的記憶體需求。下面， 權重 表示模型的所有權重矩陣及向量。

迄今為止，一個 LLM 至少有數十億引數。每個引數均為十進位制數，例如 4.5689 通常儲存成 float32、bfloat16 或 float16 格式。因此，我們能夠輕鬆算出載入 LLM 所需的記憶體:

載入 \(X\) B 引數的 FP32 模型權重需要大約 4 * \(X\) GB 視訊記憶體

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

載入有 \(X\) B 引數的 BF16/FP16 模型權重需要大約 2 * \(X\) GB 視訊記憶體

對於較短的文字輸入 (詞元數小於 1024)，推理的記憶體需求很大程度上取決於模型權重的大小。因此，現在我們假設推理的記憶體需求等於將模型載入到 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 需要 2 * 15.5 = 31 GB 視訊記憶體

迄今為止，市面上視訊記憶體最大的 GPU 晶片是 80GB 視訊記憶體的 A100。前面列出的大多數模型需要超過 80GB 才能載入，因此必然需要 張量並行 和/或 流水線並行。

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

? Transformers 開箱即用地支援簡單的流水線並行。為此，只需使用 device="auto" 載入模型，它會自動將不同層放到相應的 GPU 上，詳見 此處。
但請注意，雖然非常有效，但這種簡單的流水線並行並不能解決 GPU 空閒的問題。可參考 此處 瞭解更高階的流水線並行技術。

如果你能訪問 8 x 80GB A100 節點，你可以按如下方式載入 BLOOM:

!pip install transformers accelerate bitsandbytes optimum

# from transformers import AutoModelForCausalLM

# model = AutoModelForCausalLM.from_pretrained("bigscience/bloom", device_map="auto", pad_token_id=0)

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

本文，我們選用 bigcode/octocoder 模型，因為它可以在單個 40GB A100 GPU 上執行。請注意，下文所有的記憶體和速度最佳化同樣適用於需要模型或張量並行的模型。

由於我們以 bfloat16 精度載入模型，根據上面的速算公式，預計使用 “bigcode/octocoder” 執行推理所需的視訊記憶體約為 31 GB。我們試試吧！

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

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", torch_dtype=torch.bfloat16, device_map="auto", pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer:"

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

輸出:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

好，現在我們可以把生成的函式直接用於將位元組數轉換為千兆位元組數。

def bytes_to_giga_bytes(bytes):
  return bytes / 1024 / 1024 / 1024

我們直接呼叫 torch.cuda.max_memory_allocated 來測量 GPU 視訊記憶體的峰值佔用。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

輸出:

29.0260648727417

相當接近我們的速算結果！我們可以看到這個數字並不完全準確，因為從位元組到千位元組需要乘以 1024 而不是 1000。因此，速算公式也可以理解為“最多 \(X\) GB”。

請注意，如果我們嘗試以全 float32 精度執行模型，則需要高達 64GB 的視訊記憶體。

現在幾乎所有模型都是用 bfloat16 中訓練的，如果 你的 GPU 支援 bfloat16 的話，你就不應該以 float32 來執行推理。float32 並不會提供比訓練精度更好的推理結果。

如果你不確定 Hub 上的模型權重的精度如何，可隨時檢視模型配置檔案內的 torch_dtype 項， 如 此處。建議在使用 from_pretrained(..., torch_dtype=...) 載入模型時將精度設定為與配置檔案中的精度相同，該介面的預設精度為 float32。這樣的話，你就可以使用 float16 或 bfloat16 來推理了。

我們再定義一個 flush(...) 函式來釋放所有已分配的視訊記憶體，以便我們可以準確測量分配的 GPU 視訊記憶體的峰值。

del pipe
del model

import gc
import torch

def flush():
  gc.collect()
  torch.cuda.empty_cache()
  torch.cuda.reset_peak_memory_stats()

下一個實驗我們就可以呼叫它了。

flush()

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

from accelerate.utils import release_memory
# ...

release_memory(model)

那如果你的 GPU 沒有 32GB 視訊記憶體怎麼辦？研究發現，模型權重可以量化為 8 位元或 4 位元，而對模型輸出沒有明顯影響 (參見 Dettmers 等人的論文)。

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

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

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

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

  1. 將所有權重量化至目標精度
  2. 載入量化權重，並把 bfloat16 精度的輸入向量序列傳給模型
  3. 將權重動態反量化為 bfloat16 ，並基於 bfloat16 精度與輸入進行計算
  4. 計算後，將權重再次量化回目標精度。[譯者注: 這一步一般不需要做]

簡而言之，這意味著原來的每個 輸入資料 - 權重矩陣乘 ，其中 \(X\) 為 輸入 ， \(W\) 為權重矩陣，\(Y\) 為輸出:

\[Y = X \times W \]

都變成了:

\[Y = X \times \text{dequantize}(W); \text{quantize}(W) \]

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

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

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

# !pip install bitsandbytes

然後，只需在 from_pretrained 中新增 load_in_8bit=True 引數，即可用 8 位元量化載入模型。

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=True, pad_token_id=0)

現在，再次執行我們的示例，並測量其視訊記憶體使用情況。

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

輸出:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

很好，我們得到了與之前一樣的結果，這就說明準確性沒有損失！我們看一下這次用了多少視訊記憶體。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

輸出:

15.219234466552734

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

我們看到記憶體效率有了很大的提高，且模型的輸出沒啥退化。同時，我們也注意到推理速度出現了輕微的減慢。

刪除模型並再次刷一下視訊記憶體。

del model
del pipe

flush()

然後，我們看下 4 位元量化的 GPU 視訊記憶體消耗峰值是多少。可以用與之前相同的 API 將模型量化為 4 位元 - 這次引數設定為 load_in_4bit=True 而不是 load_in_8bit=True 。

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True, low_cpu_mem_usage=True, pad_token_id=0)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

輸出:

Here is a Python function that transforms bytes to Giga bytes:\n\n```\ndef bytes_to_gigabytes(bytes):\n return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single argument

輸出幾乎與以前相同 - 只是在程式碼片段之前缺了 python 這個詞。我們看下需要多少視訊記憶體。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

輸出:

9.543574333190918

僅需 9.5GB！對於引數量大於 150 億的模型來說，確實不算多。

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

另請注意，與 8 位元量化相比，其推理速度會更慢一些，這是由於 4 位元量化使用了更激進的量化方法，導致 \(\text{quantize}\) 和 \(\text {dequantize}\) 在推理過程中花的時間更長。

del model
del pipe

flush()

總的來說，我們發現以 8 位元精度執行 OctoCoder 將所需的 GPU 視訊記憶體 從 32GB 減少到僅 15GB，而以 4 位元精度執行模型則進一步將所需的 GPU 視訊記憶體減少到 9GB 多一點。

4 位元量化讓模型可以在 RTX3090、V100 和 T4 等大多數人都可以輕鬆獲取的 GPU 上執行。

更多有關量化的資訊以及有關如何量化模型以使其視訊記憶體佔用比 4 位元更少，我們建議大家檢視 AutoGPTQ 的實現。

總結一下，重要的是要記住，模型量化會提高記憶體效率，但會犧牲準確性，在某些情況下還會犧牲推理時間。

如果 GPU 視訊記憶體對你而言不是問題，通常不需要考慮量化。然而，如果不量化，許多 GPU 根本無法執行 LLM，在這種情況下，4 位元和 8 位元量化方案是非常有用的工具。

更詳細的使用資訊，我們強烈建議你檢視 Transformers 的量化文件。

接下來，我們看看如何用更好的演算法和改進的模型架構來提高計算和記憶體效率。

2. Flash 注意力: 速度飛躍

當今表現最好的 LLM 其基本架構大體相似，包括前饋層、啟用層、層歸一化層以及最重要的自注意力層。

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

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

我們仔細分析一下。計算長度為 \(N\) 的輸入序列 \(\mathbf{X}\) 的自注意力層的輸出 \(\mathbf{O}\) ，其公式為:

\[\textbf{O} = \text{Attn}(\mathbf{X}) = \mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) \text{ ，其中 } \mathbf{Q} = \mathbf{W}_q \mathbf{X}, \mathbf{V} = \mathbf{W}_v \mathbf{X}, \mathbf{K} = \mathbf{W}_k \mathbf{X} \]

\(\mathbf{X} = (\mathbf{x} _1, … \mathbf{x}_ {N})\) 是注意力層的輸入序列。投影 \(\mathbf{Q}\) 和 \(\mathbf{K}\) 也是 \(N\) 個向量組成的序列，其乘積 \(\mathbf{QK}^T\) 的大小為 \(N^2\) 。

LLM 通常有多個注意力頭，因此可以並行進行多個自注意力計算。
假設 LLM 有 40 個注意力頭並以 bfloat16 精度執行，我們可以計算出儲存 $ \mathbf{QK^T}$ 矩陣的記憶體需求為 \(40 \times 2 \times N^2\) 位元組。當 \(N=1000\) 時僅需要大約 50MB 的視訊記憶體，但當 \(N=16000\) 時，我們需要 19GB 的視訊記憶體，當 \(N=100,000\) 時，僅儲存 \(\mathbf{QK}^T\) 矩陣就需要近 1TB。

總之，隨著輸入上下文越來越長，預設的自注意力演算法所需的記憶體很快就會變得非常昂貴。

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

我們如何擺脫長輸入文字對記憶體的過高要求？我們需要一種新的方法讓我們在計算自注意力機制時能夠擺脫 \(QK^T\) 矩陣。 Tri Dao 等人 開發了這樣一種新演算法，並將其稱為 Flash 注意力。

簡而言之，Flash 注意力將 \(\mathbf{V} \times \text{Softmax}(\mathbf{QK}^T)\) 的計算分解成若干步驟，透過迭代多個 softmax 計算步來將輸出分成多個較小的塊進行計算:

\[\textbf{O} _i \leftarrow s^a_ {ij} \times \textbf{O} _i + s^b_ {ij} \times \mathbf{V} _{j} \times \text{Softmax}(\mathbf{QK}^T_ {i,j}) \text{，在 } i, j \text{ 上迭代} \]

其中 \(s^a_{ij}\) 和 \(s^b_{ij}\) 是隨著每個 \(i\) 和 \(j\) 迭代更新的 softmax 統計歸一化值。

請注意，整個 Flash 注意力有點複雜，這裡已經大大簡化了。如果想要深入理解，可以閱讀 Flash Attention 的論文。

要點如下:

透過跟蹤 softmax 統計歸一化值再加上一些聰明的數學技巧，與預設的自注意力層相比，Flash 注意力的計算結果 完全相同，而記憶體成本僅隨著 \(N\) 線性增加。

僅看這個公式，直覺上來講，Flash 注意力肯定比預設的自注意力公式要慢很多，因為需要進行更多的計算。確實，與普通注意力相比，Flash 注意力需要更多的 FLOP，因為需要不斷重新計算 softmax 統計歸一化值 (如果感興趣，請參閱 論文 以瞭解更多詳細資訊)。

然而，與預設注意力相比，Flash 注意力的推理速度要快得多，這是因為它能夠顯著減少對較慢的高頻寬視訊記憶體的需求，而更多使用了更快的片上記憶體 (SRAM)。

從本質上講，Flash 注意力確保所有中間寫入和讀取操作都可以使用快速 片上 SRAM 來完成，而不必訪問較慢的視訊記憶體來計算輸出向量 \(\mathbf{O}\)。

實際上，如果能用的話，我們沒有理由不用 Flash 注意力。該演算法在數學上給出相同的輸出，但速度更快且記憶體效率更高。

我們看一個實際的例子。

我們的 OctoCoder 模型現在被輸入了長得多的提示，其中包括所謂的“系統提示”。系統提示用於引導 LLM 去適應特定的使用者任務。

接下來，我們使用系統提示，引導 OctoCoder 成為更好的程式設計助手。

system_prompt = """Below are a series of dialogues between various people and an AI technical assistant.
The assistant tries to be helpful, polite, honest, sophisticated, emotionally aware, and humble but knowledgeable.
The assistant is happy to help with code questions and will do their best to understand exactly what is needed.
It also tries to avoid giving false or misleading information, and it caveats when it isn't entirely sure about the right answer.
That said, the assistant is practical really does its best, and doesn't let caution get too much in the way of being useful.

The Starcoder models are a series of 15.5B parameter models trained on 80+ programming languages from The Stack (v1.2) (excluding opt-out requests).
The model uses Multi Query Attention, was trained using the Fill-in-the-Middle objective, and with 8,192 tokens context window for a trillion tokens of heavily deduplicated data.

-----

Question: Write a function that takes two lists and returns a list that has alternating elements from each input list.

Answer: Sure. Here is a function that does that.

def alternating(list1, list2):
   results = []
   for i in range(len(list1)):
       results.append(list1[i])
       results.append(list2[i])
   return results

Question: Can you write some test cases for this function?

Answer: Sure, here are some tests.

assert alternating([10, 20, 30], [1, 2, 3]) == [10, 1, 20, 2, 30, 3]
assert alternating([True, False], [4, 5]) == [True, 4, False, 5]
assert alternating([], []) == []

Question: Modify the function so that it returns all input elements when the lists have uneven length. The elements from the longer list should be at the end.

Answer: Here is the modified function.

def alternating(list1, list2):
   results = []
   for i in range(min(len(list1), len(list2))):
       results.append(list1[i])
       results.append(list2[i])
   if len(list1) > len(list2):
       results.extend(list1[i+1:])
   else:
       results.extend(list2[i+1:])
   return results

-----
"""

為了演示需要，我們將系統提示覆制十倍，以便輸入長度足夠長以觀察 Flash 注意力帶來的記憶體節省。然後在其後加上原始提示 "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here" :

long_prompt = 10 * system_prompt + prompt

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

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", torch_dtype=torch.bfloat16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

現在，我們可以像以前一樣執行模型，同時測量其峰值 GPU 視訊記憶體需求及推理時間。

import time

start_time = time.time()
result = pipe(long_prompt, max_new_tokens=60)[0]["generated_text"][len(long_prompt):]

print(f"Generated in {time.time() - start_time} seconds.")
result

輸出:

Generated in 10.96854019165039 seconds.
Sure. Here is a function that does that.\n\ndef bytes_to_giga(bytes):\n return bytes / 1024 / 1024 / 1024\n\nAnswer: Sure. Here is a function that does that.\n\ndef

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

注意，在實際應用中，系統提示不應重複十次 —— 一次就夠了！

我們測量一下峰值 GPU 視訊記憶體需求。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

輸出:

37.668193340301514

正如我們所看到的，峰值 GPU 視訊記憶體需求現在明顯高於以前，這主要是因為輸入序列變長了。整個生成過程也需要一分多鐘的時間。

我們呼叫 flush() 來釋放 GPU 記憶體以供下一個實驗使用。

flush()

為便於比較，我們執行相同的函式，但啟用 Flash 注意力。
為此，我們將模型轉換為 BetterTransformers，這會因此而啟用 PyTorch 的 SDPA 自注意力，其實現是基於 Flash 注意力的。

model.to_bettertransformer()

現在我們執行與之前完全相同的程式碼片段，但此時 Transformers 在底層將使用 Flash 注意力。

start_time = time.time()
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    result = pipe(long_prompt, max_new_tokens=60)[0]["generated_text"][len(long_prompt):]

print(f"Generated in {time.time() - start_time} seconds.")
result

輸出:

Generated in 3.0211617946624756 seconds.
 Sure. Here is a function that does that.\n\ndef bytes_to_giga(bytes):\n return bytes / 1024 / 1024 / 1024\n\nAnswer: Sure. Here is a function that does that.\n\ndef

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

我們最後一次測量一下記憶體消耗。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

輸出:

32.617331981658936

我們幾乎一下就回到了原來的 29GB 峰值 GPU 視訊記憶體。

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

flush()

3. 架構背後的科學: 長文字輸入和聊天式 LLM 的策略選擇

到目前為止，我們已經研究了透過以下方式提高計算和記憶體效率:

  • 將權重轉換為較低精度的格式
  • 用記憶體和計算效率更高的版本替換自注意力演算法

現在讓我們看看如何改變 LLM 的架構，使其對於需要長文字輸入的任務更高效， 例如 :

  • 檢索增強問答
  • 總結
  • 聊天

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

一旦經過訓練，LLM 的基本架構就很難改變，因此提前考慮 LLM 的任務特徵並相應地最佳化模型架構非常重要。模型架構中有兩個重要元件很快就會成為長輸入序列的記憶體和/或效能瓶頸。

  • 位置嵌入 (positional embeddings)
  • 鍵值快取 (key-value cache)

我們來一一詳細探討:

3.1 改進 LLM 的位置嵌入

自注意力機制計算每個詞元間的相關係數。例如，文字輸入序列 “Hello”, “I”, “love”, “you” 的 \(\text{Softmax}(\mathbf{QK}^T)\) 矩陣看起來如下:

每個詞元都會被賦予一個機率值，表示其對另一個詞元的關注度。例如， “love” 這個詞關注 “Hello” 這個詞的機率為 0.05%，關注 “I” 的機率為 0.3%，而對自己的關注機率則為 0.65%。

基於自注意力但沒有位置嵌入的 LLM 在理解輸入文字彼此的相對位置上會遇到很大困難。這是因為在經由 \(\mathbf{QK}^T\) 來計算相關機率時，其計算是與詞元間的相對距離無關的，即該計算與詞元間的相對距離的關係為 \(O(1)\)。因此，對於沒有位置嵌入的 LLM，每個詞元似乎與所有其他詞元等距。 此時 ，區分 “Hello I love you” 和 “You love I hello” 會比較困難。

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

Attention Is All You Need 論文引入了正弦位置嵌入 $\mathbf{P} = \mathbf{p}_1, \ldots, \mathbf{p}_N $。其中每個向量 \(\mathbf{p}_i\) 為其位置 \(i\) 的正弦函式。然後將位置編碼與輸入序列向量簡單相加 \(\mathbf{\hat{X}} = \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N\) = \(\mathbf{x}_1 + \mathbf{p}_1, \ldots, \mathbf{x}_N + \mathbf{p}_N\) 從而提示模型更好地學習語序。

其他工作 (如 Devlin 等人的工作) 沒有使用固定位置嵌入，而是使用可訓練的位置編碼，在訓練期間學習位置嵌入 \(\mathbf{P}\)。

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

  1. 正弦位置嵌入以及可訓練位置嵌入都是絕對位置嵌入， 即 為每個位置 id ($ 0, \ldots, N$) 生成一個唯一的嵌入。正如 Huang et al. 和 Su et al. 的工作所示，絕對位置嵌入會導致 LLM 在處理長文字輸入時效能較差。對長文字輸入而言，如果模型能夠學習輸入詞元間的相對距離而不是它們的絕對位置，會比較好。

  2. 當使用訓練位置嵌入時，LLM 必須在固定的輸入長度 \(N\)上進行訓練，因此如果推理時的輸入長度比訓練長度更長，外插會比較麻煩。

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

RoPE 和 ALiBi 都認為，最好直接在自注意力演算法中向 LLM 提示語序，因為詞元是透過自注意力機制互相關聯的。更具體地說，應該透過修改 \(\mathbf{QK}^T\) 的計算來提示語序。

簡而言之， RoPE 指出位置資訊可以編碼為 查詢 - 鍵值對 ， 如 \(\mathbf{q}_i\) 和 \(\mathbf{x}_j\) 透過分別將每個向量根據其在句子中的位置 \(i, j\) 旋轉角度 \(\theta \times i\) 和 \(\theta \times j\):

\[\mathbf{\hat{q}}_i^T \mathbf{\hat{x}}_j = \mathbf{{q}} _i^T \mathbf{R}_ {\theta, i -j} \mathbf{{x}}_j. \]

\(\mathbf{R}_{\theta, i - j}\) 表示旋轉矩陣。 $ \theta$ 在不可訓練的預定義值，其值取決於訓練期間最大輸入序列長度。

透過這樣做，\(\mathbf{q}_i\) 和 \(\mathbf{q}_j\) 之間的機率得分僅受 \(i \ne j\) 是否成立這一條件影響，且其值僅取決於相對距離 \(i - j\)，而與每個向量的具體位置 \(i\) 和 \(j\) 無關。

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

另一個方案是 ALiBi ， 它提出了一種更簡單的相對位置編碼方案。在計算 softmax 之前，\(\mathbf{QK}^T\) 矩陣的每個元素會減去被一個預定義係數 m 縮放後的對應兩個向量間的相對距離。

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

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

  • MPT
  • BLOOM

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

RoPE 和 ALiBi 都是相對位置嵌入，其嵌入引數是 不可 訓練的，而是基於以下直覺:

  • 有關輸入文字的位置提示應直接提供給自注意力層的 \(QK^T\) 矩陣
  • 應該激勵 LLM 學習基於恆定 相對 距離的位置編碼
  • 輸入詞元間彼此距離越遠，它們的 查詢 - 鍵 機率越低。 RoPE 和 ALiBi 都降低了距離較遠詞元間的 查詢 - 鍵 機率。RoPE 透過增加 查詢 - 鍵 向量之間的夾角來減少它們的向量積。而 ALiBi 透過從向量積中減去一個更大的數來達成這個目的。

總之，打算部署在需要處理長文字輸入的任務中的 LLM 可以透過相對位置嵌入 (例如 RoPE 和 ALiBi) 來進行更好的訓練。另請注意，使用了 RoPE 和 ALiBi 的 LLM 即使是僅在固定長度 (例如 $ N_1 = 2048$) 上訓練的，其仍然可以在推理時透過位置嵌入外插來處理比 \(N_1\) 長得多的文字輸入 (如 \(N_2 = 8192 > N_1\))。

3.2 鍵值快取

使用 LLM 進行自迴歸文字生成的工作原理是把輸入序列輸入給模型，並取樣獲得下一個詞元，再將獲得的詞元新增到輸入序列後面，如此往復，直到 LLM 生成一個表示結束的詞元。

請查閱 Transformer 的文字生成教程 以更直觀地瞭解自迴歸生成的工作原理。

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

input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to("cuda")

for _ in range(5):
  next_logits = model(input_ids)["logits"][:, -1:]
  next_token_id = torch.argmax(next_logits,dim=-1)

  input_ids = torch.cat([input_ids, next_token_id], dim=-1)
  print("shape of input_ids", input_ids.shape)

generated_text = tokenizer.batch_decode(input_ids[:, -5:])
generated_text

輸出:

shape of input_ids torch.Size([1, 21])
shape of input_ids torch.Size([1, 22])
shape of input_ids torch.Size([1, 23])
shape of input_ids torch.Size([1, 24])
shape of input_ids torch.Size([1, 25])
[' Here is a Python function']

正如我們所看到的，每次我們都把剛剛取樣出的詞元新增到輸入文字中。

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

因此，當前詞元 永遠僅 依賴於其前面的詞元，更具體地說，\(\mathbf{q} _i\) 向量永遠與任何 \(j > i\) 的鍵、值向量無關聯。相反 \(\mathbf{q} _i\) 僅關注其之前的鍵、值向量 \(\mathbf{k}_ {m < i}, \mathbf{v}_ {m < i} \text{，} m \in {0, \ldots i - 1}\)。為了減少不必要的計算，因此可以把先前所有步的每一層的鍵、值向量快取下來。

接下來，我們將告訴 LLM 在每次前向傳播中都利用鍵值快取來減少計算量。在 Transformers 中，我們可以透過將 use_cache 引數傳給 forward 來利用鍵值快取，這樣的話，每次推理僅需傳當前詞元給 forward 就可以。

past_key_values = None # past_key_values is the key-value cache
generated_tokens = []
next_token_id = tokenizer(prompt, return_tensors="pt")["input_ids"].to("cuda")

for _ in range(5):
  next_logits, past_key_values = model(next_token_id, past_key_values=past_key_values, use_cache=True).to_tuple()
  next_logits = next_logits[:, -1:]
  next_token_id = torch.argmax(next_logits, dim=-1)

  print("shape of input_ids", input_ids.shape)
  print("length of key-value cache", len(past_key_values[0][0])) # past_key_values are of shape [num_layers, 0 for k, 1 for v, batch_size, length, hidden_dim]
  generated_tokens.append(next_token_id.item())

generated_text = tokenizer.batch_decode(generated_tokens)
generated_text

輸出:

shape of input_ids torch.Size([1, 20])
length of key-value cache 20
shape of input_ids torch.Size([1, 20])
length of key-value cache 21
shape of input_ids torch.Size([1, 20])
length of key-value cache 22
shape of input_ids torch.Size([1, 20])
length of key-value cache 23
shape of input_ids torch.Size([1, 20])
length of key-value cache 24
[' Here', ' is', ' a', ' Python', ' function']

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

利用鍵值快取意味著 \(\mathbf{QK}^T\) 本質上減少為 \(\mathbf{q}_c\mathbf{K}^T\)，其中 \(\mathbf{q}_c\) 是當前輸入詞元的查詢投影，它 始終 只是單個向量。

使用鍵值快取有兩個優點:

  • 與計算完整的 \(\mathbf{QK}^T\) 矩陣相比，計算量更小，計算效率顯著提高，因此推理速度也隨之提高。
  • 所需的最大記憶體不隨生成的詞元數量呈二次方增加，而僅呈線性增加。

使用者應該 始終 使用鍵值快取，因為它的生成結果相同且能顯著加快長輸入序列的生成速度。當使用文字 pipeline 或 generate 方法 時，Transformers 預設啟用鍵值快取。

請注意，鍵值快取對於聊天等需要多輪自迴歸解碼的應用程式特別有用。我們看一個例子。

User: How many people live in France?
Assistant: Roughly 75 million people live in France
User: And how many are in Germany?
Assistant: Germany has ca. 81 million inhabitants

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

  1. 第一次，鍵值快取為空，輸入提示為 "User: How many people live in France?" ，模型自迴歸生成文字 "Roughly 75 million people live in France" ，同時在每個解碼步新增鍵值快取。
  2. 第二次輸入提示為 "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many in Germany?" 。由於快取，前兩個句子的所有鍵值向量都已經計算出來。因此輸入提示僅包含 "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?" 的編碼向量序列。

這裡需要注意兩件事:

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

然而，還有一個問題。雖然 \(\mathbf{QK}^T\) 矩陣所需的峰值記憶體顯著減少，但對於長輸入序列或多輪聊天，將鍵值快取保留在記憶體中還是會非常昂貴。請記住，鍵值快取需要儲存先前所有輸入向量 \(\mathbf{x}_i \text{, for } i \in {1, \ldots, c - 1}\) 的所有層、所有注意力頭的鍵值向量。

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

config = model.config
2 * 16_000 * config.n_layer * config.n_head * config.n_embd // config.n_head

輸出:

7864320000

大約 80 億個浮點數！以 float16 精度儲存 80 億個浮點值需要大約 15 GB 的視訊記憶體，大約是模型本身權重的一半！

研究人員提出了兩種方法，用於顯著降低鍵值快取的記憶體成本:

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

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

    透過共享同一對鍵、值投影權重，鍵值向量 \(\mathbf{k}_i, \mathbf{v}_i\) 在所有注意力頭上相同，這意味著我們只需要快取 1 個鍵值投影對，而不需要 n_head 對。

    由於大多數 LLM 有 20 到 100 個注意力頭，MQA 顯著減少了鍵值快取的記憶體消耗。因此，對於本文中使用的 LLM，假設輸入序列長度為 16000，其所需的記憶體消耗從 15 GB 減少到不到 400 MB。

    除了節省記憶體之外，MQA 還可以提高計算效率。在自迴歸解碼中，需要重新載入大的鍵值向量，與當前的鍵值向量對相串接，然後將其輸入到每一步的 \(\mathbf{q}_c\mathbf{K}^T\) 計算中。對於自迴歸解碼，不斷重新載入所需的記憶體頻寬可能成為嚴重的效能瓶頸。透過減少鍵值向量的大小，需要訪問的記憶體更少，從而減少記憶體頻寬瓶頸。欲瞭解更多詳細資訊，請檢視 Noam 的論文。

    這裡的重點是，只有使用鍵值快取時，將鍵值注意力頭的數量減少到 1 才有意義。沒有鍵值快取時，模型單次前向傳播的峰值記憶體消耗保持不變，因為每個注意力頭查詢向量不同，因此每個注意力頭的 \(\mathbf{QK}^T\) 矩陣也不相同。

    MQA 已被社群廣泛採用，現已被許多流行的 LLM 所採用:

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

  2. 分組查詢注意力 (Grouped-Query-Attention，GQA)

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

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

    GQA 最近才被提出，這就是為什麼截至本文撰寫時其被採用得較少。GQA 最著名的應用是 Llama-v2。

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

總結

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

GPT3/4、Llama-2-70b、Claude、PaLM 等海量 LLM 能夠在 Hugging Face Chat 或 ChatGPT 等聊天應用中快速執行的原因是很大一部分歸功於上述精度、演算法和架構方面的改進。展望未來，GPU、TPU 等加速器只會變得更快且記憶體更大，但人們仍然應該始終確保使用最好的可用演算法和架構來獲得最大的收益 ?。

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

原文作者: Patrick von Platen

譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。

審校/排版: zhongdongy (阿東)

相關文章