在 Google Cloud 上的首次 TPU 推論

本教學將引導您如何在 TPU 上使用文字生成推論 (TGI, Text Generation Inference) (文件) 來設定並執行推論。TGI 伺服器與 OpenAI 的訊息 API 相容，並為在 TPU 上提供模型服務提供了最佳化的解決方案。

先修條件

開始之前，請確保您已準備好：

一個執行中的 TPU 執行個體（參閱 TPU 設定指南）
具備連接至 TPU 執行個體的 SSH 存取權限
一個 HuggingFace 帳號

步驟 1：初始設定

SSH 存取

首先，透過 SSH 連線至您的 TPU 執行個體。

安裝必要工具

安裝 HuggingFace Hub CLI

pip install huggingface_hub

身份驗證

登入 HuggingFace

huggingface-cli login

步驟 2：模型部署

模型選擇

我們將在本教學中使用 gemma-2b-it 模型

造訪 https://huggingface.co/google/gemma-2b-it
接受模型使用條款與條件
這將啟用模型下載存取權限

啟動 TGI 伺服器

我們將使用 Optimum-TPU 映像檔，這是由 HuggingFace 提供的 TPU 最佳化 TGI 映像檔。

docker run -p 8080:80 \
        --shm-size 16GB \
        --privileged \
        --net host \
        -e LOG_LEVEL=text_generation_router=debug \
        -v ~/hf_data:/data \
        -e HF_TOKEN=$(cat ~/.cache/huggingface/token) \
        ghcr.io/huggingface/optimum-tpu:v0.2.3-tgi \
        --model-id google/gemma-2b-it \
        --max-input-length 512 \
        --max-total-tokens 1024 \
        --max-batch-prefill-tokens 512 \
        --max-batch-total-tokens 1024

理解設定

關鍵參數說明

--shm-size 16GB --privileged --net=host：Docker 存取 TPU 所需的設定
-v ~/hf_data:/data：用於模型儲存的磁碟區掛載
--max-input-length：最大輸入序列長度
--max-total-tokens：輸入和輸出 Token 的最大總和
--max-batch-prefill-tokens：批次處理的權杖數上限
--max-batch-total-tokens：批次中 Token 的最大總和

步驟 3：發送推論請求

伺服器就緒狀態

等待日誌中出現「Connected」訊息

2025-01-11T10:40:00.256056Z  INFO text_generation_router::server: router/src/server.rs:2393: Connected

您的 TGI 伺服器現在已準備好處理請求。

從 TPU VM 進行測試

從 TPU 執行個體上的另一個終端機查詢伺服器

curl 0.0.0.0:8080/generate \
    -X POST \
    -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":20}}' \
    -H 'Content-Type: application/json'

遠端存取

若要從 TPU 執行個體外部進行查詢

External IP TPU

在 Google Cloud 控制台中找到您的 TPU 外部 IP
在請求中替換該 IP

curl 34.174.11.242:8080/generate \
    -X POST \
    -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":20}}' \
    -H 'Content-Type: application/json'

（選用）防火牆設定

您可能需要設定 GCP 防火牆規則以允許遠端存取

使用 gcloud compute firewall-rules create 來允許流量
確保 8080 連接埠是可以存取的
在生產環境中請考量安全性最佳實踐

請求參數

推論請求的關鍵參數

inputs：提示文字 (prompt)
max_new_tokens：要生成的 Token 最大數量
其他參數可在 TGI 文件中找到

下一步

請查看 TGI 使用指南，了解如何查詢您新的 TGI 伺服器。
請查看我們文件的其餘部分，以了解可用於您新 TGI 伺服器的進階設定。

Google TPUs