在 Google Cloud 上首次進行 TPU 推理

本教程將指導您如何使用文字生成推理 (TGI)（文件）在 TPU 上設定和執行推理。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` 模型

1. 訪問 https://huggingface.co/google/gemma-2b-it
2. 接受模型條款和條件
3. 這將啟用模型下載許可權

啟動 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: 最大輸入和輸出總令牌數
• --max-batch-prefill-tokens：批處理的最大令牌數
• --max-batch-total-tokens: 批處理中的最大總令牌數

步驟 3：發出推理請求

伺服器就緒

等待日誌中出現“Connected”訊息

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

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

從 TPU 虛擬機器進行測試

在 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 例項外部查詢

1. 在 Google Cloud Console 中找到您的 TPU 外部 IP
2. 替換請求中的 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 防火牆規則以允許遠端訪問

1. 使用 `gcloud compute firewall-rules create` 允許流量
2. 確保埠 8080 可訪問
3. 考慮生產環境的安全最佳實踐

請求引數

推理請求的關鍵引數

• inputs: 提示文字
• max_new_tokens: 要生成的最大令牌數
• 可在 TGI 文件中找到更多引數

後續步驟

1. 請查閱 TGI 消費指南，瞭解如何查詢您的新 TGI 伺服器。
2. 查閱我們的其餘文件，瞭解可在新 TGI 伺服器上使用的更多高階設定。