在 AMD GPU 上使用 TGI

TGI 已在 AMD Instinct MI210、MI250 與 MI300 GPU 上進行了支援與測試。未來可能會擴大支援範圍。建議透過 Docker 使用。請務必查看關於如何在 AMD GPU 上使用 Docker 的 AMD 文件。

在搭載 AMD GPU 的伺服器上，可以使用以下指令啟動 TGI：

model=teknium/OpenHermes-2.5-Mistral-7B
volume=$PWD/data # share a volume with the Docker container to avoid downloading weights every run

docker run --rm -it --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
    --device=/dev/kfd --device=/dev/dri --group-add video \
    --ipc=host --shm-size 256g --net host -v $volume:/data \
    ghcr.io/huggingface/text-generation-inference:3.3.5-rocm \
    --model-id $model

TGI 伺服器啟動後即可供用戶端查詢，請務必參閱「使用 TGI」指南。

TunableOp

TGI 的 AMD GPU Docker 映像檔整合了 PyTorch 的 TunableOp，它允許進行額外的預熱（warmup），以便從 rocBLAS 或 hipBLASLt 中選出效能最佳的矩陣乘法（GEMM）核心。

根據實驗，在 MI300X 上使用 ROCm 6.1 和 PyTorch 2.3 並啟用 TunableOp 時，我們觀察到延遲改善了 6-8%。

TunableOp 預設為啟用狀態，預熱過程可能需要 1-2 分鐘。如果您希望停用 TunableOp，請在啟動 TGI 的 Docker 容器時傳入 --env PYTORCH_TUNABLEOP_ENABLED="0"。

Flash Attention 實作

ROCm 提供了兩種 Flash Attention 實作方式，第一種是基於 Composable Kernel (CK) 的 ROCm/flash-attention，第二種則是 Triton 實作。

預設使用 Composable Kernel 實作。不過，Triton 實作在 MI250 和 MI300 上的延遲略低，但它需要進行預熱，且由於每個新的 Prompt 長度都需要重新預熱，這可能會造成負擔。如有需要，可在啟動 TGI Docker 容器時使用 --env ROCM_USE_FLASH_ATTN_V2_TRITON="0" 來停用 FA Triton 實作。

自訂 PagedAttention

為了在 ROCm 上獲得更好的效能，我們提供了一個自訂的 Paged Attention 核心，且預設為啟用。若要停用它並改回使用 PagedAttention v2 核心，請設定環境變數 ROCM_USE_CUSTOM_PAGED_ATTN=0。

此自訂核心支援 bf16 和 fp16 資料類型、16 的區塊大小（block size）、128 的頭大小（head size）、最大 16k 的上下文長度（context length），以及 1 到 16 之間的 GQA 比率。對於其他設定，我們將使用 PagedAttention v2 核心。

未支援的功能

ROCm 版本的 TGI 目前不支援以下功能，未來可能會擴大支援：

• 載入 AWQ 檢查點（checkpoints）。
• 滑動視窗注意力機制（Sliding Window Attention）的核心（適用於 Mistral）。