在 Google Cloud TPU 執行個體上部署文字生成推論伺服器 (TGI)

文字生成推論 (Text-Generation-Inference, TGI) 支援在 TPU 上提供大型語言模型 (LLM) 服務，而 Optimum TPU 提供了針對 TPU 硬體全面優化的專用 TGI 執行環境。

TGI 還提供與 OpenAI 相容的 API，便於與眾多工具進行整合。

如需支援的模型列表，請查看支援模型頁面。

在 Cloud TPU 執行個體上部署 TGI

本指南假設您已經有一個正在執行的 Cloud TPU 執行個體。如果沒有，請參閱我們的部署指南。

您有兩種部署 TGI 的選項：

使用我們預先建置的 TGI 映像檔（建議使用）
手動建置映像檔以獲取最新功能

選項 1：使用預先建置的映像檔

optimum-tpu 映像檔可於 ghcr.io/huggingface/optimum-tpu:v0.2.3-tgi 取得。請查看 optimum-tpu 容器說明文件以獲取最新的 TGI 映像檔。服務教學指南也詳細說明了如何從預先建置的映像檔啟動 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=<your_hf_token_here> \
        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

您需要將替換為您的 HuggingFace 存取權杖（Access Token），您可以[在此處](https://huggingface.co/settings/tokens)取得。

如果您已經透過 `huggingface-cli login` 登入，則可以設定 HF_TOKEN=$(cat ~/.cache/huggingface/token) 以獲得更方便的操作體驗。

您也可以使用 GCP 提供的映像檔，詳情請參考 optimum-tpu 容器頁面。

選項 2：手動建置映像檔

若要使用最新功能（optimum-tpu 的 main 分支）或進行自訂修改，請自行建置映像檔。

複製 (clone) 儲存庫

git clone https://github.com/huggingface/optimum-tpu.git

建置映像檔

make tpu-tgi

執行容器

HF_TOKEN=<your_hf_token_here>
MODEL_ID=google/gemma-2b-it

sudo docker run --net=host \
                --privileged \
                -v $(pwd)/data:/data \
                -e HF_TOKEN=${HF_TOKEN} \
                huggingface/optimum-tpu:latest \
                --model-id ${MODEL_ID} \
                --max-concurrent-requests 4 \
                --max-input-length 32 \
                --max-total-tokens 64 \
                --max-batch-size 1

對服務執行請求

您可以使用 /generate 或 /generate_stream 路由來查詢模型。

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

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

Google TPUs

在 Google Cloud TPU 執行個體上部署文字生成推論伺服器 (TGI)

在 Cloud TPU 執行個體上部署 TGI

選項 1：使用預先建置的映像檔

選項 2：手動建置映像檔

對服務執行請求