最佳化LLM的速度和記憶體

像GPT3/4、Falcon 和 Llama 這樣的大型語言模型 (LLM) 在處理以人類為中心的任務方面正在迅速發展，並已成為現代知識型產業的重要工具。然而，在實際任務中部署這些模型仍然具有挑戰性。

• 為了展現接近人類的文字理解和生成能力，LLM 目前需要由數十億個引數組成 (參見 Kaplan et al, Wei et. al)。這必然會增加推理的記憶體需求。
• 在許多實際任務中，LLM 需要被賦予大量的上下文資訊。這就要求模型能夠在推理過程中管理非常長的輸入序列。

這些挑戰的核心在於增強LLM的計算和記憶體能力，尤其是在處理大量輸入序列時。

在本指南中，我們將介紹高效部署LLM的有效技術。

1. 降低精度： 研究表明，以降低的數值精度（即8位和4位）操作可以在不顯著降低模型效能的情況下獲得計算優勢。

2. Flash Attention： Flash Attention 是注意力演算法的一種變體，它不僅提供了更節省記憶體的方法，而且由於優化了GPU記憶體利用率而提高了效率。

3. 架構創新： 考慮到LLM在推理期間總是以相同的方式部署，即具有長輸入上下文的自迴歸文字生成，因此已經提出了專門的模型架構，以實現更高效的推理。其中最重要的模型架構進步是Alibi、Rotary embeddings、多查詢注意力 (MQA) 和 分組查詢注意力 (GQA)。

在本指南中，我們將從張量的角度對自迴歸生成進行分析。我們將深入探討採用低精度的優缺點，全面探索最新的注意力演算法，並討論改進的LLM架構。在此過程中，我們將執行實際示例來展示每個功能改進。

1. 降低精度

LLM 的記憶體需求可以透過將 LLM 視為一組權重矩陣和向量，並將文字輸入視為一系列向量來最好地理解。下文中的術語“權重”將用於表示所有模型權重矩陣和向量。

截至本指南撰寫之時，LLM 至少包含數十億個引數。每個引數都由一個十進位制陣列成，例如 `4.5689`，通常以 float32、bfloat16 或 float16 格式儲存。這使我們能夠輕鬆計算將 LLM 載入到記憶體所需的記憶體量。

以 float32 精度載入具有 X 億引數的模型權重，大約需要 4 * X GB 的 VRAM。

然而，如今模型很少以完整的 float32 精度進行訓練，而是通常以 bfloat16 精度或較少見的 float16 精度進行訓練。因此，經驗法則變為：

載入具有 X 億引數的模型權重，大約需要 2 * X GB 的 bfloat16/float16 精度 VRAM。

對於較短的文字輸入（小於 1024 個 token），推理的記憶體需求主要由載入權重的記憶體需求決定。因此，目前我們假設推理的記憶體需求等於將模型載入到 GPU VRAM 的記憶體需求。

以下是一些以 bfloat16 格式載入模型所需的大致 VRAM 量示例：

• GPT3 需要 2 * 175 GB = 350 GB VRAM
• Bloom 需要 2 * 176 GB = 352 GB VRAM
• Llama-2-70b 需要 2 * 70 GB = 140 GB VRAM
• Falcon-40b 需要 2 * 40 GB = 80 GB VRAM
• MPT-30b 需要 2 * 30 GB = 60 GB VRAM
• bigcode/starcoder 需要 2 * 15.5 = 31 GB VRAM

截至本文撰寫之時，市場上最大的 GPU 晶片是 A100 和 H100，提供 80GB 的 VRAM。上面列出的大多數模型僅載入就需要超過 80GB，因此必然需要張量並行和/或流水線並行。

🤗 Transformers 現在支援對在其各自配置類中具有 `base_tp_plan` 的受支援模型進行張量並行。在此處瞭解更多關於張量並行的資訊。此外，如果您有興趣以張量並行友好的方式編寫模型，請隨時檢視text-generation-inference 庫。

樸素的流水線並行是開箱即用的。為此，只需使用 `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，因為它可以在單個 40 GB A100 GPU 裝置晶片上執行。請注意，我們接下來將應用的所有記憶體和速度最佳化同樣適用於需要模型並行或張量並行的模型。

由於模型以 bfloat16 精度載入，根據我們上面的經驗法則，我們預計使用 `bigcode/octocoder` 執行推理所需的記憶體約為 31 GB VRAM。讓我們試一試。

我們首先載入模型和分詞器，然後將兩者傳遞給 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 精度執行模型，將需要高達 64 GB 的 VRAM。

現在幾乎所有模型都以 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，它考慮了各種硬體後端，如 XPU、MLU、NPU、MPS 等。

from accelerate.utils import release_memory
# ...

release_memory(model)

如果你的 GPU 沒有 32 GB 的 VRAM 怎麼辦？研究發現，模型權重可以量化為 8 位或 4 位，而不會顯著損失效能（參見 Dettmers et al.）。甚至可以將模型量化到 3 位或 2 位，效能損失在可接受範圍內，如最近的 GPTQ 論文 🤯 所示。

不深入過多細節，量化方案旨在降低權重的精度，同時儘量保持模型的推理結果儘可能準確（即儘可能接近 bfloat16）。請注意，量化對於文字生成尤其有效，因為我們只關心選擇 *最有可能的下一個 token 集*，而不關心下一個 token *logit* 分佈的確切值。重要的是下一個 token *logit* 分佈大致保持不變，這樣 `argmax` 或 `topk` 操作才能給出相同的結果。

有各種量化技術，我們在此不詳細討論，但總的來說，所有量化技術都按如下方式工作：

• 1. 將所有權重量化到目標精度
• 2. 載入量化後的權重，並以 bfloat16 精度傳遞輸入序列向量
• 3. 動態地將權重反量化為 bfloat16，以與 bfloat16 精度下的輸入向量執行計算

簡而言之，這意味著輸入-權重矩陣乘法，其中$X$是輸入，$W$是一個權重矩陣，以及$Y$是輸出$$Y = X * W$$

被更改為$$Y = X * \text{dequantize}(W)$$

對於每個矩陣乘法。反量化和再量化是順序執行的，適用於所有權重矩陣，因為輸入透過網路圖執行。

因此，使用量化權重時，推理時間通常不會減少，反而會增加。理論夠了，讓我們試一試！要使用 Transformers 量化權重，您需要確保已安裝 `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

顯著減少！我們現在只有略高於 15 GB，因此可以在像 4090 這樣的消費級 GPU 上執行此模型。我們看到了記憶體效率的顯著提升，並且模型輸出幾乎沒有退化。但是，我們也可以注意到推理速度略有下降。

我們刪除模型並再次重新整理記憶體。

del model
del pipe

flush()

讓我們看看 4 位量化會帶來怎樣的峰值 GPU 記憶體消耗。將模型量化到 4 位可以使用與之前相同的 API 完成——這次透過傳遞 `load_in_4bit=True` 而不是 `load_in_8bit=True`。

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=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 VRAM 從 32GB 減少到僅 15GB，而以 4 位精度執行模型則將所需的 GPU VRAM 進一步減少到僅 9GB 多一點。

4位量化允許模型在 RTX3090、V100 和 T4 等 GPU 上執行，這些 GPU 對大多數人來說都相當容易獲得。

有關量化的更多資訊以及如何量化模型以使所需的 GPU VRAM 少於 4 位，我們建議查閱 `AutoGPTQ` 實現。

總而言之，重要的是要記住模型量化以提高記憶體效率為代價，有時也會犧牲準確性和推理時間。

如果 GPU 記憶體對您的用例不是限制，通常沒有必要考慮量化。然而，許多 GPU 在沒有量化方法的情況下根本無法執行 LLM，在這種情況下，4 位和 8 位量化方案是非常有用的工具。

有關更詳細的使用資訊，我們強烈建議檢視 Transformers Quantization Docs。接下來，讓我們看看如何透過使用更好的演算法和改進的模型架構來提高計算和記憶體效率。

2. Flash Attention

當今效能最佳的 LLM 共享或多或少相同的基本架構，該架構由前饋層、啟用層、層歸一化層以及最關鍵的自注意力層組成。

自注意力層是大型語言模型 (LLM) 的核心，因為它們使模型能夠理解輸入 token 之間的上下文關係。然而，自注意力層的峰值 GPU 記憶體消耗在計算和記憶體複雜性上與輸入 token 數量（也稱為序列長度）呈二次增長，我們將其表示為$N$。雖然這對於較短的輸入序列（最多 1000 個輸入 token）並不明顯，但對於較長的輸入序列（大約 16000 個輸入 token）來說，它成為一個嚴重的問題。

讓我們仔細看看。計算自注意力層輸出$\mathbf{O}$的公式，對於長度為$\mathbf{X}$的輸入是：$\mathbf{O}$的輸入$\mathbf{X}$長度為$N$是$$\textbf{O} = \text{Attn}(\mathbf{X}) = \mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) \text{ with } \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 * 2 * N^2$位元組。對於$N=1000$只需要大約 50 MB 的 VRAM，然而，對於$N=16000$我們將需要 19 GB 的 VRAM，而對於$N=100,000$我們將需要將近 1TB 來儲存$\mathbf{QK}^T$矩陣。

長話短說，對於大型輸入上下文，預設的自注意力演算法很快就會變得記憶體開銷過大。

隨著 LLM 在文字理解和生成方面的改進，它們被應用於日益複雜的任務。模型曾經處理幾句話的翻譯或摘要，現在可以管理整個頁面，這需要處理大量輸入的能力。

我們如何擺脫對長輸入長度的過高記憶體要求？我們需要一種新的計算自注意力機制的方法，它能消除$QK^T$矩陣。Tri Dao 等人開發了這樣一種新演算法，並將其命名為Flash Attention。

簡而言之，Flash Attention 將$\mathbf{V} \times \text{Softmax}(\mathbf{QK}^T$) 計算拆開，透過迭代多個 Softmax 計算步驟來計算輸出的較小塊$$\textbf{O}_i \leftarrow s^a_{ij} * \textbf{O}_i + s^b_{ij} * \mathbf{V}_{j} \times \text{Softmax}(\mathbf{QK}^T_{i,j}) \text{ for multiple } i, j \text{ iterations}$$

，其中$s^a_{ij}$和$s^b_{ij}$是每次迭代都需要重新計算的 Softmax 歸一化統計量。$i$和$j$ .

請注意，整個 Flash Attention 要複雜一些，並且在此處進行了極大的簡化，因為深入探討超出了本指南的範圍。讀者可以參考寫得很好的Flash Attention 論文以獲取更多詳細資訊。

這裡的關鍵要點是：

透過跟蹤 Softmax 歸一化統計資料並運用一些巧妙的數學方法，Flash Attention 能夠產生與預設自注意力層數值上相同的輸出，而記憶體成本僅隨$N$ .

從公式上看，人們直觀地會認為 Flash Attention 必須比預設的自注意力公式慢得多，因為它需要進行更多的計算。確實，與普通注意力相比，Flash Attention 需要更多的 FLOPs，因為 Softmax 歸一化統計資料必須不斷重新計算（如果感興趣，請參閱論文以獲取更多詳細資訊）。

然而，Flash Attention 在推理方面比預設注意力快得多，這歸因於它能夠顯著減少對較慢、高頻寬 GPU 記憶體（VRAM）的需求，轉而專注於較快的片上記憶體（SRAM）。

本質上，Flash Attention 確保所有中間寫入和讀取操作都可以使用快速的片上SRAM 記憶體完成，而無需訪問較慢的 VRAM 記憶體來計算輸出向量$\mathbf{O}$ .

在實踐中，如果可用，目前絕對沒有理由不使用 Flash Attention。該演算法在數學上給出相同的輸出，並且速度更快，記憶體效率更高。

我們來看一個實際的例子。

我們的 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 Attention 的記憶體節省。我們附加原始文字提示 "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)

現在，讓我們像以前一樣執行模型，不使用 Flash Attention，並測量峰值 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 個 token 的截止。這並不奇怪，因為我們為了演示目的將系統提示重複了十次，從而提示模型重複自己。

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

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

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

輸出:

37.668193340301514

正如我們所看到的，現在峰值 GPU 記憶體需求顯著高於開始時，這主要是由於輸入序列更長。此外，生成現在需要一分多鐘。

我們呼叫 `flush()` 來釋放 GPU 記憶體，以便進行下一次實驗。

flush()

為了進行比較，讓我們執行相同的函式，但啟用 Flash Attention。為此，我們將模型轉換為BetterTransformer，從而啟用 PyTorch 的SDPA 自注意力，而後者又能夠使用 Flash Attention。

model.to_bettertransformer()

現在我們執行與之前完全相同的程式碼片段，Transformers 將在底層利用 Flash Attention。

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 Attention，我們可以觀察到非常顯著的速度提升。

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

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

輸出:

32.617331981658936

我們幾乎回到了最初的 29GB 峰值 GPU 記憶體。

我們可以觀察到，在使用 Flash Attention 傳遞非常長的輸入序列時，我們僅比開始時傳遞短輸入序列多使用了大約 100MB 的 GPU 記憶體。

flush()

有關如何使用 Flash Attention 的更多資訊，請參閱此文件頁面。

3. 架構創新

目前為止，我們已經探討了如何透過以下方式提高計算和記憶體效率：

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

現在讓我們看看如何改變 LLM 的架構，使其對需要長文字輸入的任務（例如：）最有效和高效。

• 檢索增強型問答，
• 摘要，
• 聊天

請注意，*聊天*不僅要求LLM能夠處理長文字輸入，還需要LLM能夠高效處理使用者和助手之間（例如 ChatGPT）的來回對話。

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

• 位置嵌入
• 鍵值快取

讓我們更詳細地介紹每個元件。

3.1 改進LLM的位置嵌入

自注意力將每個標記與其他標記關聯起來。例如，文字輸入序列“Hello”、“I”、“love”、“you”的$\text{Softmax}(\mathbf{QK}^T)$矩陣可能如下所示：

每個單詞標記都被賦予一個機率質量，表示它關注所有其他單詞標記的程度，因此與其他所有單詞標記關聯起來。例如，單詞“love”以 5% 的機率關注“Hello”，以 30% 的機率關注“I”，以 65% 的機率關注它自己。

一個基於自注意力但沒有位置嵌入的 LLM 在理解文字輸入彼此之間的位置時會遇到很大困難。這是因為由$\mathbf{QK}^T$將每個詞元與$O(1)$計算中的所有其他詞元關聯起來，無論它們之間的相對位置距離如何。因此，對於沒有位置嵌入的 LLM，每個詞元似乎與其他所有詞元具有相同的距離，例如，區分“Hello I love you”和“You love I hello”將非常具有挑戰性。$O(1)$計算，無論它們彼此之間的相對位置距離如何。因此，對於沒有位置嵌入的LLM，每個token似乎與其他所有token具有相同的距離，例如，區分“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 等人和Su 等人所示，絕對位置嵌入導致大型語言模型（LLM）在長文字輸入上的效能不佳。對於長文字輸入，如果模型學習輸入標記之間的相對位置距離而不是它們的絕對位置，則會更有優勢。
2. 當使用學習到的位置嵌入時，大型語言模型（LLM）必須在固定輸入長度上進行訓練$N$，這使得其難以推廣到比訓練時更長的輸入長度。

最近，能夠解決上述問題的相對位置嵌入變得越來越流行，其中最引人注目的是：

• 旋轉位置嵌入（RoPE）
• ALiBi

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

在此不深入探討過多細節，RoPE 指出位置資訊可以編碼到查詢-鍵對中，例如$\mathbf{q}_i$和$\mathbf{x}_j$透過將每個向量旋轉一個角度$\theta * i$和$\theta * j$分別與$i, 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$ .

RoPE 目前被用於許多最重要的LLM中，例如：

• Falcon
• Llama
• PaLM

作為替代方案，ALiBi 提出了一種更簡單的相對位置編碼方案。輸入詞元之間的相對距離，以一個由預定義值 m 縮放的負整數形式，在 softmax 計算之前直接新增到 matrix 的每個查詢-鍵條目中。$\mathbf{QK}^T$矩陣 right before the softmax computation.

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

ALiBi目前被用於許多最重要的LLM中，例如：

• MPT
• BLOOM

RoPE和ALiBi位置編碼都可以外推到訓練時未見的輸入長度，儘管已表明ALiBi的開箱即用外推效果遠優於RoPE。對於ALiBi，只需增加下三角位置矩陣的值以匹配輸入序列的長度即可。對於RoPE，如果保持與訓練時相同的$\theta$當輸入文字長度遠超訓練時所見長度時，這將導致糟糕的結果，參見Press 等人。然而，社群已經發現了一些有效的技巧來調整$\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$向量從不與任何鍵、值向量相關聯$\mathbf{k}_j, \mathbf{v}_j$如果$j > i$。相反$\mathbf{q}_i$只關注之前的鍵值向量$\mathbf{k}_{m < i}, \mathbf{v}_{m < i} \text{ , for } m \in \{0, \ldots i - 1\}$。為了減少不必要的計算，可以快取每個層的所有先前時間步的鍵值向量。

接下來，我們將告訴 LLM 使用鍵值快取，並在每次前向傳遞時檢索和轉發它。在 Transformers 中，我們可以透過將 `use_cache` 標誌傳遞給 `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", next_token_id.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, 1])
length of key-value cache 20
shape of input_ids torch.Size([1, 1])
length of key-value cache 21
shape of input_ids torch.Size([1, 1])
length of key-value cache 22
shape of input_ids torch.Size([1, 1])
length of key-value cache 23
shape of input_ids torch.Size([1, 1])
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$矩陣相比，執行的計算量更少。這導致推理速度的提高。
• 所需的峰值記憶體不會隨生成的詞元數量呈平方級增長，而只會呈線性增長。

始終應該使用鍵值快取，因為它可以帶來相同的結果，並且對於較長的輸入序列，顯著加快速度。在Transformers中，當使用文字管道或`generate`方法時，鍵值快取預設啟用。我們有一整篇關於快取的指南，請參見此處。

3.2.1 多輪對話

鍵值快取對於需要多次自迴歸解碼的應用程式（如聊天）特別有用。我們來看一個例子。

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. 第一次，鍵值快取為空，輸入提示為“使用者：法國有多少人口？”模型自迴歸生成文字“法國大約有7500萬人口”，並在每個解碼步驟增加鍵值快取。
2. 第二次，輸入提示是“使用者：法國有多少人口？\n助理：法國大約有7500萬人口。\n使用者：德國有多少人口？”。由於快取的存在，前兩句的所有鍵值向量都已計算。因此，輸入提示只包含“使用者：德國有多少人口？”。在處理縮短的輸入提示時，其計算出的鍵值向量會與第一次解碼的鍵值快取連線起來。然後，第二個助理的回答“德國大約有8100萬居民”將使用包含“使用者：法國有多少人口？\n助理：法國大約有7500萬人口。\n使用者：德國有多少人口？”編碼的鍵值向量的鍵值快取進行自迴歸生成。

此處應注意兩點

1. 對於部署在聊天中的大型語言模型（LLM），保留所有上下文至關重要，以便LLM理解對話的所有先前上下文。例如，對於上述例子，當用戶詢問“德國有多少人口”時，LLM需要理解使用者指的是人口。
2. 鍵值快取對於聊天非常有用，因為它允許我們持續增長編碼的聊天曆史，而無需從頭重新編碼聊天曆史（例如，在使用編碼器-解碼器架構時就會出現這種情況）。

在 `transformers` 中，當 `return_dict_in_generate=True` 被傳入（除了預設的 `use_cache=True` 之外）時，`generate` 呼叫將返回 `past_key_values`。請注意，它尚未透過 `pipeline` 介面提供。

# Generation as usual
prompt = system_prompt + "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(**model_inputs, max_new_tokens=60, return_dict_in_generate=True)
decoded_output = tokenizer.batch_decode(generation_output.sequences)[0]

# Piping the returned `past_key_values` to speed up the next conversation round
prompt = decoded_output + "\nQuestion: How can I modify the function above to return Mega bytes instead?\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(
  **model_inputs,
  past_key_values=generation_output.past_key_values,
  max_new_tokens=60,
  return_dict_in_generate=True
)
tokenizer.batch_decode(generation_output.sequences)[0][len(prompt):]

輸出:

 is a modified version of the function that returns Mega bytes instead.

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

Answer: The function takes a number of bytes as input and returns the number of

太棒了，沒有額外的時間用於重新計算注意力層的相同鍵和值！然而，這裡有一個問題。儘管所需的峰值記憶體$\mathbf{QK}^T$矩陣顯著減少，但對於長輸入序列或多輪對話，將鍵值快取儲存在記憶體中可能會非常耗記憶體。請記住，鍵值快取需要為所有先前的輸入向量儲存鍵值向量$\mathbf{x}_i \text{, for } i \in \{1, \ldots, c - 1\}$對於所有自注意力層和所有注意力頭。

我們來計算之前使用的 LLM `bigcode/octocoder` 的鍵值快取需要儲存的浮點值數量。浮點值數量是序列長度的兩倍，再乘以注意力頭數量，再乘以注意力頭維度，最後乘以層數。對於我們的 LLM，在假定輸入序列長度為 16000 的情況下，計算結果為：

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

輸出:

7864320000

大約80億個浮點值！以`float16`精度儲存80億個浮點值需要大約15GB的RAM，這大約是模型權重本身的一半！研究人員提出了兩種方法，可以顯著降低儲存鍵值快取的記憶體成本，這將在接下來的小節中進行探討。

3.2.2 多頭查詢注意力（MQA）

多查詢注意力由 Noam Shazeer 在其論文《快速 Transformer 解碼：一個寫頭足矣》中提出。正如標題所示，Noam 發現，無需使用 `n_head` 鍵值投影權重，而是可以使用一個在所有注意力頭之間共享的單頭值投影權重對，而模型的效能不會顯著下降。

透過使用單個頭值投影權重對，鍵值向量$\mathbf{k}_i, \mathbf{v}_i$在所有注意力頭中必須相同，這意味著我們只需在快取中儲存一對鍵值投影，而不是`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 所使用

• Falcon
• PaLM
• MPT
• BLOOM

此外，本筆記本中使用的檢查點——`bigcode/octocoder`——也使用了 MQA。

3.2.3 分組查詢注意力（GQA）

分組查詢注意力，由 Google 的 Ainslie 等人提出，發現使用 MQA 通常會導致質量下降，而普通的多鍵值頭投影則不然。該論文認為，透過不那麼大幅度地減少查詢頭投影權重，可以保持更好的模型效能。它建議使用 `n < n_head` 個鍵值投影權重，而不是僅使用單個鍵值投影權重。透過將 `n` 選擇為遠小於 `n_head` 的值，例如 2、4 或 8，幾乎可以保留 MQA 的所有記憶體和速度收益，同時犧牲更少的模型容量，從而可以說效能下降更少。

此外，GQA 的作者發現，現有模型檢查點可以透過僅使用原始預訓練計算量的 5% 進行升級訓練，以具備 GQA 架構。儘管原始預訓練計算量的 5% 仍然是一個巨大的數字，但 GQA 升級訓練 允許現有檢查點可用於更長的輸入序列。

GQA 是最近才提出的，因此在撰寫本筆記本時，其採用率較低。GQA 最顯著的應用是Llama-v2。

綜上所述，如果LLM部署在需要處理大型輸入序列的自迴歸解碼任務中（例如聊天），強烈建議使用GQA或MQA。

結論

研究界不斷提出新的巧妙方法來加速日益增長的LLM的推理時間。例如，一個有前途的研究方向是推測性解碼，其中“簡單詞元”由更小、更快的語言模型生成，只有“困難詞元”才由LLM本身生成。更詳細的討論超出了本筆記本的範圍，但可以在這篇精彩的部落格文章中閱讀。

像 GPT3/4、Llama-2-70b、Claude、PaLM 這樣的大型 LLM 之所以能在 Hugging Face Chat 或 ChatGPT 等聊天介面中如此迅速地執行，很大程度上得益於上述在精度、演算法和架構方面的改進。展望未來，GPU、TPU 等加速器只會變得更快並允許更多的記憶體，但我們仍然應該始終確保使用最佳可用演算法和架構，以獲得最大的效益 🤗