在 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 访问令牌

如果您已经通过 huggingface-cli login 登录，那么您可以设置 HF_TOKEN=$(cat ~/.cache/huggingface/token) 以获得更多便利

您也可以使用 optimum-tpu 容器页面中引用的 GCP 提供的镜像

选项 2：手动构建镜像

为了获取最新功能 (optimum-tpu 的 main 分支) 或自定义修改，请自行构建镜像

克隆仓库

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'