Saxml を備えた GKE で TPU を使用して Gemma オープンモデルを提供する

Standard

このガイドでは、Saxml を備えた Google Kubernetes Engine（GKE）で Tensor Processing Unit（TPU）を使用して、Gemma オープンモデル大規模言語モデル（LLM）を提供する方法について説明します。このガイドでは、2B パラメータと 7B パラメータの指示でチューニングした Gemma モデルを Cloud Storage にダウンロードし、Saxml を実行するコンテナを使用して GKE Standard クラスタにデプロイします。

このガイドは、モデルを Saxml にデプロイするときに Kubernetes の機能によって実現されるスケーラビリティ、復元力、費用対効果が必要な場合に適しています。

ML モデルを費用対効果の高い方法で迅速に構築して提供するために、統合されたマネージド 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: 1 つの TPU チップを表す 1x1 トポロジを持つ TPU v5e ノードプールでホストされる指示チューニング済みモデル。ノードのマシンタイプは ct5lp-hightpu-1t です。
Gemma 7B: 4 つの TPU チップを表す 2x2 トポロジを持つ TPU v5e ノードプールでホストされる指示チューニング済みモデル。ノードのマシンタイプは ct5lp-hightpu-4t です。

Saxml

Saxml は、推論に Paxml、JAX、PyTorch モデルを提供する試験運用版のシステムです。Saxml システムには、次のコンポーネントが含まれています。

Saxml セルまたは Sax クラスタ: 管理サーバーとモデルサーバーのグループ。管理サーバーはモデルサーバーを追跡し、公開されたモデルをモデルサーバーに割り当てて提供します。また、クライアントが特定の公開モデルを提供するモデルサーバーを特定できるようにします。
Saxml クライアント: Saxml システムのユーザー向けプログラミングインターフェース。Saxml クライアントには、コマンドラインツール（saxutil）と、Python、C++、Go のクライアントライブラリのスイートが含まれています。

このチュートリアルでは、Saxml HTTP サーバーも使用します。Saxml HTTP サーバーは、Saxml Python クライアントライブラリをカプセル化し、Saxml システムを操作する REST API を公開するカスタム HTTP サーバーです。REST API には、モデルの公開、一覧取得、公開停止、予測の生成を行うエンドポイントが含まれます。

目標

このチュートリアルは、JAX を使用している生成 AI をご利用のお客様、Gemma を提供するために Kubernetes コンテナのオーケストレーション機能を使用することを検討している GKE の新規または既存のユーザー（ML エンジニア、MLOps（DevOps）エンジニア、プラットフォーム管理者など）を対象としています。

このチュートリアルでは、次の手順について説明します。

モデルの特性に基づいて推奨される TPU トポロジを持つ GKE Standard クラスタを準備します。
GKE に Saxml コンポーネントをデプロイします。
Gemma 2B または Gemma 7B パラメータモデルを取得して公開します。
公開されたモデルを提供して操作します。

アーキテクチャ

このセクションでは、このチュートリアルで使用する GKE アーキテクチャについて説明します。このアーキテクチャは、TPU をプロビジョニングし、Gemma 2B または 7B モデルをデプロイして提供するための Saxml コンポーネントをホストする GKE Standard クラスタで構成されています。次の図は、このアーキテクチャのコンポーネントを示しています。

このチュートリアルでデプロイするアーキテクチャの図

このアーキテクチャには次のコンポーネントが含まれています。

GKE Standard ゾーンクラスタ。
提供する Gemma モデルに依存する単一ホストの TPU スライスノードプール:
- Gemma 2B: 1x1 トポロジの TPU v5e で構成されています。Saxml モデルサーバーの 1 つのインスタンスが、このノードプールを使用するように構成されています。
- Gemma 7B: 2x2 トポロジの TPU v5e で構成されています。Saxml モデルサーバーの 1 つのインスタンスが、このノードプールを使用するように構成されています。
Saxml 管理サーバーと Saxml HTTP サーバーがデプロイされているデフォルトの CPU ノードプール。
2 つの Cloud Storage バケット:
- 一方の Cloud Storage バケットには、管理サーバーによって管理される状態が保存されます。
- もう一方の Cloud Storage バケットには Gemma モデルのチェックポイントが保存されます。

このアーキテクチャには次の特徴があります。

一般公開の Artifact Registry は、Saxml コンポーネントのコンテナイメージを管理します。
GKE クラスタは GKE 用 Workload Identity 連携を使用します。すべての Saxml コンポーネントは、IAM サービスアカウントを統合する Workload Identity 連携を使用して、Cloud Storage バケットなどの外部サービスにアクセスします。
Saxml コンポーネントによって生成されたログは、Cloud Logging に統合されます。
Cloud Monitoring を使用して、このチュートリアルで作成する GKE ノードプールのパフォーマンス指標を分析できます。

始める前に

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 required 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 required API.

Enable the API

Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin
Check for the roles
1. In the Google Cloud console, go to the IAM page.
  Go to IAM
2. Select the project.
3. In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.
4. For all rows that specify or include you, check the Role colunn to see whether the list of roles includes the required roles.
Grant the roles
1. In the Google Cloud console, go to the IAM page.
  [IAM] に移動
2. プロジェクトを選択します。
3. [ アクセスを許可] をクリックします。
4. [新しいプリンシパル] フィールドに、ユーザー ID を入力します。これは通常、Google アカウントのメールアドレスです。
5. [ロールを選択] リストでロールを選択します。
6. 追加のロールを付与するには、 [別のロールを追加] をクリックして各ロールを追加します。
7. [保存] をクリックします。

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 用 Workload Identity 連携を使用する 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 トポロジと 1 つのノードを持つ 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 用 Workload Identity 連携を使用する 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 トポロジと 1 つのノードを持つ 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 バケットを作成する

Saxml 管理サーバーとモデルチェックポイントの状態を管理する 2 つの Cloud Storage バケットを作成します。

Cloud Shell で次のコマンドを実行します。

Saxml 管理サーバーの構成を保存する Cloud Storage バケットを作成します。
```
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 用 Workload Identity 連携を使用してワークロードアクセスを構成する

アプリケーションに Kubernetes ServiceAccount を割り当て、割り当てた Kubernetes ServiceAccount を IAM サービスアカウントとして機能するように構成します。

クラスタと通信を行うように kubectl を構成します。

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

アプリケーションで使用する Kubernetes ServiceAccount を作成します。
```
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

2 つのサービスアカウントの間に IAM ポリシーバインディングを追加して、Kubernetes ServiceAccount が IAM サービスアカウントの権限を借用できるようにします。このバインドにより、Kubernetes 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]"
```
Kubernetes サービスアカウントに IAM サービスアカウントのメールアドレスでアノテーションを付けます。
```
kubectl annotate serviceaccount default \
    iam.gke.io/gcp-service-account=wi-sax@${PROJECT_ID}.iam.gserviceaccount.com
```

モデルへのアクセス権を取得する

GKE にデプロイするために Gemma モデルへのアクセス権を取得するには、Kaggle プラットフォームにログインし、ライセンス同意契約に署名して、Kaggle API トークンを入手する必要があります。このチュートリアルでは、Kaggle 認証情報に Kubernetes Secret を使用します。

Gemma を使用するには同意契約に署名する必要があります。手順は次のとおりです。

Kaggle.com のモデルの同意ページにアクセスします。
Kaggle にログインしていない場合はログインします。
[アクセス権限をリクエスト] をクリックします。
[同意に使用するアカウントを選択] セクションで、[Kaggle アカウントを使用して確認] を選択して、同意に Kaggle アカウントを使用します。
モデルの利用規約に同意します。

アクセストークンを生成する

Kaggle からモデルにアクセスするには、Kaggle API トークンが必要です。

トークンをまだ生成していない場合は、次の手順に沿って生成します。

ブラウザで [Kaggle の設定] に移動します。
[API] セクションで [新しいトークンを作成] をクリックします。

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}

Kaggle 認証情報を保存する Secret を作成します。

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
```

モデルサーバーの Deployment のステータスを確認します。

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
```

モデルサーバーの Deployment のステータスを確認します。

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 の Service を作成します。

次の 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 Job を実行します。使用する 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
```

Job が完了するまで待ちます。

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

Job のログを表示します。

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
```

Job が完了するまで待ちます。

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

Job のログを表示します。

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

Job はチェックポイントを gs://CHECKPOINTS_BUCKET_NAME/gemma_7b_it/checkpoint_00000000 にアップロードします。

Saxml HTTP サーバーを公開する

Saxml HTTP サーバーには、Saxml HTTP サーバーをデプロイするときに作成した ClusterIP Service を介してアクセスできます。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 モデルを公開する

次に、TPU スライスのノードプールで実行されるモデルサーバーに Gemma モデルを公開できます。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 Deployment のモデルサーバー Pod でログを監視して、進行状況をモニタリングします。
```
kubectl logs --follow deployment/sax-model-server-v5e-1x1
```
このデプロイが完了するまでに、最長で 5 分を要する場合があります。次のようなメッセージが表示されるまで待ちます。
```
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 Deployment のモデルサーバー 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 というメッセージが表示された場合は、コンテナがモデルデータのダウンロードを完了していない可能性があります。モデルがサービス提供の準備ができていることを示す Connected というメッセージがないか、再度 Pod のログを確認します。
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 HTTP API など、Saxml の GitHub リポジトリを確認する。
Vertex AI Model Garden を確認する。
GKE プラットフォームのオーケストレーション機能を使用して、最適化された AI / ML ワークロードを実行する方法を確認する。
Google Cloud に関するリファレンスアーキテクチャ、図、ベストプラクティスを確認する。Cloud アーキテクチャセンターをご覧ください。

Saxml を備えた GKE で TPU を使用して Gemma オープンモデルを提供する

背景

Gemma

TPU

Saxml

目標

アーキテクチャ

始める前に

Check for the roles

Grant the roles

Gemma の環境を準備する

Cloud Shell を起動する

GKE Standard クラスタを作成する

Gemma 2B-it

Gemma 7B-it

Cloud Storage バケットを作成する

GKE 用 Workload Identity 連携を使用してワークロード アクセスを構成する

モデルへのアクセス権を取得する

ライセンス同意契約に署名する

アクセス トークンを生成する

アクセス トークンを 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

問題のトラブルシューティング

クリーンアップ

デプロイされたリソースを削除する

次のステップ

GKE 用 Workload Identity 連携を使用してワークロードアクセスを構成する

アクセストークンを生成する

アクセストークンを Cloud Shell にアップロードする