MaxDiffusion を備えた GKE で TPU を使用して Stable Diffusion XL（SDXL）を提供する

Autopilot Standard

---

このチュートリアルでは、MaxDiffusion を備えた Google Kubernetes Engine（GKE）上で Tensor Processing Unit（TPU）を使用して SDXL 画像生成モデルを提供する方法について説明します。このチュートリアルでは、Hugging Face からモデルをダウンロードし、MaxDiffusion を実行するコンテナを使用して Autopilot クラスタまたは Standard クラスタにデプロイします。

このガイドは、AI / ML ワークロードをデプロイして提供する際に、マネージド Kubernetes での詳細な制御、カスタマイズ、拡張性、復元力、ポータビリティ、費用対効果が求められる場合の出発点として適しています。ML モデルを費用対効果の高い方法で迅速に構築して提供するために、統合されたマネージド AI プラットフォームが必要な場合は、Vertex AI デプロイ ソリューションをお試しになることをおすすめします。

背景

GKE で TPU を使用して MaxDiffusion で SDXL をサービングすることで、マネージド Kubernetes のメリット（費用効率、拡張性、高可用性など）をすべて活用した、本番環境対応の堅牢なサービング ソリューションを構築できます。このセクションでは、このチュートリアルで使用されている重要なテクノロジーについて説明します。

Stable Diffusion XL（SDXL）

Stable Diffusion XL（SDXL）は、推論用に MaxDiffusion がサポートする潜在的な拡散モデル（LDM）の一種です。生成 AI の場合、LDM を使用してテキストの説明から高品質の画像を生成できます。LDM は、画像検索や画像キャプションなどのアプリケーションに役立ちます。

SDXL は、シャーディング アノテーションを使用して、単一ホストまたはマルチホストの推論をサポートします。これにより、SDXL を複数のマシンでトレーニングして実行できるため、効率性が向上します。

詳細については、Stability AI リポジトリによる生成モデルと SDXL の論文をご覧ください。

TPU

TPU は、Google が独自に開発した特定用途向け集積回路（ASIC）であり、TensorFlow、PyTorch、JAX などのフレームワークを使用して構築された機械学習モデルと AI モデルを高速化するために使用されます。

GKE で TPU を使用する前に、次の学習プログラムを完了することをおすすめします。

1. Cloud TPU システム アーキテクチャで、現在の TPU バージョンの可用性について学習する。
2. GKE の TPU についてを確認する。

このチュートリアルでは、SDXL モデルのサービングについて説明します。GKE は、低レイテンシでプロンプトをサービングするモデルの要件に基づいて構成された TPU トポロジを使用して、単一ホストの TPU v5e ノードにモデルをデプロイします。このガイドでは、モデルで 1x1 トポロジの TPU v5e チップを使用します。

MaxDiffusion

MaxDiffusion は、Python と Jax で記述された、XLA デバイス（TPU や GPU など）で実行されるさまざまな潜在的な拡散モデルのリファレンス実装のコレクションです。MaxDiffusion は、研究環境と本番環境の両方での Diffusion プロジェクトの開始点です。

詳細については、MaxDiffusion リポジトリをご覧ください。

目標

このチュートリアルは、JAX を使用している生成 AI をご利用のお客様、SDXL の新規または既存のユーザー、ML エンジニア、MLOps（DevOps）エンジニア、LLM のサービングに Kubernetes コンテナのオーケストレーション機能を使用することに関心をお持ちのプラットフォーム管理者を対象としています。

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

1. モデルの特性に基づいて推奨される TPU トポロジを持つ GKE Autopilot または Standard クラスタを作成します。
2. SDXL 推論コンテナ イメージをビルドします。
3. GKE に SDXL 推論サーバーをデプロイします。
4. ウェブアプリを介してモデルをサービングして操作します。

アーキテクチャ

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

次の図は、このアーキテクチャのコンポーネントを示しています。

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

• GKE Autopilot または Standard リージョン クラスタ。
• MaxDiffusion のデプロイで SDXL モデルをホストする 1 つの単一ホストの TPU スライス ノードプール。
• ClusterIP タイプのロードバランサを持つ Service コンポーネント。この Service は、すべての MaxDiffusion HTTP レプリカにインバウンド トラフィックを分散します。
• インバウンド トラフィックを分散し、モデル サービング トラフィックを ClusterIP Service にリダイレクトする外部 LoadBalancer Service を備えた WebApp HTTP サーバー。

始める前に

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.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the 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.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the 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, roles/artifactregistry.admin

  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 column 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. [保存] をクリックします。
  • TPU v5e PodSlice Lite チップに十分な割り当てがあることを確認します。このチュートリアルでは、オンデマンド インスタンスを使用します。

  環境を準備する

  このチュートリアルでは、Cloud Shell を使用してGoogle Cloudでホストされているリソースを管理します。Cloud Shell には、このチュートリアルに必要な kubectl や gcloud CLI などのソフトウェアがプリインストールされています。

  Cloud Shell を使用して環境を設定するには、次の操作を行います。

  1. Google Cloud コンソールで （Cloud Shell をアクティブにする）をクリックして、Google Cloud コンソールで Cloud Shell セッションを起動します。これにより、 Google Cloud コンソールの下部ペインでセッションが起動します。

  2. デフォルトの環境変数を設定します。

     gcloud config set project PROJECT_ID
     gcloud config set billing/quota_project PROJECT_ID
     export PROJECT_ID=$(gcloud config get project)
     export CLUSTER_NAME=CLUSTER_NAME
     export REGION=REGION_NAME
     export ZONE=ZONE
     

     次の値を置き換えます。

     • PROJECT_ID: 実際の Google Cloud プロジェクト ID。
     • CLUSTER_NAME: GKE クラスタの名前。
     • REGION_NAME: GKE クラスタ、Cloud Storage バケット、TPU ノードが配置されているリージョン。TPU v5e マシンタイプを使用できるゾーン（us-west1、us-west4、us-central1、us-east1、us-east5、europe-west4 など）が含まれているリージョンです。
     • （Standard クラスタのみ）ZONE: TPU リソースが使用可能なゾーン（us-west4-a など）。Autopilot クラスタの場合は、リージョンのみを指定します。ゾーンを指定する必要はありません。

  3. サンプル リポジトリのクローンを作成し、チュートリアル ディレクトリを開きます。

     git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples
     cd kubernetes-engine-samples/ai-ml/maxdiffusion-tpu
     WORK_DIR=$(pwd)
     gcloud artifacts repositories create gke-llm --repository-format=docker --location=$REGION
     gcloud auth configure-docker $REGION-docker.pkg.dev
     

  Google Cloud リソースを作成して構成する

  次の手順で、必要なリソースを作成します。

  注: 一部のアクセラレータについては、使用にあたって容量の予約の作成が必要になる場合があります。リソースの予約方法と予約済みリソースの使用方法については、予約済みゾーンリソースの使用をご覧ください。

  GKE クラスタを作成する

  GKE Autopilot クラスタまたは GKE Standard クラスタの GPU で SDXL をサービングできます。フルマネージドの Kubernetes エクスペリエンスを実現するには、Autopilot クラスタを使用することをおすすめします。ワークロードに最適な GKE の運用モードを選択するには、GKE の運用モードを選択するをご覧ください。

  Autopilot

  1. 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 クラスタを作成します。

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

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

  標準

  1. GKE 用 Workload Identity 連携を使用するリージョン GKE Standard クラスタを作成します。

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

     クラスタの作成には数分かかることもあります。

  2. 次のコマンドを実行して、クラスタのノードプールを作成します。

     gcloud container node-pools create maxdiffusion-tpu-nodepool \
       --cluster=${CLUSTER_NAME} \
       --machine-type=ct5lp-hightpu-1t \
       --num-nodes=1 \
       --region=${REGION} \
       --node-locations=${ZONE} \
       --spot
     

     GKE は、1x1 トポロジと 1 つのノードを持つ TPU v5e ノードプールを作成します。

     さまざまなトポロジのノードプールを作成するには、TPU 構成を計画する方法を学習します。このチュートリアルのサンプル値（cloud.google.com/gke-tpu-topology や google.com/tpu など）を必ず更新してください。

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

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

  SDXL 推論コンテナをビルドする

  次の手順に沿って、SDXL 推論サーバーのコンテナ イメージをビルドします。

  1. build/server/cloudbuild.yaml マニフェストを開きます。

     steps:
     - name: 'gcr.io/cloud-builders/docker'
       args: [ 'build', '-t', '$LOCATION-docker.pkg.dev/$PROJECT_ID/gke-llm/max-diffusion:latest', '.' ]
     images:
     - '$LOCATION-docker.pkg.dev/$PROJECT_ID/gke-llm/max-diffusion:latest'

  2. ビルドを実行して、推論コンテナ イメージを作成します。

     cd $WORK_DIR/build/server
     gcloud builds submit . --region=$REGION
     

     出力には、コンテナ イメージのパスが含まれます。

  SDXL 推論サーバーをデプロイする

  このセクションでは、SDXL 推論サーバーをデプロイします。このチュートリアルでは、サーバーをデプロイするために Kubernetes Deployment を使用します。Deployment は、クラスタ内のノードに分散された Pod の複数のレプリカを実行できる Kubernetes API オブジェクトです。

  1. serve_sdxl_v5e.yaml マニフェストを確認します。

     apiVersion: apps/v1
     kind: Deployment
     metadata:
       name: stable-diffusion-deployment
     spec:
       selector:
         matchLabels:
           app: max-diffusion-server
       replicas: 1  # number of nodes in node-pool
       template:
         metadata:
           labels:
             app: max-diffusion-server
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-topology: 1x1 #  target topology
             cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
             #cloud.google.com/gke-spot: "true"
           volumes:
           - name: dshm
             emptyDir:
                   medium: Memory
           containers:
           - name: serve-stable-diffusion
             image: REGION-docker.pkg.dev/PROJECT_ID/gke-llm/max-diffusion:latest
             env:
             - name: MODEL_NAME
               value: 'stable_diffusion'
             ports:
             - containerPort: 8000
             resources:
               requests:
                 google.com/tpu: 1  # TPU chip request
               limits:
                 google.com/tpu: 1  # TPU chip request
             volumeMounts:
                 - mountPath: /dev/shm
                   name: dshm
     
     ---
     apiVersion: v1
     kind: Service
     metadata:
       name: max-diffusion-server
       labels:
         app: max-diffusion-server
     spec:
       type: ClusterIP
       ports:
         - port: 8000
           targetPort: 8000
           name: http-max-diffusion-server
           protocol: TCP
       selector:
         app: max-diffusion-server

  2. マニフェスト内のプロジェクト ID を更新します。

     cd $WORK_DIR
     perl -pi -e 's|PROJECT_ID|PROJECT_ID|g' serve_sdxl_v5e.yaml
     perl -pi -e 's|REGION|REGION_NAME|g' serve_sdxl_v5e.yaml
     

  3. 次のようにマニフェストを適用します。

     kubectl apply -f serve_sdxl_v5e.yaml
     

     出力は次のようになります。

     deployment.apps/max-diffusion-server created
     

  4. モデルのステータスを確認します。

     kubectl get deploy --watch
     

     出力は次のようになります。

     NAME                          READY   UP-TO-DATE   AVAILABLE   AGE
     stable-diffusion-deployment   1/1     1            1           8m21s
     

  5. ClusterIP アドレスを取得します。

     kubectl get service max-diffusion-server
     

     出力には ClusterIP フィールドが含まれます。CLUSTER-IP の値をメモします。

  6. デプロイを検証します。

      export ClusterIP=CLUSTER_IP
      kubectl run curl --image=curlimages/curl \
         -it --rm --restart=Never \
         -- "$ClusterIP:8000"
     

     CLUSTER_IP は、前にメモした CLUSTER-IP の値に置き換えます。出力は次のようになります。

     {"message":"Hello world! From FastAPI running on Uvicorn with Gunicorn."}
     pod "curl" deleted
     

  7. Deployment のログを表示します。

     kubectl logs -l app=max-diffusion-server
     

     Deployment が完了すると、出力は次のようになります。

     2024-06-12 15:45:45,459 [INFO] __main__: replicate params:
     2024-06-12 15:45:46,175 [INFO] __main__: start initialized compiling
     2024-06-12 15:45:46,175 [INFO] __main__: Compiling ...
     2024-06-12 15:45:46,175 [INFO] __main__: aot compiling:
     2024-06-12 15:45:46,176 [INFO] __main__: tokenize prompts:2024-06-12 15:48:49,093 [INFO] __main__: Compiled in 182.91802048683167
     INFO:     Started server process [1]
     INFO:     Waiting for application startup.
     INFO:     Application startup complete.
     

  ウェブアプリ クライアントをデプロイする

  このセクションでは、SDXL モデルを提供する webapp クライアントをデプロイします。

  1. build/webapp/cloudbuild.yaml マニフェストを確認します。

     steps:
     - name: 'gcr.io/cloud-builders/docker'
       args: [ 'build', '-t', '$LOCATION-docker.pkg.dev/$PROJECT_ID/gke-llm/max-diffusion-web:latest', '.' ]
     images:
     - '$LOCATION-docker.pkg.dev/$PROJECT_ID/gke-llm/max-diffusion-web:latest'

  2. ビルドを実行し、build/webapp ディレクトリにクライアント コンテナ イメージを作成します。

     cd $WORK_DIR/build/webapp
     gcloud builds submit . --region=$REGION
     

     出力には、コンテナ イメージのパスが含まれます。

  3. serve_sdxl_client.yaml マニフェストを開きます。

     apiVersion: apps/v1
     kind: Deployment
     metadata:
       name: max-diffusion-client
     spec:
       selector:
         matchLabels:
           app: max-diffusion-client
       template:
         metadata:
           labels:
             app: max-diffusion-client
         spec:
           containers:
           - name: webclient
             image: REGION-docker.pkg.dev/PROJECT_ID/gke-llm/max-diffusion-web:latest
             env:
               - name: SERVER_URL
                 value: "http://ClusterIP:8000"
             resources:
               requests:
                 memory: "128Mi"
                 cpu: "250m"
               limits:
                 memory: "256Mi"
                 cpu: "500m"
             ports:
             - containerPort: 5000
     ---
     apiVersion: v1
     kind: Service
     metadata:
       name: max-diffusion-client-service
     spec:
       type: LoadBalancer
       selector:
         app: max-diffusion-client
       ports:
       - port: 8080
         targetPort: 5000

  4. マニフェストでプロジェクト ID を編集します。

     cd $WORK_DIR
     perl -pi -e 's|PROJECT_ID|PROJECT_ID|g' serve_sdxl_client.yaml
     perl -pi -e 's|ClusterIP|CLUSTER_IP|g' serve_sdxl_client.yaml
     perl -pi -e 's|REGION|REGION_NAME|g' serve_sdxl_client.yaml
     

  5. 次のようにマニフェストを適用します。

     kubectl apply -f serve_sdxl_client.yaml
     

  6. LoadBalancer IP アドレスを取得します。

     kubectl get service max-diffusion-client-service
     

     出力には LoadBalancer フィールドが含まれます。EXTERNAL-IP の値をメモします。

  成功: GKE で MaxDiffusion と TPU v5e を使用して SDXL モデルを正常にデプロイしました。

  ウェブページを使用してモデルを操作する

  1. ウェブブラウザから次の URL にアクセスします。

     http://EXTERNAL_IP:8080
     

     EXTERNAL_IP は、以前にメモした EXTERNAL_IP の値に置き換えます。

  2. チャット インターフェースを使用して SDXL を操作します。プロンプトを追加して [送信] をクリックします。例:

     Create a detailed image of a fictional historical site, capturing its unique architecture and cultural significance
     

  出力は、次の例のようなモデル生成画像です。

  

  クリーンアップ

  このチュートリアルで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

  プロジェクトを削除する

  注意: プロジェクトを削除すると、次のような影響があります。

  • プロジェクト内のすべてのものが削除されます。このドキュメントのタスクで既存のプロジェクトを使用した場合、それを削除すると、そのプロジェクトで行った他の作業もすべて削除されます。
  • カスタム プロジェクト ID が失われます。このプロジェクトを作成したときに、将来使用するカスタム プロジェクト ID を作成した可能性があります。そのプロジェクト ID を使用した URL（たとえば、appspot.com）を保持するには、プロジェクト全体ではなくプロジェクト内の選択したリソースだけを削除します。

  複数のアーキテクチャ、チュートリアル、クイックスタートを実施する予定がある場合は、プロジェクトを再利用すると、プロジェクトの割り当て上限を超えないようにすることができます。

  1. In the Google Cloud console, go to the Manage resources page.

     Go to Manage resources

  2. In the project list, select the project that you want to delete, and then click Delete.
  3. In the dialog, type the project ID, and then click Shut down to delete the project.

  個々のリソースの削除

  次のセクションで説明するように、プロジェクトを保持して個々のリソースを削除します。次のコマンドを実行し、プロンプトに従います。

  gcloud container clusters delete ${CLUSTER_NAME} --region=${REGION}
  

  次のステップ

  • 他の TPU トポロジでチュートリアルを構成します。その他の TPU トポロジの詳細については、TPU 構成の計画をご覧ください。
  • このチュートリアルでクローンを作成したサンプル リポジトリで、MaxDiffusion 推論サーバーのサンプルコードを確認します。
  • GKE の TPU の詳細を確認する。
  • JetStream の GitHub リポジトリを確認する。
  • Vertex AI Model Garden を確認する。