このガイドでは、Google Kubernetes Engine（GKE）で Tensor Processing Unit（TPU）を使用し、JetStream と MaxText で Gemma 大規模言語モデル（LLM）をサービングする方法について説明します。このガイドでは、Gemma 7B パラメータ指示用にチューニングされたモデルの重みを Cloud Storage にダウンロードし、JetStream を実行するコンテナを使用して GKE Autopilot または Standard クラスタにデプロイします。
このガイドは、モデルを JetStream にデプロイするときに Kubernetes の機能によって実現されるスケーラビリティ、復元力、費用対効果が必要な場合に適しています。
背景
GKE で TPU を使用して JetStream で Gemma をサービングすることで、マネージド Kubernetes のメリット（費用効率、スケーラビリティ、高可用性など）をすべて活用した、本番環境対応の堅牢なサービング ソリューションを構築できます。このセクションでは、このチュートリアルで使用されている重要なテクノロジーについて説明します。
Gemma
Gemma は、オープン ライセンスでリリースされ一般公開されている、軽量の生成 AI モデルのセットです。これらの AI モデルは、アプリケーション、ハードウェア、モバイル デバイス、ホスト型サービスで実行できます。Gemma モデルはテキスト生成に使用できますが、特殊なタスク用にチューニングすることもできます。
詳しくは、Gemma のドキュメントをご覧ください。
TPU
TPU は、Google が独自に開発した特定用途向け集積回路（ASIC）であり、TensorFlow、PyTorch、JAX などのフレームワークを使用して構築された機械学習モデルと AI モデルを高速化するために使用されます。
GKE で TPU を使用する前に、次の学習プログラムを完了することをおすすめします。
- Cloud TPU システム アーキテクチャで、現在の TPU バージョンの可用性について学習する。
- GKE の TPU についてを確認する。
このチュートリアルでは、Gemma 7B モデルのサービングについて説明します。GKE は、低レイテンシでプロンプトをサービングするモデルの要件に基づいて構成された TPU トポロジを使用して、単一ホストの TPUv5e ノードにモデルをデプロイします。
JetStream
JetStream は、Google が開発したオープンソースの推論サービング フレームワークです。JetStream を使用すると、TPU と GPU で高性能、高スループット、メモリ最適化された推論が可能になります。継続的なバッチ処理や量子化技術などの高度なパフォーマンス最適化により、LLM のデプロイを容易にします。JetStream では、PyTorch / XLA と JAX TPU のサービングにより、最適なパフォーマンスを実現できます。
これらの最適化の詳細については、JetStream PyTorch と JetStream MaxText のプロジェクト リポジトリをご覧ください。
MaxText
MaxText は、Flax、Orbax、Optax などのオープンソースの JAX ライブラリ上に構築された、パフォーマンス、スケーラビリティ、適応性に優れた JAX LLM 実装です。MaxText のデコーダ専用の LLM 実装は Python で記述されています。XLA コンパイラの活用により、カスタム カーネルを構築しなくても高いパフォーマンスを実現できます。
MaxText がサポートする最新のモデルとパラメータ サイズの詳細については、MaxtText プロジェクト リポジトリをご覧ください。
目標
このチュートリアルは、JAX を使用している生成 AI をご利用のお客様、GKE の新規または既存のユーザー、ML エンジニア、MLOps（DevOps）エンジニア、LLM のサービングに Kubernetes コンテナのオーケストレーション機能を使用することに関心をお持ちのプラットフォーム管理者を対象としています。
このチュートリアルでは、次の手順について説明します。
- モデルの特性に基づいて推奨される TPU トポロジを持つ GKE Autopilot または Standard クラスタを準備します。
- GKE に JetStream コンポーネントをデプロイします。
- Gemma 7B 指示用にチューニングされたモデルを取得して公開します。
- 公開されたモデルをサービングして操作します。
アーキテクチャ
このセクションでは、このチュートリアルで使用する GKE アーキテクチャについて説明します。このアーキテクチャは、TPU をプロビジョニングし、モデルをデプロイしてサービングするための JetStream コンポーネントをホストする GKE Autopilot または Standard クラスタで構成されています。
次の図は、このアーキテクチャのコンポーネントを示しています。
このアーキテクチャには次のコンポーネントが含まれています。
- GKE Autopilot または Standard リージョン クラスタ。
- JetStream デプロイをホストする 2 つの単一ホスト TPU スライス ノードプール。
- Service コンポーネントは、インバウンド トラフィックをすべての
JetStream HTTPレプリカに分散します。 JetStream HTTPは、JetStream の必須フォーマットのラッパーとしてリクエストを受け取り、JetStream の GRPC クライアントに送信する HTTP サーバーです。Maxengineは、継続的バッチ処理で推論を行う JetStream サーバーです。
始める前に
- 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.
-
Google Cloud Console の [プロジェクト セレクタ] ページで、Google Cloud プロジェクトを選択または作成します。
-
必要な API を有効にします。
-
Google Cloud Console の [プロジェクト セレクタ] ページで、Google Cloud プロジェクトを選択または作成します。
-
必要な API を有効にします。
-
プロジェクトに次のロールがあることを確認します。 roles/container.admin, roles/iam.serviceAccountAdmin
ロールを確認する
-
Google Cloud コンソールの [IAM] ページに移動します。
[IAM] に移動 - プロジェクトを選択します。
-
[プリンシパル] 列で、自分のメールアドレスを含む行を見つけます。
自分のメールアドレスがその列にない場合、ロールは割り当てられていません。
- 自分のメールアドレスを含む行の [ロール] 列で、ロールのリストに必要なロールが含まれているかどうかを確認します。
ロールを付与する
-
Google Cloud コンソールの [IAM] ページに移動します。
[IAM] に移動 - プロジェクトを選択します。
- [ アクセスを許可] をクリックします。
- [新しいプリンシパル] フィールドに、自分のメールアドレスを入力します。
- [ロールを選択] リストでロールを選択します。
- 追加のロールを付与するには、 [別のロールを追加] をクリックして各ロールを追加します。
- [保存] をクリックします。
-
- 8 個の TPU v5e PodSlice Lite チップに十分な割り当てがあることを確認します。このチュートリアルでは、オンデマンド インスタンスを使用します。
- Kaggle アカウントを作成します（まだアカウントを保有されていない場合）。
モデルへのアクセス権を取得する
GKE にデプロイするために Gemma モデルへのアクセス権を取得するには、まずライセンス同意契約に署名する必要があります。
ライセンス同意契約に署名する
Gemma を使用するには同意契約に署名する必要があります。手順は次のとおりです。
- Kaggle.com の Gemma モデルの同意ページにアクセスします。
- Kaggle にログインしていない場合はログインします。
- [アクセス権限をリクエスト] をクリックします。
- [同意に使用するアカウントを選択] セクションで、[Kaggle アカウントを使用して確認] を選択して、同意に Kaggle アカウントを使用します。
- モデルの利用規約に同意します。
アクセス トークンを生成する
Kaggle からモデルにアクセスするには、Kaggle API トークンが必要です。
トークンをまだ生成していない場合は、次の手順に沿って生成します。
- ブラウザで [Kaggle の設定] に移動します。
- [API] セクションで [新しいトークンを作成] をクリックします。
kaggle.json という名前のファイルがダウンロードされます。
環境を準備する
このチュートリアルでは、Cloud Shell を使用して Google Cloud でホストされるリソースを管理します。Cloud Shell には、このチュートリアルに必要な kubectl や gcloud CLI などのソフトウェアがプリインストールされています。
Cloud Shell を使用して環境を設定するには、次の操作を行います。
Google Cloud コンソールで、Google Cloud コンソールの
[Cloud Shell をアクティブにする] をクリックして、Cloud Shell セッションを起動します。これにより、Google Cloud コンソールの下部ペインでセッションが起動されます。デフォルトの環境変数を設定します。
gcloud config set project PROJECT_ID export PROJECT_ID=$(gcloud config get project) export CLUSTER_NAME=CLUSTER_NAME export BUCKET_NAME=BUCKET_NAME export REGION=REGION export LOCATION=LOCATION次の値を置き換えます。
- PROJECT_ID: Google Cloud のプロジェクト ID。
- CLUSTER_NAME: GKE クラスタの名前。
- BUCKET_NAME: Cloud Storage バケットの名前。
gs://プレフィックスを指定する必要はありません。 - REGION_NAME: GKE クラスタ、Cloud Storage バケット、TPU ノードが配置されているリージョン。TPU v5e マシンタイプを使用できるゾーン（
us-west1、us-west4、us-central1、us-east1、us-east5、europe-west4など）が含まれているリージョンです。 - （Standard クラスタのみ）LOCATION: TPU リソースが使用可能なゾーン（
us-west4-aなど）。Autopilot クラスタの場合は、リージョンのみを指定します。ゾーンを指定する必要はありません。
Google Cloud リソースを作成して構成する
次の手順に沿って、必要なリソースを作成します。
GKE クラスタを作成する
GKE Autopilot クラスタまたは GKE Standard クラスタの GPU で Gemma をサービングできます。フルマネージドの Kubernetes エクスペリエンスを実現するには、Autopilot クラスタを使用することをおすすめします。ワークロードに最適な GKE の運用モードを選択するには、GKE の運用モードを選択するをご覧ください。
Autopilot
Cloud Shell で、次のコマンドを実行します。
gcloud container clusters create-auto ${CLUSTER_NAME} \
--project=${PROJECT_ID} \
--region=${REGION} \
--release-channel=rapid \
--cluster-version=1.29
GKE は、デプロイされたワークロードからのリクエストに応じた CPU ノードと TPU ノードを持つ Autopilot クラスタを作成します。
Standard
GKE 用 Workload Identity 連携を使用するリージョン 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=${REGION}クラスタの作成には数分かかることもあります。
次のコマンドを実行して、クラスタのノードプールを作成します。
gcloud container node-pools create gemma-7b-tpu-nodepool \ --cluster=${CLUSTER_NAME} \ --machine-type=ct5lp-hightpu-8t \ --project=${PROJECT_ID} \ --num-nodes=2 \ --region=${REGION} \ --node-locations=${LOCATION}GKE は、
2x4トポロジと 2 つのノードを持つ TPU v5e ノードプールを作成します。
Cloud Storage バケットを作成する
Cloud Shell で、次のコマンドを実行します。
gcloud storage buckets create gs://${BUCKET_NAME} --location=${REGION}
これにより、Kaggle からダウンロードしたモデルファイルを格納する Cloud Storage バケットが作成されます。
アクセス トークンを 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=${REGION}Kaggle 認証情報を保存する Secret を作成します。
kubectl create secret generic kaggle-secret \ --from-file=kaggle.json
GKE 用 Workload Identity 連携を使用してワークロード アクセスを構成する
アプリケーションに Kubernetes ServiceAccount を割り当て、IAM サービス アカウントとして機能するようにその Kubernetes ServiceAccount を構成します。
アプリケーションの IAM サービス アカウントを作成します。
gcloud iam service-accounts create wi-jetstreamCloud Storage を管理する IAM サービス アカウントの IAM ポリシー バインディングを追加します。
gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member "serviceAccount:wi-jetstream@${PROJECT_ID}.iam.gserviceaccount.com" \ --role roles/storage.objectUser gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member "serviceAccount:wi-jetstream@${PROJECT_ID}.iam.gserviceaccount.com" \ --role roles/storage.insightsCollectorService2 つのサービス アカウントの間に IAM ポリシー バインディングを追加して、Kubernetes ServiceAccount が IAM サービス アカウントの権限を借用できるようにします。このバインドにより、Kubernetes ServiceAccount が IAM サービス アカウントとして機能します。
gcloud iam service-accounts add-iam-policy-binding wi-jetstream@${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-jetstream@${PROJECT_ID}.iam.gserviceaccount.com
モデルのチェックポイントを変換する
このセクションでは、次の処理を行う Job を作成します。
- Kaggle から基本的な Orbax チェックポイントをダウンロードします。
- チェックポイントを Cloud Storage バケットにアップロードします。
- チェックポイントを MaxText 対応のチェックポイントに変換します。
- サービングに使用するチェックポイントのスキャンを解除します。
モデルのチェックポイント変換 Job をデプロイする
次の手順に沿って、Gemma 7B モデルのチェックポイント ファイルをダウンロードして変換します。
次のマニフェストを
job-7b.yamlとして作成します。apiVersion: batch/v1 kind: Job metadata: name: data-loader-7b spec: ttlSecondsAfterFinished: 30 template: spec: restartPolicy: Never containers: - name: inference-checkpoint image: us-docker.pkg.dev/cloud-tpu-images/inference/inference-checkpoint:v0.2.0 args: - -b=BUCKET_NAME - -m=google/gemma/maxtext/7b-it/2 volumeMounts: - mountPath: "/kaggle/" name: kaggle-credentials readOnly: true resources: requests: google.com/tpu: 8 limits: google.com/tpu: 8 nodeSelector: cloud.google.com/gke-tpu-topology: 2x4 cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice volumes: - name: kaggle-credentials secret: defaultMode: 0400 secretName: kaggle-secret次のようにマニフェストを適用します。
kubectl apply -f job-7b.yamlJob のログを表示します。
kubectl logs -f jobs/data-loader-7bJob が完了すると、出力は次のようになります。
Successfully generated decode checkpoint at: gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items + echo -e '\nCompleted unscanning checkpoint to gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items' Completed unscanning checkpoint to gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items
JetStream をデプロイする
このセクションでは、Gemma モデルをサービングする JetStream コンテナをデプロイします。
次の手順に沿って、Gemma 7B 指示用にチューニングされたモデルをデプロイします。
次の
jetstream-gemma-deployment.yamlマニフェストを作成します。apiVersion: apps/v1 kind: Deployment metadata: name: maxengine-server spec: replicas: 2 selector: matchLabels: app: maxengine-server template: metadata: labels: app: maxengine-server spec: nodeSelector: cloud.google.com/gke-tpu-topology: 2x4 cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice containers: - name: maxengine-server image: us-docker.pkg.dev/cloud-tpu-images/inference/maxengine-server:v0.2.0 args: - model_name=gemma-7b - tokenizer_path=assets/tokenizer.gemma - per_device_batch_size=4 - max_prefill_predict_length=1024 - max_target_length=2048 - async_checkpointing=false - ici_fsdp_parallelism=1 - ici_autoregressive_parallelism=-1 - ici_tensor_parallelism=1 - scan_layers=false - weight_dtype=bfloat16 - load_parameters_path=gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items ports: - containerPort: 9000 resources: requests: google.com/tpu: 8 limits: google.com/tpu: 8 - name: jetstream-http image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.0 ports: - containerPort: 8000 --- apiVersion: v1 kind: Service metadata: name: jetstream-http-svc spec: selector: app: maxengine-server ports: - protocol: TCP port: 8000 targetPort: 8000このマニフェストでは、次の主要なプロパティが設定されています。
tokenizer_path: モデルのトークナイザのパス。load_parameters_path: チェックポイントが保存される Cloud Storage バケット内のパス。per_device_batch_size: デバイスごとのデコード バッチサイズ。1 つの TPU チップは 1 つのデバイスに相当します。max_prefill_predict_length: 自動回帰を行う場合のプリフィルの最大長。max_target_length: 最大シーケンス長。model_name: モデル名（gemma-7b）。ici_fsdp_parallelism: 完全にシャーディングされたデータ並列処理（FSDP）のシャードの数。ici_tensor_parallelism: テンソル並列処理のシャードの数。ici_autoregressive_parallelism: 自己回帰並列処理のシャードの数。scan_layers: スキャンレイヤのブール値フラグ（ブール値）。weight_dtype: 重みのデータ型（bfloat16）。
次のようにマニフェストを適用します。
kubectl apply -f jetstream-gemma-deployment.yamlDeployment を確認します。
kubectl get deployment出力は次のようになります。
NAME READY UP-TO-DATE AVAILABLE AGE maxengine-server 2/2 2 2 ##sHTTP サーバーログを表示して、モデルが読み込まれ、コンパイルされていることを確認します。このオペレーションが完了するまでに数分かかることがあります。
kubectl logs deploy/maxengine-server -f -c jetstream-http出力は次のようになります。
kubectl logs deploy/maxengine-server -f -c jetstream-http INFO: Started server process [1] INFO: Waiting for application startup. INFO: Application startup complete. INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)MaxEngine ログを表示して、コンパイルが完了したことを確認します。
kubectl logs deploy/maxengine-server -f -c maxengine-server出力は次のようになります。
2024-03-29 17:09:08,047 - jax._src.dispatch - DEBUG - Finished XLA compilation of jit(initialize) in 0.26236414909362793 sec 2024-03-29 17:09:08,150 - root - INFO - ---------Generate params 0 loaded.---------
モデルをサービングする
このセクションでは、モデルを操作します。
ポート転送をセットアップする
JetStream Deployment には、前の手順で作成した ClusterIP Service を介してアクセスできます。ClusterIP Service にはクラスタ内からのみアクセスできます。したがって、クラスタの外部から Service にアクセスするには、次の手順を完了します。
ポート転送セッションを確立するには、次のコマンドを実行します。
kubectl port-forward svc/jetstream-http-svc 8000:8000
curl を使用してモデルを操作する
新しいターミナルを開いて次のコマンドを実行し、JetStream HTTP サーバーにアクセスできることを確認します。
curl --request POST \ --header "Content-type: application/json" \ -s \ localhost:8000/generate \ --data \ '{ "prompt": "What are the top 5 programming languages", "max_tokens": 200 }'モデルのウォームアップにより、最初のリクエストが完了するまでに数秒かかることがあります。出力は次のようになります。
{ "response": "\nfor data science in 2023?\n\n**1. Python:**\n- Widely used for data science due to its simplicity, readability, and extensive libraries for data wrangling, analysis, visualization, and machine learning.\n- Popular libraries include pandas, scikit-learn, and matplotlib.\n\n**2. R:**\n- Statistical programming language widely used for data analysis, visualization, and modeling.\n- Popular libraries include ggplot2, dplyr, and caret.\n\n**3. Java:**\n- Enterprise-grade language with strong performance and scalability.\n- Popular libraries include Spark, TensorFlow, and Weka.\n\n**4. C++:**\n- High-performance language often used for data analytics and machine learning models.\n- Popular libraries include TensorFlow, PyTorch, and OpenCV.\n\n**5. SQL:**\n- Relational database language essential for data wrangling and querying large datasets.\n- Popular tools" }
（省略可）Gradio のチャット インターフェースでモデルを操作する
このセクションでは、指示用にチューニングされたモデルを操作できるウェブチャット アプリケーションを作成します。
Gradio は、chatbot のユーザー インターフェースを作成する ChatInterface ラッパーを含む Python ライブラリです。
チャット インターフェースをデプロイする
Cloud Shell で、次の Ingress マニフェストを
gradio.yamlとして保存します。apiVersion: apps/v1 kind: Deployment metadata: name: gradio labels: app: gradio spec: replicas: 1 selector: matchLabels: app: gradio template: metadata: labels: app: gradio spec: containers: - name: gradio image: us-docker.pkg.dev/google-samples/containers/gke/gradio-app:v1.0.3 resources: requests: cpu: "512m" memory: "512Mi" limits: cpu: "1" memory: "512Mi" env: - name: CONTEXT_PATH value: "/generate" - name: HOST value: "http://jetstream-http-svc:8000" - name: LLM_ENGINE value: "max" - name: MODEL_ID value: "gemma" - name: USER_PROMPT value: "<start_of_turn>user\nprompt<end_of_turn>\n" - name: SYSTEM_PROMPT value: "<start_of_turn>model\nprompt<end_of_turn>\n" ports: - containerPort: 7860 --- apiVersion: v1 kind: Service metadata: name: gradio spec: selector: app: gradio ports: - protocol: TCP port: 8080 targetPort: 7860 type: ClusterIP次のようにマニフェストを適用します。
kubectl apply -f gradio.yamlDeployment が利用可能になるまで待ちます。
kubectl wait --for=condition=Available --timeout=300s deployment/gradio
チャット インターフェースを使用する
Cloud Shell で、次のコマンドを実行します。
kubectl port-forward service/gradio 8080:8080これにより、Cloud Shell から Gradio サービスへのポート転送が作成されます。
Cloud Shell タスクバーの右上にある [
ウェブでプレビュー] ボタンをクリックします。[ポート 8080 でプレビュー] をクリックします。ブラウザで新しいタブが開きます。Gradio のチャット インターフェースを使用して Gemma を操作します。プロンプトを追加して [送信] をクリックします。
問題のトラブルシューティング
Empty reply from serverというメッセージが表示された場合は、コンテナがモデルデータのダウンロードを完了していない可能性があります。モデルがサービス提供の準備ができていることを示すConnectedというメッセージがないか、再度 Pod のログを確認します。Connection refusedが表示された場合は、ポート転送が有効であることを確認します。
クリーンアップ
このチュートリアルで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。
デプロイされたリソースを削除する
このガイドで作成したリソースに対して Google Cloud アカウントで課金されないようにするには、次のコマンドを実行し、プロンプトに従います。
gcloud container clusters delete ${CLUSTER_NAME} --region=${REGION}
gcloud iam service-accounts delete wi-jetstream@PROJECT_ID.iam.gserviceaccount.com
gcloud storage rm --recursive gs://BUCKET_NAME
次のステップ
- GKE で Gemma モデルを実行する方法と、GKE プラットフォームのオーケストレーション機能を使用して最適化された AI / ML ワークロードを実行する方法を確認する。
- GKE の TPU の詳細を確認する。
- JetStream の GitHub リポジトリを確認する。
- Vertex AI Model Garden を確認する。