最佳化推理

使用大型語言模型（LLM）進行推理可能具有挑戰性，因為它們需要儲存和處理數十億個引數。載入一個 70B 引數的 Llama 2 模型，需要 256GB 記憶體用於全精度權重，128GB 記憶體用於半精度權重。而目前最強大的 GPU（A100 和 H100）只有 80GB 記憶體。

除了記憶體需求之外，推理速度也很慢，因為 LLM 需要反覆呼叫以生成下一個 token。隨著生成過程的進行，輸入序列會不斷增長，處理時間也越來越長。

本指南將向您展示如何最佳化 LLM 推理以加速生成並減少記憶體使用。

嘗試使用 Text Generation Inference (TGI)，這是一個 Hugging Face 庫，專門用於部署和提供高度最佳化的 LLM 進行推理。

靜態 kv-快取和 torch.compile

LLM 會為每個輸入 token 計算鍵值（kv）對，並且每次都會執行相同的 kv 計算，因為生成的輸出成為輸入的一部分。然而，每次都執行相同的 kv 計算效率不高。

kv-快取會儲存過去的鍵和值，而不是每次都重新計算它們。因此，kv-快取是動態的，並且隨著每個生成步驟而增長，這使得您無法利用 torch.compile，這是一種將 PyTorch 程式碼融合到最佳化核心中的強大最佳化方法。

靜態 kv-快取透過將 kv-快取大小預分配到最大值來解決這個問題，這樣您就可以將其與 torch.compile 結合使用，從而將速度提高多達 4 倍。您的速度提升可能因模型大小（模型越大，速度提升越小）和硬體而異。

關注此問題以跟蹤哪些模型（Llama、Gemma、Mistral 等）支援靜態 kv-快取和 torch.compile。

根據您的任務，有幾種方法可以使用靜態 kv-快取。

對於基本用例，將 cache_implementation 設定為 "static"（推薦）。
對於多輪生成或自定義生成迴圈，請直接初始化和處理 StaticCache。
對於更獨特的硬體或用例，將整個 generate() 函式編譯成單個圖可能更好。

無論您如何使用靜態 kv-快取和 torch.compile，請使用 pad_to_multiple_of 將輸入左填充到有限的一組值，以避免與形狀相關的重新編譯。

1. 快取實現

2. StaticCache

3. 編譯整個生成函式

解碼策略

解碼也可以最佳化以加速生成。您可以使用一個輕量級的輔助模型來比 LLM 本身更快地生成候選 token，或者您可以使用這種解碼策略的一個變體，它在以輸入為基礎的任務中特別有效。

推測解碼

要了解更深入的解釋，請查閱輔助生成：低延遲文字生成的新方向部落格文章！

對於每個輸入 token，模型權重在前向傳播期間每次都會載入，這在模型擁有數十億引數時既慢又麻煩。推測解碼透過使用第二個更小、更快的輔助模型來生成候選 token，這些 token 在單個前向傳播中由更大的模型進行驗證，從而緩解了這種減速。如果驗證的 token 是正確的，LLM 基本可以“免費”獲得它們，而無需自己生成它們。準確性沒有下降，因為驗證前向傳播確保生成與 LLM 自己生成時相同的輸出。

為了獲得最大的加速，輔助模型應該比 LLM 小得多，以便它可以快速生成 token。輔助模型和 LLM 模型還必須共享相同的 tokenizer，以避免重新編碼和解碼 token。

推測解碼僅支援貪婪搜尋和取樣解碼策略，並且不支援批處理輸入。

透過載入輔助模型並將其傳遞給 generate() 來啟用推測解碼。

貪婪搜尋

取樣

提示查詢解碼

提示查詢解碼是推測解碼的一種變體，也相容貪婪搜尋和取樣。提示查詢在輸入接地任務（例如摘要）中特別有效，因為提示和輸出之間通常存在重疊詞。這些重疊的 N-gram 用作 LLM 候選 token。

要啟用提示查詢解碼，請在 prompt_lookup_num_tokens 引數中指定應重疊的 token 數量。然後將此引數傳遞給 generate()。

貪婪解碼

取樣

注意力機制

Transformer 模型的一個已知問題是，自注意力機制在計算和記憶體方面隨輸入 token 的數量呈二次增長。這一限制在處理更長序列的 LLM 中更為明顯。為了解決這個問題，可以嘗試 FlashAttention2 或 PyTorch 的縮放點積注意力 (SDPA)，它們是更節省記憶體的注意力實現。

FlashAttention-2

FlashAttention 和 FlashAttention-2 將注意力計算分解為更小的塊，並減少對 GPU 記憶體的中間讀/寫操作次數，以加速推理。FlashAttention-2 在原始 FlashAttention 演算法的基礎上進行了改進，透過在序列長度維度上並行化，並更好地劃分硬體上的工作以減少同步和通訊開銷。

要使用 FlashAttention-2，請在 from_pretrained() 中將 attn_implementation 設定為 "flash_attention_2"。

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

quant_config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2b",
    quantization_config=quant_config,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)

PyTorch 縮放點積注意力

縮放點積注意力 (SDPA) 在 PyTorch 2.0 中自動啟用，它支援 FlashAttention、xFormers 和 PyTorch 的 C++ 實現。如果您使用的是 CUDA 後端，SDPA 會選擇效能最佳的注意力演算法。對於其他後端，SDPA 預設為 PyTorch C++ 實現。

只要您安裝了最新版本的 PyTorch，SDPA 就會自動支援 FlashAttention-2。

使用 torch.nn.attention.sdpa_kernel 上下文管理器來顯式啟用或停用四種注意力演算法中的任何一種。例如，使用 SDPBackend.FLASH_ATTENTION 啟用 FlashAttention。

import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2b",
    torch_dtype=torch.bfloat16,
)

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    outputs = model.generate(**inputs)

量化

量化透過以較低精度儲存模型權重來減小模型大小。這會降低記憶體使用量，並在 GPU 記憶體受限時使 LLM 推理更易於訪問。

如果您不受 GPU 限制，則不必對模型進行量化，因為它可能會略微增加延遲（AWQ 和融合 AWQ 模組除外），因為需要額外的步驟來量化和反量化權重。

有許多量化庫可用（有關更多詳細資訊，請參閱量化指南），例如 Quanto、AQLM、VPTQ、AWQ 和 AutoGPTQ。歡迎嘗試它們，看看哪個最適合您的用例。我們還建議閱讀《🤗 Transformers 中原生支援的量化方案概述》部落格文章，其中比較了 AutoGPTQ 和 bitsandbytes。

使用下面的模型記憶體計算器估算和比較載入模型所需的記憶體。例如，嘗試估算載入 Mistral-7B-v0.1 所需的記憶體。

要以半精度載入模型，請在 from_pretrained() 中將 torch_dtype 引數設定為 torch.bfloat16。這需要 13.74GB 記憶體。

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1", torch_dtype=torch.bfloat16, device_map="auto",
)

要載入量化模型（8 位或 4 位），請嘗試 bitsandbytes 並將 load_in_4bit 或 load_in_8bit 引數設定為 True。以 8 位載入模型只需要 6.87 GB 記憶體。

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch

quant_config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1", quantization_config=quant_config, device_map="auto"
)

< > 在 GitHub 上更新