通过 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。

准备工作

  1. 登录您的 Google Cloud 账号。如果您是 Google Cloud 新手，请创建一个账号来评估我们的产品在实际场景中的表现。新客户还可获享 $300 赠金，用于运行、测试和部署工作负载。

  2. 在 Google Cloud Console 中的项目选择器页面上，选择或创建一个 Google Cloud 项目。

    转到“项目选择器”

  3. 确保您的 Google Cloud 项目已启用结算功能。

  4. 启用 Vertex AI API and Artifact Registry API API。

    启用 API

    5.

  5. 在 Google Cloud Console 中的项目选择器页面上，选择或创建一个 Google Cloud 项目。

    转到“项目选择器”

  6. 确保您的 Google Cloud 项目已启用结算功能。

  7. 启用 Vertex AI API and Artifact Registry API API。

    启用 API

    8.

  8. 在 Google Cloud 控制台中，激活 Cloud Shell。

    激活 Cloud Shell

    Cloud Shell 会话随即会在 Google Cloud 控制台的底部启动，并显示命令行提示符。Cloud Shell 是一个已安装 Google Cloud CLI 且已为当前项目设置值的 Shell 环境。该会话可能需要几秒钟时间来完成初始化。

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

  1. 安装 Google Cloud CLI。

  2. 如需初始化 gcloud CLI，请运行以下命令： 

    gcloud init
  3. 按照 Artifact Registry 文档安装 Docker。
  4. 确保您有足够的配额用于 Vertex AI Prediction 的 5 个 TPU v5e 芯片。
  5. 如果您还没有 Kaggle 账号，请创建一个。

获取对模型的访问权限

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

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

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

生成一个访问令牌

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

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

  1. 在浏览器中，转到 Kaggle 设置。

  2. 在 API 部分下，点击 Create New Token（创建新令牌）。

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

将访问令牌上传到 Cloud Shell

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

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

创建 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
gsutil -m cp -R /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/

# 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
gsutil -m cp -R /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/

创建 Artifact Registry 仓库

创建 Artifact Registry 仓库以存储您将在下一部分中创建的容器映像。

为您的项目启用 Artifact Registry API 服务。

gcloud services enable artifactregistry.googleapis.com

在 shell 中运行以下命令以创建 Artifact Registry 代码库：

gcloud artifacts repositories create saxml \
 --repository-format=docker \
 --location=LOCATION \
 --description="Saxml Docker repository"

将 LOCATION 替换为 Artifact Registry 存储容器映像的区域。稍后，您必须在与此区域匹配的区域端点创建 Vertex AI 模型资源，因此请为 TPU 选择 Vertex AI 在其中具有区域端点的区域，例如 us-west1。

将容器映像推送到 Artifact Registry

us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest 提供了预构建的 Saxml 容器。将其复制到 Artifact Registry。 配置 Docker 以访问 Artifact Registry。然后，将您的容器映像推送到 Artifact Registry 代码库。

  1. 如需向本地 Docker 安装授予推送到所选地区中 Artifact Registry 的权限，请在 shell 中运行以下命令：

    gcloud auth configure-docker LOCATION-docker.pkg.dev
    • 将 LOCATION 替换为创建了仓库的区域。

  2. 如需将您刚刚构建的容器映像复制到 Artifact Registry，请在 shell 中运行以下命令：

    docker tag us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest

  3. 如需将您刚刚构建的容器映像推送到 Artifact Registry，请在 shell 中运行以下命令：

    docker push LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest

    如上一节所述，替换以下内容：

    • LOCATION：您的 Artifact Registry 仓库的区域。
    • PROJECT_ID：您的 Google Cloud 项目的 ID

部署模型

上传模型

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

Gemma 2B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-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=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-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 资源：

  1. 如需从端点取消部署模型并删除端点，请在 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 替换为您在上一部分中创建了模型的地区。

  2. 如需删除模型，请在 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 替换为您在上一部分中创建了模型的地区。

  3. 如需删除 Artifact Registry 代码库和容器映像，请在 shell 中运行以下命令：

    gcloud artifacts repositories delete saxml \
  --location=LOCATION \
  --quiet

    将 LOCATION 替换为您在上一部分中创建了 Artifact Registry 代码库的地区。

限制

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

后续步骤

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