Google uses AI technology to translate content into your preferred language. AI translations can contain errors.

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 を使用する前に、次の学習プログラムを完了することをおすすめします。

Cloud TPU システムアーキテクチャで、現在の TPU バージョンの可用性について学習する。
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 コンテナのオーケストレーション機能を使用することに関心をお持ちのプラットフォーム管理者を対象としています。

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

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

アーキテクチャ

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

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

GKE 上の TPU v5e で MaxDiffusion を提供するアーキテクチャの例。

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

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

始める前に

Google Cloud アカウントにログインします。 Google Cloudを初めて使用する場合は、アカウントを作成して、実際のシナリオでの Google プロダクトのパフォーマンスを評価してください。新規のお客様には、ワークロードの実行、テスト、デプロイができる無料クレジット $300 分を差し上げます。

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required APIs.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the APIs

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required APIs.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the APIs

プロジェクトに次のロール（複数の場合あり）が割り当てられていることを確認します。 roles/container.admin、roles/iam.serviceAccountAdmin、roles/artifactregistry.admin、roles/cloudbuild.builds.editor
ロールを確認する
1. Google Cloud コンソールで、[IAM] ページに移動します。
  IAM に移動
2. プロジェクトを選択します。
3. [プリンシパル] 列で、自分または自分が所属するグループの行をすべて確認します。所属するグループについては、管理者にお問い合わせください。
4. 自分のメールアドレスを含む行の [ロール] 列で、ロールのリストに必要なロールが含まれているかどうか確認します。
ロールを付与する
1. Google Cloud コンソールで、[IAM] ページに移動します。
  IAM に移動
2. プロジェクトを選択します。
3. [ アクセスを許可] をクリックします。
4. [新しいプリンシパル] フィールドに、ユーザー ID を入力します。これは通常、Google アカウントのメールアドレスです。
5. [ロールを選択] をクリックし、ロールを検索します。
6. 追加のロールを付与するには、 [別のロールを追加] をクリックして各ロールを追加します。
7. [保存] をクリックします。

TPU v5e PodSlice Lite チップに十分な割り当てがあることを確認します。このチュートリアルでは、オンデマンドインスタンスを使用します。

環境を準備する

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

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

Google Cloud コンソールで（Cloud Shell をアクティブにする）をクリックして、Google Cloud コンソールで Cloud Shell セッションを起動します。これにより、 Google Cloud コンソールの下部ペインでセッションが起動します。
デフォルトの環境変数を設定します。
```
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 CLUSTER_VERSION=CLUSTER_VERSION
export REGION=REGION_NAME
export ZONE=ZONE
```
次の値を置き換えます。
- PROJECT_ID: 実際の Google Cloud プロジェクト ID。
- CLUSTER_NAME: GKE クラスタの名前。
- CLUSTER_VERSION : GKE のバージョン。TPU Trillium（v6e）をサポートする GKE バージョンを指定する必要があります。詳細については、GKE での TPU の可用性を検証するをご覧ください。
- 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 クラスタの場合は、リージョンのみを指定します。ゾーンを指定する必要はありません。

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

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

Cloud Shell で、次のコマンドを実行します。
```
gcloud container clusters create-auto ${CLUSTER_NAME} \
  --project=${PROJECT_ID} \
  --location=${REGION} \
  --release-channel=rapid \
  --cluster-version=${CLUSTER_VERSION}
```
GKE は、デプロイされたワークロードからのリクエストに応じた CPU ノードと TPU ノードを持つ Autopilot クラスタを作成します。

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

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

標準

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}

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

次のコマンドを実行して、クラスタのノードプールを作成します。
```
gcloud container node-pools create maxdiffusion-tpu-nodepool \
  --cluster=${CLUSTER_NAME} \
  --machine-type=ct5lp-hightpu-1t \
  --num-nodes=1 \
  --location=${REGION} \
  --node-locations=${ZONE} \
  --spot
```
GKE は、1x1 トポロジと 1 つのノードを持つ TPU v5e ノードプールを作成します。

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

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

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

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

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

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'

ビルドを実行して、推論コンテナイメージを作成します。
```
cd $WORK_DIR/build/server
gcloud builds submit . --region=$REGION
```
出力には、コンテナイメージのパスが含まれます。

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

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

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

マニフェスト内のプロジェクト 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

次のようにマニフェストを適用します。
```
kubectl apply -f serve_sdxl_v5e.yaml
```
出力は次のようになります。
```
deployment.apps/max-diffusion-server created
```

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

kubectl get deploy --watch

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

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

ClusterIP アドレスを取得します。
```
kubectl get service max-diffusion-server
```
出力には ClusterIP フィールドが含まれます。CLUSTER-IP の値をメモします。

デプロイを検証します。

 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

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 クライアントをデプロイします。

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'

ビルドを実行し、build/webapp ディレクトリにクライアントコンテナイメージを作成します。
```
cd $WORK_DIR/build/webapp
gcloud builds submit . --region=$REGION
```
出力には、コンテナイメージのパスが含まれます。

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

マニフェストでプロジェクト 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

次のようにマニフェストを適用します。
```
kubectl apply -f serve_sdxl_client.yaml
```
LoadBalancer IP アドレスを取得します。
```
kubectl get service max-diffusion-client-service
```
出力には LoadBalancer フィールドが含まれます。EXTERNAL-IP の値をメモします。

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

ウェブブラウザから次の URL にアクセスします。
```
http://EXTERNAL_IP:8080
```
EXTERNAL_IP は、以前にメモした EXTERNAL_IP の値に置き換えます。
チャットインターフェースを使用して SDXL を操作します。プロンプトを追加して [送信] をクリックします。例:
```
Create a detailed image of a fictional historical site, capturing its unique architecture and cultural significance
```

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

SDXL で生成された画像

クリーンアップ

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

プロジェクトを削除する

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

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

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

Google Cloud コンソールで [リソースの管理] ページに移動します。
[リソースの管理] に移動
プロジェクトリストで、削除するプロジェクトを選択し、[削除] をクリックします。
ダイアログでプロジェクト ID を入力し、[シャットダウン] をクリックしてプロジェクトを削除します。

個々のリソースの削除

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

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

次のステップ

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

MaxDiffusion を備えた GKE で TPU を使用して Stable Diffusion XL（SDXL）を提供する コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。

背景

Stable Diffusion XL（SDXL）

TPU

MaxDiffusion

目標

アーキテクチャ

始める前に

ロールを確認する

ロールを付与する

環境を準備する

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

GKE クラスタを作成する

Autopilot

標準

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

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

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

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

クリーンアップ

プロジェクトを削除する

個々のリソースの削除

次のステップ

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

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