通过 Saxml 使用 GKE 中的 TPU 应用 Gemma 开放模型

Standard

本指南介绍如何通过 Saxml，使用 Google Kubernetes Engine (GKE) 中的张量处理单元 (TPU) 应用 Gemma 开放模型大语言模型 (LLM)。在本教程中，您会将 2B 和 7B 参数指令调优 Gemma 模型下载到 Cloud Storage，并使用运行 Saxml 的容器将它们部署到 GKE Standard 集群。

如果您在 Saxml 上部署模型时需要利用 Kubernetes 功能提供的可伸缩性、弹性和成本效益，那么本指南是一个很好的起点。如果您需要统一的代管式 AI 平台来经济高效地快速构建和应用机器学习模型，我们建议您试用我们的 Vertex AI 部署解决方案。

背景

您可以通过 Saxml 使用 GKE 中的 TPU 应用 Gemma，从而实现一个可直接用于生产环境的强大推理服务解决方案，具备代管式 Kubernetes 的所有优势，包括高效的可伸缩性和更高的可用性。本部分介绍本教程中使用的关键技术。

Gemma

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

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

TPU

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

使用 GKE 中的 TPU 之前，我们建议您完成以下学习路线：

了解 Cloud TPU 系统架构中的当前 TPU 版本可用性。
了解 GKE 中的 TPU。

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

Gemma 2B：在具有 1x1 拓扑（表示一个 TPU 芯片）的 TPU v5e 节点池中托管的指令调优模型。节点的机器类型为 ct5lp-hightpu-1t。
Gemma 7B：在具有 2x2 拓扑（表示四个 TPU 芯片）的 TPU v5e 节点池中托管的指令调优模型。节点的机器类型为 ct5lp-hightpu-4t。

Saxml

Saxml 是一个实验性系统，应用 Paxml、JAX 和 PyTorch 模型以便进行推理。Saxml 系统包含以下组件：

Saxml 单元或 Sax 集群：由一个管理服务器和一组模型服务器组成。管理服务器会跟踪模型服务器、将已发布的模型分配给模型服务器进行应用，以及帮助客户端找到应用特定已发布模型的模型服务器。
Saxml 客户端：面向用户的 Saxml 系统编程接口。Saxml 客户端包含一个命令行工具 (saxutil) 以及一套采用 Python、C++ 和 Go 的客户端库。

在本教程中，您还会使用 Saxml HTTP 服务器。 Saxml HTTP 服务器是自定义 HTTP 服务器，它会封装 Saxml Python 客户端库并公开 REST API 以与 Saxml 系统互动，包括用于发布、列出、取消发布模型和生成预测的端点。

目标

本教程适用于使用 JAX 的生成式 AI 客户、GKE 的新用户或现有用户、机器学习工程师、MLOps (DevOps) 工程师或是对使用 Kubernetes 容器编排功能应用 Gemma 感兴趣的平台管理员。

本教程介绍以下步骤：

根据模型特征准备一个具有推荐 TPU 拓扑的 GKE Standard 集群。
在 GKE 上部署 Saxml 组件。
获取并发布 Gemma 2B 或 Gemma 7B 参数模型。
应用已发布的模型并与之互动。

架构

本部分介绍本教程中使用的 GKE 架构。此架构包含一个 GKE Standard 集群，该集群预配 TPU 并托管 Saxml 组件，以部署和应用 Gemma 2B 或 7B 模型。下图展示了此架构的组件：

本教程中部署的架构的示意图

此架构包括以下组件：

一个 GKE Standard 可用区级集群。
一个单主机 TPU 切片节点池，依赖于您要应用的 Gemma 模型：
- Gemma 2B：使用具有 1x1 拓扑的 TPU v5e 进行配置。一个 Saxml 模型服务器实例配置为使用此节点池。
- Gemma 7B：使用具有 2x2 拓扑的 TPU v5e 进行配置。一个 Saxml 模型服务器实例配置为使用此节点池。
在其中部署 Saxml 管理服务器和 Saxml HTTP 服务器的默认 CPU 节点池。
两个 Cloud Storage 存储桶：
- 一个 Cloud Storage 存储桶存储由管理服务器管理的状态。
- 一个 Cloud Storage 存储桶存储 Gemma 模型检查点。

此架构具有以下特征：

一个公共 Artifact Registry 管理 Saxml 组件的容器映像。
GKE 集群使用适用于 GKE 的工作负载身份联合。所有 Saxml 组件都使用集成了 IAM 服务账号的工作负载身份联合来访问 Cloud Storage 存储桶等外部服务。
Saxml 组件生成的日志集成到 Cloud Logging 中。
您可以使用 Cloud Monitoring 分析本教程创建的 GKE 节点池的性能指标。

准备工作

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

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

转到“项目选择器”

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

启用所需的 API。

启用 API

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

转到“项目选择器”

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

启用所需的 API。

启用 API

确保您拥有项目的以下一个或多个角色： roles/container.admin, roles/iam.serviceAccountAdmin
检查角色
1. 在 Google Cloud 控制台中，前往 IAM 页面。
  转到 IAM
2. 选择项目。
3. 在主账号列中，找到您的电子邮件地址所在的行。
  
  如果您的电子邮件地址不在此列，则表示您没有任何角色。
4. 在您的电子邮件地址所在的行对应的角色列中，检查角色列表是否包含所需的角色。
授予角色
1. 在 Google Cloud 控制台中，前往 IAM 页面。
  转到 IAM
2. 选择项目。
3. 点击 授予访问权限。
4. 在新的主账号字段中，输入您的电子邮件地址。
5. 在选择角色列表中，选择一个角色。
6. 如需授予其他角色，请点击 添加其他角色，然后添加其他各个角色。
7. 点击 Save（保存）。

确保您具有足够的配额，以便用于 5 个 TPU v5e 芯片。在本教程中，您将使用按需实例。
如果您还没有 Kaggle 账号，请创建一个。

为 Gemma 准备环境

启动 Cloud Shell

在本教程中，您将使用 Cloud Shell 来管理 Google Cloud 上托管的资源。Cloud Shell 中预安装了本教程所需的软件，包括 kubectl 和 gcloud CLI。

在 Google Cloud 控制台中，启动 Cloud Shell 实例：
打开 Cloud Shell
设置默认环境变量：
```
gcloud config set project PROJECT_ID
export PROJECT_ID=$(gcloud config get project)
export LOCATION=LOCATION
export CLUSTER_NAME=saxml-tpu
```
替换以下值：
- PROJECT_ID：您的 Google Cloud 项目 ID。
- LOCATION：可使用 TPU v5e 机器类型的 Compute Engine 可用区的名称。

创建 GKE Standard 集群

在本部分中，您将创建 GKE 集群和节点池。

Gemma 2B-it

使用 Cloud Shell 执行以下操作：

创建使用适用于 GKE 的工作负载身份联合的 Standard 集群：

gcloud container clusters create ${CLUSTER_NAME} \
    --enable-ip-alias \
    --machine-type=e2-standard-4 \
    --num-nodes=2 \
    --release-channel=rapid \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --location=${LOCATION}

集群创建可能需要几分钟的时间。

创建具有 1x1 拓扑和一个节点的 TPU v5e 节点池：

gcloud container node-pools create tpu-v5e-1x1 \
    --cluster=${CLUSTER_NAME} \
    --machine-type=ct5lp-hightpu-1t \
    --num-nodes=1 \
    --location=${LOCATION}

您在此节点池中应用 Gemma 2B 模型。

Gemma 7B-it

使用 Cloud Shell 执行以下操作：

创建使用适用于 GKE 的工作负载身份联合的 Standard 集群：

gcloud container clusters create ${CLUSTER_NAME} \
    --enable-ip-alias \
    --machine-type=e2-standard-4 \
    --num-nodes=2 \
    --release-channel=rapid \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --location=${LOCATION}

集群创建可能需要几分钟的时间。

创建具有 2x2 拓扑和一个节点的 TPU v5e 节点池：

gcloud container node-pools create tpu-v5e-2x2 \
    --cluster=${CLUSTER_NAME} \
    --machine-type=ct5lp-hightpu-4t \
    --num-nodes=1 \
    --location=${LOCATION}

您在此节点池中应用 Gemma 7B 模型。

创建 Cloud Storage 存储桶

创建两个 Cloud Storage 存储桶以管理 Saxml 管理服务器的状态和模型检查点。

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

创建一个 Cloud Storage 存储桶以存储 Saxml 管理服务器配置。
```
gcloud storage buckets create gs://ADMIN_BUCKET_NAME
```
将 ADMIN_BUCKET_NAME 替换为存储 Saxml 管理服务器的 Cloud Storage 存储桶的名称。
创建一个 Cloud Storage 存储桶以存储模型检查点：
```
gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME
```
将 CHECKPOINTS_BUCKET_NAME 替换为存储模型检查点的 Cloud Storage 存储桶的名称。

使用适用于 GKE 的工作负载身份联合配置工作负载访问权限

为应用分配 Kubernetes ServiceAccount，并将该 Kubernetes ServiceAccount 配置为充当 IAM 服务账号。

配置 kubectl 以与您的集群通信：

gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${LOCATION}

为您的应用创建 Kubernetes 服务账号：
```
gcloud iam service-accounts create wi-sax
```

为您的 IAM 服务账号添加 IAM 政策绑定，以便对 Cloud Storage 执行读写操作：

gcloud projects add-iam-policy-binding ${PROJECT_ID} \
    --member "serviceAccount:wi-sax@${PROJECT_ID}.iam.gserviceaccount.com" \
    --role roles/storage.objectUser

gcloud projects add-iam-policy-binding ${PROJECT_ID} \
    --member "serviceAccount:wi-sax@${PROJECT_ID}.iam.gserviceaccount.com" \
    --role roles/storage.insightsCollectorService

通过在两个服务账号之间添加 IAM 政策绑定，允许 Kubernetes 服务账号模拟 IAM 服务账号。此绑定允许 Kubernertes ServiceAccount 充当 IAM 服务账号：

gcloud iam service-accounts add-iam-policy-binding wi-sax@${PROJECT_ID}.iam.gserviceaccount.com \
    --role roles/iam.workloadIdentityUser \
    --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/default]"

使用 IAM 服务账号的电子邮件地址为 Kubernetes 服务账号添加注解：

kubectl annotate serviceaccount default \
    iam.gke.io/gcp-service-account=wi-sax@${PROJECT_ID}.iam.gserviceaccount.com

获取对模型的访问权限

如需获取对 Gemma 模型的访问权限以便部署到 GKE，您必须登录 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 文件。
点击上传。

为 Kaggle 凭据创建 Kubernetes Secret

在 Cloud Shell 中，执行以下操作：

配置 kubectl 以与您的集群通信：

gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${LOCATION}

创建一个 Secret 以存储 Kaggle 凭据：

kubectl create secret generic kaggle-secret \
    --from-file=kaggle.json

部署 Saxml

在本部分中，您将部署 Saxml 管理服务器、模型服务器和 HTTP 服务器。

部署 Saxml 管理服务器

创建以下 saxml-admin-server.yaml 清单：

apiVersion: apps/v1
kind: Deployment
metadata:
  name: sax-admin-server
spec:
  replicas: 1
  selector:
    matchLabels:
      app: sax-admin-server
  template:
    metadata:
      labels:
        app: sax-admin-server
    spec:
      hostNetwork: false
      containers:
      - name: sax-admin-server
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.2.0
        securityContext:
          privileged: true
        ports:
        - containerPort: 10000
        env:
        - name: GSBUCKET
          value: ADMIN_BUCKET_NAME

将 ADMIN_BUCKET_NAME 替换为您在创建 Cloud Storage 存储桶部分中创建的存储桶的名称。请勿添加 gs:// 前缀。

应用清单：

kubectl apply -f saxml-admin-server.yaml

验证管理服务器部署：

kubectl get deployment

输出应类似如下所示：

NAME                              READY   UP-TO-DATE   AVAILABLE   AGE
sax-admin-server                  1/1     1            1           ##s

部署 Saxml 模型服务器

请按照以下说明为 Gemma 2B 或 Gemma 7B 模型部署模型服务器。

Gemma 2B-it

创建以下 saxml-model-server-1x1.yaml 清单：

apiVersion: apps/v1
kind: Deployment
metadata:
  name: sax-model-server-v5e-1x1
spec:
  replicas: 1
  selector:
    matchLabels:
      app: gemma-server
  strategy:
    type: Recreate
  template:
    metadata:
      labels:
        app: gemma-server
        ai.gke.io/model: gemma-2b-it
        ai.gke.io/inference-server: saxml
        examples.ai.gke.io/source: user-guide
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 1x1
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      hostNetwork: false
      restartPolicy: Always
      containers:
      - name: inference-server
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.2.0
        args:
        - "--jax_platforms=tpu"
        - "--platform_chip=tpuv5e"
        - "--platform_topology=1x1"
        - "--port=10001"
        - "--sax_cell=/sax/test"
        ports:
        - containerPort: 10001
        securityContext:
          privileged: true
        env:
        - name: SAX_ROOT
          value: "gs://ADMIN_BUCKET_NAME/sax-root"
        resources:
          requests:
            google.com/tpu: 1
          limits:
            google.com/tpu: 1

将 ADMIN_BUCKET_NAME 替换为您在创建 Cloud Storage 存储桶部分中创建的存储桶的名称。请勿添加 gs:// 前缀。

应用清单：

kubectl apply -f saxml-model-server-1x1.yaml

验证模型服务器部署的状态：

kubectl get deployment

输出应类似如下所示：

NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server                                  1/1     Running   0          ##m
sax-model-server-v5e-1x1                          1/1     Running   0          ##s

Gemma 7B-it

创建以下 saxml-model-server-2x2.yaml 清单：

apiVersion: apps/v1
kind: Deployment
metadata:
  name: sax-model-server-v5e-2x2
spec:
  replicas: 1
  selector:
    matchLabels:
      app: gemma-server
  strategy:
    type: Recreate
  template:
    metadata:
      labels:
        app: gemma-server
        ai.gke.io/model: gemma-7b-it
        ai.gke.io/inference-server: saxml
        examples.ai.gke.io/source: user-guide
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x2
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      hostNetwork: false
      restartPolicy: Always
      containers:
      - name: inference-server
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.2.0
        args:
        - "--jax_platforms=tpu"
        - "--platform_chip=tpuv5e"
        - "--platform_topology=2x2"
        - "--port=10001"
        - "--sax_cell=/sax/test"
        ports:
        - containerPort: 10001
        securityContext:
          privileged: true
        env:
        - name: SAX_ROOT
          value: "gs://ADMIN_BUCKET_NAME/sax-root"
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

将 ADMIN_BUCKET_NAME 替换为您在创建 Cloud Storage 存储桶部分中创建的存储桶的名称。请勿添加 gs:// 前缀。

应用清单：

kubectl apply -f saxml-model-server-2x2.yaml

验证模型服务器部署的状态：

kubectl get deployment

输出应类似如下所示：

NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server                                  1/1     Running   0          ##m
sax-model-server-v5e-2x2                          1/1     Running   0          ##s

部署 Saxml HTTP 服务器

在本部分中，您将部署 Saxml HTTP 服务器并创建用于访问该服务器的集群 IP 服务。

创建以下 saxml-http.yaml 清单：

apiVersion: apps/v1
kind: Deployment
metadata:
  name: sax-http
spec:
  replicas: 1
  selector:
    matchLabels:
      app: sax-http
  template:
    metadata:
      labels:
        app: sax-http
    spec:
      hostNetwork: false
      containers:
      - name: sax-http
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-http:v1.2.0
        imagePullPolicy: Always
        ports:
        - containerPort: 8888
        env:
        - name: SAX_ROOT
          value: "gs://ADMIN_BUCKET_NAME/sax-root"
---
apiVersion: v1
kind: Service
metadata:
  name: sax-http-svc
spec:
  selector:
    app: sax-http
  ports:
  - protocol: TCP
    port: 8888
    targetPort: 8888
  type: ClusterIP

将 ADMIN_BUCKET_NAME 替换为存储 Saxml 管理服务器的 Cloud Storage 存储桶的名称。

应用清单：
```
kubectl apply -f saxml-http.yaml
```
验证 Saxml HTTP 服务器部署的状态：
```
kubectl get deployment
```

Gemma 2B-it

输出应类似如下所示：

NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server                                  1/1     Running   0          ##m
sax-model-server-v5e-1x1                          1/1     Running   0          ##m
sax-http                                          1/1     Running   0          ##s

Gemma 7B-it

输出应类似如下所示：

NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server                                  1/1     Running   0          ##m
sax-model-server-v5e-2x2                          1/1     Running   0          ##m
sax-http                                          1/1     Running   0          ##s

下载模型检查点

在本部分中，您将运行一个提取、下载和存储模型检查点的 Kubernetes 作业。根据您要使用的 Gemma 模型执行相应步骤：

Gemma 2B-it

创建以下 job-2b.yaml 清单：

apiVersion: v1
kind: ConfigMap
metadata:
  name: fetch-model-scripts
data:
  fetch_model.sh: |-
    #!/usr/bin/bash -x
    pip install kaggle --break-system-packages && \

    MODEL_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $2}') && \
    VARIATION_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $4}') && \

    mkdir -p /data/${MODEL_NAME}_${VARIATION_NAME} &&\
    kaggle models instances versions download ${MODEL_PATH} --untar -p /data/${MODEL_NAME}_${VARIATION_NAME} && \
    echo -e "\nCompleted extraction to /data/${MODEL_NAME}_${VARIATION_NAME}" && \

    gcloud storage rsync --recursive --no-clobber /data/${MODEL_NAME}_${VARIATION_NAME} gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME} && \
    echo -e "\nCompleted copy of data to gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME}"
---
apiVersion: batch/v1
kind: Job
metadata:
  name: data-loader-2b
  labels:
    app: data-loader-2b
spec:
  ttlSecondsAfterFinished: 120
  template:
    metadata:
      labels:
        app: data-loader-2b
    spec:
      restartPolicy: OnFailure
      containers:
      - name: gcloud
        image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
        command:
        - /scripts/fetch_model.sh
        env:
        - name: BUCKET_NAME
          value: CHECKPOINTS_BUCKET_NAME
        - name: KAGGLE_CONFIG_DIR
          value: /kaggle
        - name: MODEL_PATH
          value: "google/gemma/pax/2b-it/2"
        volumeMounts:
        - mountPath: "/kaggle/"
          name: kaggle-credentials
          readOnly: true
        - mountPath: "/scripts/"
          name: scripts-volume
          readOnly: true
      volumes:
      - name: kaggle-credentials
        secret:
          defaultMode: 0400
          secretName: kaggle-secret
      - name: scripts-volume
        configMap:
          defaultMode: 0700
          name: fetch-model-scripts

将 CHECKPOINTS_BUCKET_NAME 替换为您在创建 Cloud Storage 存储桶部分中创建的存储桶的名称。请勿添加 gs:// 前缀。

应用清单：
```
kubectl apply -f job-2b.yaml
```

等待作业完成：

kubectl wait --for=condition=complete --timeout=180s job/data-loader-2b

输出类似于以下内容：

job.batch/data-loader-2b condition met

验证作业是否已成功完成：

kubectl get job/data-loader-2b

输出类似于以下内容：

NAME             COMPLETIONS   DURATION   AGE
data-loader-2b   1/1           ##s        #m##s

查看作业的日志：

kubectl logs --follow job/data-loader-2b

检查点会上传到 gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000。

Gemma 7B-it

创建以下 job-7b.yaml 清单：

apiVersion: v1
kind: ConfigMap
metadata:
  name: fetch-model-scripts
data:
  fetch_model.sh: |-
    #!/usr/bin/bash -x
    pip install kaggle --break-system-packages && \

    MODEL_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $2}') && \
    VARIATION_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $4}') && \

    mkdir -p /data/${MODEL_NAME}_${VARIATION_NAME} &&\
    kaggle models instances versions download ${MODEL_PATH} --untar -p /data/${MODEL_NAME}_${VARIATION_NAME} && \
    echo -e "\nCompleted extraction to /data/${MODEL_NAME}_${VARIATION_NAME}" && \

    gcloud storage rsync --recursive --no-clobber /data/${MODEL_NAME}_${VARIATION_NAME} gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME} && \
    echo -e "\nCompleted copy of data to gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME}"
---
apiVersion: batch/v1
kind: Job
metadata:
  name: data-loader-7b
  labels:
    app: data-loader-7b
spec:
  ttlSecondsAfterFinished: 120
  template:
    metadata:
      labels:
        app: data-loader-7b
    spec:
      restartPolicy: OnFailure
      containers:
      - name: gcloud
        image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
        command:
        - /scripts/fetch_model.sh
        env:
        - name: BUCKET_NAME
          value: CHECKPOINTS_BUCKET_NAME
        - name: KAGGLE_CONFIG_DIR
          value: /kaggle
        - name: MODEL_PATH
          value: "google/gemma/pax/7b-it/2"
        volumeMounts:
        - mountPath: "/kaggle/"
          name: kaggle-credentials
          readOnly: true
        - mountPath: "/scripts/"
          name: scripts-volume
          readOnly: true
      volumes:
      - name: kaggle-credentials
        secret:
          defaultMode: 0400
          secretName: kaggle-secret
      - name: scripts-volume
        configMap:
          defaultMode: 0700
          name: fetch-model-scripts

将 CHECKPOINTS_BUCKET_NAME 替换为您在创建 Cloud Storage 存储桶部分中创建的存储桶的名称。请添加 gs:// 前缀。

应用清单：
```
kubectl apply -f job-7b.yaml
```

等待作业完成：

kubectl wait --for=condition=complete --timeout=360s job/data-loader-7b

输出类似于以下内容：

job.batch/data-loader-7b condition met

验证作业是否已成功完成：

kubectl get job/data-loader-7b

输出类似于以下内容：

NAME             COMPLETIONS   DURATION   AGE
data-loader-7b   1/1           ##s        #m##s

查看作业的日志：

kubectl logs --follow job/data-loader-7b

检查点会上传到 gs://CHECKPOINTS_BUCKET_NAME/gemma_7b_it/checkpoint_00000000。

公开 Saxml HTTP 服务器

您可以通过在上一步中创建的 ClusterIP Service 访问 Saxml HTTP 服务器。只能从集群内部访问 ClusterIP Service。因此，如需从集群外部访问该 Service，请完成以下步骤：

建立端口转发会话：

kubectl port-forward service/sax-http-svc 8888:8888

通过打开新终端并运行以下命令，验证您是否可以访问 Saxml HTTP 服务器：
```
curl -s localhost:8888
```
输出类似于以下内容：
```
{
    "Message": "HTTP Server for SAX Client"
}
```

Saxml HTTP 服务器会封装与 Saxml 系统之间的客户端接口，并通过一组 REST API 来公开。您将使用这些 API 发布、管理 Gemma 2B 和 Gemma 7B 模型并与这些模型对接。

发布 Gemma 模型

设置 Saxml 组件后，您可以将 Gemma 模型发布到在 TPU 切片节点池中运行的模型服务器。您可以使用 Saxml HTTP 服务器的 publish API 发布模型。如需详细了解 Saxml HTTP 服务器的 API，请参阅 Saxml HTTP API。请按照以下步骤发布 Gemma 2B 或 7B 参数模型。

Gemma 2B-it

确保您的端口转发会话仍然处于活动状态：
```
curl -s localhost:8888
```

发布 Gemma 2B 参数：

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/publish \
--data \
'{
    "model": "/sax/test/gemma2bfp16",
    "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
    "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
    "replicas": "1"
}'

输出类似于以下内容：

{
    "model": "/sax/test/gemma2bfp16",
    "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
    "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
    "replicas": 1
}

模型的部署可能需要几分钟的时间。

通过观察 sax-model-server-v5e-1x1 部署的模型服务器 Pod 中的日志来监控进度。

kubectl logs --follow deployment/sax-model-server-v5e-1x1

此部署最多可能需要五分钟才能完成。请等待显示类似如下内容的消息：

I0125 15:34:31.685555 139063071708736 servable_model.py:699] loading completed.
I0125 15:34:31.686286 139063071708736 model_service_base.py:532] Successfully loaded model for key: /sax/test/gemma2bfp16

通过显示模型信息来验证您是否可以访问模型：

curl --request GET \
--header "Content-type: application/json" \
-s \
localhost:8888/listcell \
--data \
'{
    "model": "/sax/test/gemma2bfp16"
}'

输出类似于以下内容：

{
    "model": "/sax/test/gemma2bfp16",
    "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
    "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
    "max_replicas": 1,
    "active_replicas": 1
}

Gemma 7B-it

确保您的端口转发会话仍然处于活动状态：
```
curl -s localhost:8888
```

发布 Gemma 7B 参数：

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/publish \
--data \
'{
    "model": "/sax/test/gemma7bfp16",
    "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
    "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
    "replicas": "1"
}'

输出类似于以下内容：

{
    "model": "/sax/test/gemma7bfp16",
    "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
    "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
    "replicas": 1
}

模型的部署可能需要几分钟的时间。

通过观察 sax-model-server-v5e-2x2 部署的模型服务器 Pod 中的日志来监控进度。

kubectl logs --follow deployment/sax-model-server-v5e-2x2

请等待显示类似如下内容的消息：

I0125 15:34:31.685555 139063071708736 servable_model.py:699] loading completed.
I0125 15:34:31.686286 139063071708736 model_service_base.py:532] Successfully loaded model for key: /sax/test/gemma7bfp16

通过显示模型信息来验证模型是否已发布：

curl --request GET \
--header "Content-type: application/json" \
-s \
localhost:8888/listcell \
--data \
'{
    "model": "/sax/test/gemma7bfp16"
}'

输出类似于以下内容：

{
    "model": "/sax/test/gemma7bfp16",
    "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
    "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
    "max_replicas": 1,
    "active_replicas": 1
}

使用模型

您可以与 Gemma 2B 或 7B 模型互动。使用 Saxml HTTP 服务器的 generate API 可向模型发送提示。

Gemma 2B-it

使用 Saxml HTTP 服务器的 generate 端点提交提示请求：

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/generate \
--data \
'{
  "model": "/sax/test/gemma2bfp16",
  "query": "What are the top 5 most popular programming languages?"
}'

以下是模型响应的示例。输出可能因您提交的每个提示而异：

[
    [
        "\n\n1. **Python**\n2. **JavaScript**\n3. **Java**\n4. **C++**\n5. **Go**",
        -3.0704939365386963
    ]
]

您可以使用不同的 query 参数运行该命令。您还可以使用 generate API 修改其他参数，例如 temperature、top_k、topc_p。如需详细了解 Saxml HTTP 服务器的 API，请参阅 Saxml HTTP API。

Gemma 7B-it

使用 Saxml HTTP 服务器的 generate 端点提交提示请求：

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/generate \
--data \
'{
  "model": "/sax/test/gemma7bfp16",
  "query": "What are the top 5 most popular programming languages?"
}'

以下是模型响应的示例。输出可能因您提交的每个提示而异：

[
    [
        "\n\n**1. JavaScript**\n\n* Most widely used language on the web.\n* Used for front-end development, such as websites and mobile apps.\n* Extensive libraries and frameworks available.\n\n**2. Python**\n\n* Known for its simplicity and readability.\n* Versatile, used for various tasks, including data science, machine learning, and web development.\n* Large and active community.\n\n**3. Java**\n\n* Object-oriented language widely used in enterprise applications.\n* Used for web applications, mobile apps, and enterprise software.\n* Strong ecosystem and support.\n\n**4. Go**\n\n",
        -16.806324005126953
    ]
]

取消发布模型

请按照以下步骤取消发布模型：

Gemma 2B-it

如需取消发布 Gemma 2B-it 模型，请运行以下命令：

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/unpublish \
--data \
'{
    "model": "/sax/test/gemma2bfp16"
}'

输出可能类似以下内容：

{
    "model": "/sax/test/gemma2bfp16"
}

您可以使用在 query 参数中传递的不同提示来运行该命令。

Gemma 7B-it

如需取消发布 Gemma 7B-it 模型，请运行以下命令：

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/unpublish \
--data \
'{
    "model": "/sax/test/gemma7bfp16"
}'

输出可能类似以下内容：

{
    "model": "/sax/test/gemma7bfp16"
}

您可以使用在 query 参数中传递的不同提示来运行该命令。

问题排查

如果您收到 Empty reply from server 消息，则容器可能尚未完成模型数据下载。再次检查 Pod 的日志中是否包含 Connected 消息，该消息表明模型已准备好进行应用。
如果您看到 Connection refused，请验证您的端口转发已启用。

清理

为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用，请删除包含这些资源的项目，或者保留项目但删除各个资源。

删除已部署的资源

为避免因您在本指南中创建的资源导致您的 Google Cloud 账号产生费用，请运行以下命令：

gcloud container clusters delete ${CLUSTER_NAME} --location=${LOCATION}
gcloud iam service-accounts delete --quiet wi-sax@${PROJECT_ID}.iam.gserviceaccount.com
gcloud storage rm --recursive gs://ADMIN_BUCKET_NAME
gcloud storage rm --recursive gs://CHECKPOINTS_BUCKET_NAME

替换以下内容：

ADMIN_BUCKET_NAME：存储 Saxml 管理服务器的 Cloud Storage 存储桶的名称。
CHECKPOINTS_BUCKET_NAME：存储模型检查点的 Cloud Storage 存储桶的名称。

后续步骤

详细了解 GKE 中的 TPU。
探索 Saxml GitHub 代码库。
探索 Vertex AI Model Garden。
了解如何使用 GKE 平台编排功能运行经过优化的 AI/机器学习工作负载。
探索有关 Google Cloud 的参考架构、图表和最佳做法。查看我们的 Cloud Architecture Center。

通过 Saxml 使用 GKE 中的 TPU 应用 Gemma 开放模型

背景

Gemma

TPU

Saxml

目标

架构

准备工作

检查角色

授予角色

为 Gemma 准备环境

启动 Cloud Shell

创建 GKE Standard 集群

Gemma 2B-it

Gemma 7B-it

创建 Cloud Storage 存储桶

使用适用于 GKE 的工作负载身份联合配置工作负载访问权限

获取对模型的访问权限

签署许可同意协议

生成一个访问令牌

将访问令牌上传到 Cloud Shell

为 Kaggle 凭据创建 Kubernetes Secret

部署 Saxml

部署 Saxml 管理服务器

部署 Saxml 模型服务器

Gemma 2B-it

Gemma 7B-it

部署 Saxml HTTP 服务器

Gemma 2B-it

Gemma 7B-it

下载模型检查点

Gemma 2B-it

Gemma 7B-it

公开 Saxml HTTP 服务器

发布 Gemma 模型

Gemma 2B-it

Gemma 7B-it

使用模型

Gemma 2B-it

Gemma 7B-it

取消发布模型

Gemma 2B-it

Gemma 7B-it

问题排查

清理

删除已部署的资源

后续步骤