此页面由 Cloud Translation API 翻译。

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

本教程演示了如何通过 Saxml 服务框架，使用 GKE 中的 TPU 部署和提供 Gemma 2 大语言模型 (LLM)。本教程为您提供了一个基础，让您能够了解和探索用于在托管式 Kubernetes 环境中进行推理的实际 LLM 部署。您会将使用 Saxml 的预构建容器部署到 GKE。您还需要配置 GKE，以便在运行时从 Cloud Storage 加载 Gemma 2B 和 7B 权重。

本教程适用于机器学习 (ML) 工程师、平台管理员和操作者，以及对使用 Kubernetes 容器编排功能提供 LLM 感兴趣的数据和 AI 专家。如需详细了解我们在Google Cloud 内容中提及的常见角色和示例任务，请参阅常见的 GKE 用户角色和任务。

在阅读本页面之前，请确保您熟悉以下内容：

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

如果您需要统一的代管式 AI 平台来经济高效地快速构建和应用机器学习模型，我们建议您试用我们的 Vertex AI 部署解决方案。

背景

本部分介绍本教程中使用的关键技术。

Gemma

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

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

TPU

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

本教程应用 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 系统互动。REST API 包含用于发布、列出、取消发布模型和生成预测的端点。

为 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 会下载一个名为 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 服务器。本教程使用 Kubernetes Deployment 清单。Deployment 是一个 Kubernetes API 对象，可让您运行在集群节点中分布的多个 Pod 副本。

部署 Saxml 管理服务器

在本部分中，您将部署 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 作业。Kubernetes 中的 Job 控制器会创建一个或多个 Pod，并确保它们成功执行特定任务。

执行与您要使用的 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

验证 Job 是否已成功完成：

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

Job 会将检查点上传到 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

验证 Job 是否已成功完成：

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

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

公开 Saxml HTTP 服务器

您可以通过在部署 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 模型

接下来，您可以将 Gemma 模型发布到在 TPU 切片节点池中运行的模型服务器。您可以使用 Saxml HTTP 服务器的 publish API 发布模型。请按照以下步骤发布 Gemma 2B 或 7B 参数模型。

如需详细了解 Saxml HTTP 服务器的 API，请参阅 Saxml HTTP API。

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，请验证您的端口转发已启用。

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

问题排查

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