GPU

GPU 是機器學習的標準硬體，因為它們針對記憶體頻寬和並行性進行了最佳化。隨著現代模型規模的不斷擴大，確保 GPU 能夠高效處理並提供最佳效能變得前所未有的重要。

本指南將演示幾種最佳化 GPU 推理的方法。下面所示的最佳化方法可以相互結合，以實現更好的效能，並且它們也適用於分散式 GPU。

bitsandbytes

bitsandbytes 是一個支援 8 位和 4 位量化的量化庫。量化表示與原始全精度格式相比，權重具有較低的精度。它減少了記憶體需求，並使大型模型更容易適應記憶體。

首先確保安裝了 bitsandbytes 和 Accelerate。

pip install bitsandbytes accelerate

8位

4位

Optimum

Optimum 是一個 Hugging Face 庫，專注於最佳化各種硬體上的模型效能。它支援 ONNX Runtime (ORT)，這是一種模型加速器，適用於包括 NVIDIA GPU 和使用 ROCm 堆疊的 AMD GPU 在內的各種硬體和框架。

ORT 使用最佳化技術，將常見操作融合到單個節點中，並進行常量摺疊以減少計算次數。ORT 還將計算量最大的操作放在 GPU 上，其餘操作放在 CPU 上，以智慧地分配兩個裝置之間的工作負載。

Optimum 提供了 ORTModel 類用於載入 ONNX 模型。根據下表設定 provider 引數。

提供者	硬體
CUDAExecutionProvider	支援 CUDA 的 GPU
ROCMExecutionProvider	AMD Instinct、Radeon Pro、Radeon GPU
TensorrtExecutionProvider	TensorRT

例如，載入 distilbert/distilbert-base-uncased-finetuned-sst-2-english 檢查點進行序列分類。此檢查點包含一個 model.onnx 檔案。如果檢查點沒有 model.onnx 檔案，請設定 export=True 以動態將檢查點轉換為 ONNX 格式。

from optimum.onnxruntime import ORTModelForSequenceClassification

ort_model = ORTModelForSequenceClassification.from_pretrained(
  "distilbert/distilbert-base-uncased-finetuned-sst-2-english",
  #export=True,
  provider="CUDAExecutionProvider",
)

現在您可以在 Pipeline 中使用模型進行推理。

from optimum.pipelines import pipeline
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased-finetuned-sst-2-english")
pipeline = pipeline(task="text-classification", model=ort_model, tokenizer=tokenizer, device="cuda:0")
result = pipeline("Both the music and visual were astounding, not to mention the actors performance.")

在 NVIDIA GPU 上的加速推理和 AMD GPU 上的加速推理指南中瞭解更多關於將 ORT 與 Optimum 結合使用的詳細資訊。

BetterTransformer

BetterTransformer 是直接在 GPU 等硬體級別上執行專門的 Transformers 函式的 *快速路徑*。快速路徑執行主要包含兩個部分。

將多個操作融合到一個核心中，以實現更快、更高效的執行
透過巢狀張量跳過填充令牌的不必要計算

一些 BetterTransformer 功能正在被引入 Transformers，預設支援原生 torch.nn.functional.scaled_dot_product_attention (SDPA)。BetterTransformer 比 Transformers SDPA 整合具有更廣泛的覆蓋範圍，但您可以預期越來越多的架構將原生支援 Transformers 中的 SDPA。

BetterTransformer 透過 Optimum 提供，使用 to_bettertransformer()。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("bigscience/bloom")
model = model.to_bettertransformer()

呼叫 reverse_bettertransformer() 並先儲存它，以將模型恢復到原始的 Transformers 模型。

model = model.reverse_bettertransformer()
model.save_pretrained("saved_model")

請參閱 PyTorch 2.0 中 🤗 解碼器模型的開箱即用加速和記憶體節省中的 BetterTransformer 和縮放點積注意力效能基準。如果您有興趣瞭解更多，BetterTransformer 部落格文章也更詳細地討論了快速路徑執行。

縮放點積注意力 (SDPA)

PyTorch 的 torch.nn.functional.scaled_dot_product_attention (SDPA) 是縮放點積注意力機制的原生實現。SDPA 是 transformer 模型中使用的注意力機制的一種更高效、更最佳化的版本。

有三種受支援的實現可用。

FlashAttention2 僅支援 fp16 或 bf16 torch 型別的模型。請務必先將模型轉換為適當的型別。
xFormers 或記憶體高效注意力能夠支援 fp32 torch 型別的模型。
縮放點積注意力的 C++ 實現

當有可用的實現時，SDPA 預設用於 PyTorch v2.1.1 及更高版本。您可以透過在 from_pretrained() 中設定 attn_implementation="sdpa" 來顯式啟用 SDPA。某些注意力引數，例如 head_mask 和 output_attentions=True，不受支援，並會返回警告，指出 Transformers 將回退到（較慢的）eager 實現。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B", device_map="auto", attn_implementation="sdpa")

SDPA 選擇可用的效能最佳的實現，但您也可以使用 torch.nn.attention.sdpa_kernel 作為上下文管理器顯式選擇實現。下面的示例展示瞭如何使用 enable_flash=True 啟用 FlashAttention2 實現。

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

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B", device_map="auto").to("cuda")

input_text = "Hello, my llama is cute"
inputs = tokenizer(input_text, return_tensors="pt").to("cuda")

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

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

如果您遇到以下 RuntimeError，請嘗試安裝 PyTorch 的每夜版本，它對 FlashAttention 具有更廣泛的覆蓋範圍。

RuntimeError: No available kernel. Aborting execution.

pip3 install -U --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu118

FlashAttention

FlashAttention 也可作為獨立包使用。它可以透過以下方式顯著加速推理：

額外地對序列長度上的注意力計算進行並行化
在 GPU 執行緒之間劃分工作，以減少它們之間的通訊和共享記憶體讀寫

首先為您正在使用的硬體安裝 FlashAttention。

NVIDIA

AMD

透過在 from_pretrained() 中設定 attn_implementation="flash_attention_2" 來啟用 FlashAttention2。FlashAttention2 僅支援 fp16 或 bf16 torch 型別的模型。請務必先將模型轉換為適當的資料型別。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B", device_map="auto", torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2")

基準測試

FlashAttention2 顯著加速了推理，特別是對於長序列的輸入。然而，由於 FlashAttention2 不支援使用填充令牌計算注意力分數，如果序列包含填充令牌，則必須手動對批次推理的注意力分數進行填充和取消填充。缺點是使用填充令牌進行批次生成速度較慢。

短序列長度

長序列長度

為了避免這種減速，在訓練期間使用 FlashAttention2，序列中不帶填充令牌。打包資料集或連線序列直到達到最大序列長度。

tiiuae/falcon-7b

meta-llama/Llama-7b-hf

< > 在 GitHub 上更新