在 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 服务器上使用的更多高级设置。