在 Google Cloud 上进行首次 TPU 推理

本教程将指导您使用 Text Generation Inference (TGI) (文档) 在 TPU 上进行设置和运行推理。TGI 服务器与 OpenAI messages 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 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 实例外部进行查询

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 Consuming Guide，了解如何查询您的新 TGI 服务器。
2. 查看我们的其他文档，了解可在您的新 TGI 服务器上使用的高级设置。