將 TGI 與 AMD GPU 結合使用

TGI 在 AMD Instinct MI210、MI250 和 MI300 GPU 上受到支援和測試。未來可能會擴充套件支援範圍。建議透過 Docker 使用。請務必查閱 AMD 文件，瞭解如何將 Docker 與 AMD GPU 結合使用。

在由 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.4-rocm \
    --model-id $model

啟動的 TGI 伺服器隨後可以從客戶端查詢，請務必檢視使用 TGI 指南。

TunableOp

TGI 的 AMD GPU Docker 映象集成了 PyTorch 的 TunableOp，它允許進行額外的預熱，以從 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 上略低，但需要預熱，這可能會帶來問題，因為它需要為每個新的提示長度再次進行。如果需要，可以在啟動 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，頭部大小為 128，最大上下文長度為 16k，GQA 比率在 1 到 16 之間。對於其他配置，我們使用 PagedAttention v2 核心。

不支援的功能

TGI 的 ROCm 版本目前不支援以下功能，未來可能會擴充套件支援範圍

• 載入 AWQ 檢查點。
• 滑動視窗注意力 (Mistral) 核心