通过 Saxml 在 Vertex AI Prediction 上使用 TPU 服务 Gemma 开放模型

本指南介绍了如何通过 Saxml 在 Vertex AI Prediction 上使用张量处理单元 (TPU) 服务 Gemma 开放模型大语言模型 (LLM)。在本指南中，您需要将 2B 和 7B 参数指令调整后的 Gemma 模型下载到 Cloud Storage，并将其部署到在 TPU 上运行 Saxml 的 Vertex AI Prediction 上。

背景

通过 Saxml 在 Vertex AI Prediction 上使用 TPU 服务 Gemma。您可以利用托管 AI 解决方案，该解决方案负责底层基础架构，并提供经济高效的方法来服务 LLM。本部分介绍本教程中使用的关键技术。

Gemma

Gemma 是一组公开提供的轻量级生成式人工智能 (AI) 模型（根据开放许可发布）。这些 AI 模型可以在应用、硬件、移动设备或托管服务中运行。您可以使用 Gemma 模型生成文本，但也可以针对专门任务对这些模型进行调优。

如需了解详情，请参阅 Gemma 文档。

Saxml

Saxml 是一个实验性系统，应用 Paxml、JAX 和 PyTorch 模型以便进行推理。在本教程中，我们将介绍如何在对 Saxml 更具成本效益的 TPU 上服务 Gemma。GPU 的设置过程类似。Saxml 提供了用于为 Vertex AI Prediction 构建容器的脚本，我们将在本教程中使用这一脚本。

TPU

TPU 是 Google 定制开发的应用专用集成电路 (ASIC)，用于加速数据处理框架（例如 TensorFlow、PyTorch 和 JAX）。

本教程应用 Gemma 2B 和 Gemma 7B 模型。Vertex AI Prediction 在以下单主机 TPU v5e 节点池上托管这些模型：

Gemma 2B：托管在 TPU v5e 节点池中，其 1x1 拓扑表示一个 TPU 芯片。节点的机器类型为 ct5lp-hightpu-1t。
Gemma 7B：托管在 TPU v5e 节点池中，其 2x2 拓扑表示四个 TPU 芯片。节点的机器类型为 ct5lp-hightpu-4t。

准备工作

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the Vertex AI API.

Enable the API

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the Vertex AI API.

Enable the API

In the Google Cloud console, activate Cloud Shell.

Activate Cloud Shell

At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

本教程假定您使用 Cloud Shell 与 Google Cloud 进行交互。如果您想使用其他 shell 取代 Cloud Shell，请执行以下额外的配置：

Install the Google Cloud CLI.
To initialize the gcloud CLI, run the following command:
```
gcloud init
```
确保您有足够的 TPU v5e 芯片配额以用于 Vertex AI Prediction。默认情况下，此配额为 0。对于 1x1 拓扑，它必须为 1。对于 2x2，它必须为 4。如需同时运行这两个拓扑，它必须为 5。
如果您还没有 Kaggle 账号，请创建一个。

获取对模型的访问权限

请注意，Cloud Shell 可能没有足够的资源来下载模型权重。如果是这样，您可以创建 Vertex AI Workbench 实例来执行该任务。

如需获取对 Gemma 模型的访问权限以便部署到 Vertex AI Prediction，您必须登录 Kaggle 平台，签署许可同意协议，并获取 Kaggle API 令牌。在本教程中，您会将 Kubernetes Secret 用于 Kaggle 凭据。

您必须签署同意协议才能使用 Gemma。请按照以下说明操作：

访问 Kaggle.com 上的模型同意页面。
如果您尚未登录 Kaggle，请进行登录。
点击申请访问权限。
在 Choose Account for Consent（选择进行同意的账号）部分中，选择 Verify via Kaggle Account（通过 Kaggle 账号验证），以使用您的 Kaggle 账号进行同意。
接受模型条款及条件。

生成一个访问令牌

如需通过 Kaggle 访问模型，您需要 Kaggle API 令牌。

如果您还没有令牌，请按照以下步骤生成新令牌：

在浏览器中，转到 Kaggle 设置。
在 API 部分下，点击 Create New Token（创建新令牌）。

系统将下载名为 kaggle.json 的文件。

将访问令牌上传到 Cloud Shell

在 Cloud Shell 中，您可以将 Kaggle API 令牌上传到 Google Cloud 项目：

在 Cloud Shell 中，点击更多 > 上传。
选择“文件”，然后点击选择文件。
打开 kaggle.json 文件。
点击上传。

创建 Cloud Storage 存储桶

创建 Cloud Storage 存储桶以存储模型检查点。

在 Cloud Shell 中，运行以下命令：

gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME

将 CHECKPOINTS_BUCKET_NAME 替换为存储模型检查点的 Cloud Storage 存储桶的名称。

将模型复制到 Cloud Storage 存储桶

在 Cloud Shell 中，运行以下命令：

pip install kaggle --break-system-packages

# For Gemma 2B
mkdir -p /data/gemma_2b-it
kaggle models instances versions download google/gemma/pax/2b-it/1 --untar -p /data/gemma_2b-it
gcloud storage cp /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/ --recursive

# For Gemma 7B
mkdir -p /data/gemma_7b-it
kaggle models instances versions download google/gemma/pax/7b-it/1 --untar -p /data/gemma_7b-it
gcloud storage cp /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/ --recursive

部署模型

上传模型

如需上传使用 Saxml 容器的 Model 资源，请运行以下 gcloud ai models upload 命令：

Gemma 2B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma2BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Gemma 7B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma7BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

替换以下内容：

PROJECT_ID：您的 Google Cloud 项目的 ID
LOCATION_ID：您在其中使用 Vertex AI 的区域。请注意，TPU 仅可用于 us-west1。
DEPLOYED_MODEL_NAME：DeployedModel 的名称。您还可以将 Model 的显示名用于 DeployedModel。

创建端点

您必须先将模型部署到端点，然后才能使用模型进行在线预测。如果要将模型部署到现有端点，则可以跳过此步骤。以下示例使用 gcloud ai endpoints create 命令：

gcloud ai endpoints create \
  --region=LOCATION \
  --display-name=ENDPOINT_NAME

替换以下内容：

LOCATION_ID：您在其中使用 Vertex AI 的区域。
ENDPOINT_NAME：端点的显示名称。

Google Cloud CLI 工具可能需要几秒钟才能创建端点。

将模型部署到端点

端点准备就绪后，您便可以将模型部署到端点。

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

MODEL_ID=$(gcloud ai models list \
   --region=LOCATION \
   --filter=display_name=DEPLOYED_MODEL_NAME \
   --format="value(name)")

gcloud ai endpoints deploy-model $ENDPOINT_ID \
  --region=LOCATION \
  --model=$MODEL_ID \
  --display-name=DEPLOYED_MODEL_NAME \
  --machine-type=ct5lp-hightpu-4t \
  --traffic-split=0=100

替换以下内容：

LOCATION_ID：您在其中使用 Vertex AI 的区域。
ENDPOINT_NAME：端点的显示名称。
DEPLOYED_MODEL_NAME：DeployedModel 的名称。您还可以将 Model 的显示名用于 DeployedModel。

Gemma 2B 可以部署在较小的 ct5lp-hightpu-1t 机器上，在这种情况下，您应该在上传模型时指定 --platform_topology=1x1。

Google Cloud CLI 工具可能需要几分钟才能将模型部署到端点。成功部署模型后，此命令将输出以下内容：

  Deployed a model to the endpoint xxxxx. Id of the deployed model: xxxxx.

从已部署的模型获取在线预测结果

如需通过 Vertex AI Prediction 端点调用模型，请使用标准推断请求 JSON 对象设置预测请求的格式。

以下示例使用 gcloud ai endpoints predict 命令：

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

gcloud ai endpoints predict $ENDPOINT_ID \
  --region=LOCATION \
  --http-headers=Content-Type=application/json \
  --json-request instances.json

替换以下内容：

LOCATION_ID：您在其中使用 Vertex AI 的区域。
ENDPOINT_NAME：端点的显示名称。
instances.json 的格式如下：{"instances": [{"text_batch": "<your prompt>"},{...}]}

清理

为避免产生额外的 Vertex AI 费用和 Artifact Registry 费用，请删除您在本教程中创建的 Google Cloud 资源：

如需从端点取消部署模型并删除端点，请在 shell 中运行以下命令：

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

DEPLOYED_MODEL_ID=$(gcloud ai endpoints describe $ENDPOINT_ID \
   --region=LOCATION \
   --format="value(deployedModels.id)")

gcloud ai endpoints undeploy-model $ENDPOINT_ID \
  --region=LOCATION \
  --deployed-model-id=$DEPLOYED_MODEL_ID

gcloud ai endpoints delete $ENDPOINT_ID \
   --region=LOCATION \
   --quiet

将 LOCATION 替换为您在上一部分中创建了模型的地区。

如需删除模型，请在 shell 中运行以下命令：

MODEL_ID=$(gcloud ai models list \
   --region=LOCATION \
   --filter=display_name=DEPLOYED_MODEL_NAME \
   --format="value(name)")

gcloud ai models delete $MODEL_ID \
   --region=LOCATION \
   --quiet

将 LOCATION 替换为您在上一部分中创建了模型的地区。

限制

在 Vertex AI Prediction 上，只有 us-west1 支持 Cloud TPU。如需了解详情，请参阅位置。

后续步骤

了解如何部署其他 Saxml 模型，例如 Llama2 和 GPT-J。