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 サーバー。

始める前に

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. [保存] をクリックします。

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

環境を準備する

このチュートリアルでは、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 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 クラスタの場合は、リージョンのみを指定します。ゾーンを指定する必要はありません。

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

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} \
  --region=${REGION} \
  --release-channel=rapid \
  --cluster-version=1.29
```
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 \
  --region=${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 推論サーバーのコンテナイメージをビルドします。

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 推論サーバーをデプロイする

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
sed -i "s|PROJECT_ID|$PROJECT_ID|g" serve_sdxl_v5e.yaml
sed -i "s|REGION|$REGION|g" serve_sdxl_v5e.yaml

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

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

watch kubectl get deploy

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

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
sed -i "s|PROJECT_ID|$PROJECT_ID|g" serve_sdxl_client.yaml
sed -i "s|ClusterIP|$ClusterIP|g" serve_sdxl_client.yaml
sed -i "s|REGION|$REGION|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）を保持するには、プロジェクト全体ではなくプロジェクト内の選択したリソースだけを削除します。

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

In the Google Cloud console, go to the Manage resources page.
Go to Manage resources
In the project list, select the project that you want to delete, and then click Delete.
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 を確認する。

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

背景

Stable Diffusion XL（SDXL）

TPU

MaxDiffusion

目標

アーキテクチャ

始める前に

Check for the roles

Grant the roles

環境を準備する

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

GKE クラスタを作成する

Autopilot

標準

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

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

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

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

クリーンアップ

プロジェクトの削除

個々のリソースを削除する

次のステップ

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