透過 MaxDiffusion 在 GKE 上使用 TPU 提供 Stable Diffusion XL (SDXL)

自動駕駛 標準

---

本教學課程說明如何使用 MaxDiffusion，在 Google Kubernetes Engine (GKE) 上透過張量處理單元 (TPU) 服務 SDXL 圖像生成模型。在本教學課程中，您會從 Hugging Face 下載模型，並使用執行 MaxDiffusion 的容器，將模型部署至 Autopilot 或標準叢集。

如果您在部署及提供 AI/機器學習工作負載時，需要代管型 Kubernetes 的精細控制、自訂、可擴充性、復原能力、可攜性和成本效益，這份指南就是您的最佳起點。如果您需要統一管理的 AI 平台，以便快速建構及提供機器學習模型，並有效控管成本，建議您試用 Vertex AI 部署解決方案。

背景

透過 MaxDiffusion 在 GKE 上使用 TPU 提供 SDXL，即可建構可用於正式環境的穩固服務解決方案，同時享有代管型 Kubernetes 的所有優點，包括成本效益、擴充性和高可用性。本節說明本教學課程中使用的主要技術。

Stable Diffusion XL (SDXL)

Stable Diffusion XL (SDXL) 是 MaxDiffusion 支援用於推論的潛在擴散模型 (LDM)。生成式 AI 方面，您可以使用 LDM 根據文字描述生成高品質圖片。LDM 適用於圖片搜尋和圖片說明文字等應用程式。

SDXL 支援單一或多主機推論，並提供分片註解。 這樣一來，SDXL 就能在多部機器上訓練及執行，進而提升效率。

詳情請參閱 Stability AI 存放區的生成模型和 SDXL 文件。

TPU

TPU 是 Google 開發的客製化特殊應用積體電路 (ASIC)，用於加速機器學習和 AI 模型，這些模型是使用 TensorFlow、PyTorch 和 JAX 等架構建構而成。

在 GKE 中使用 TPU 之前，建議您先完成下列學習路徑：

1. 如要瞭解目前可用的 TPU 版本，請參閱 Cloud TPU 系統架構。
2. 瞭解 GKE 中的 TPU。

本教學課程涵蓋 SDXL 模型服務。GKE 會在單一主機 TPU v5e 節點上部署模型，並根據模型需求設定 TPU 拓撲，以低延遲提供提示。在本指南中，模型會使用 TPU v5e 晶片和 1x1 拓撲。

MaxDiffusion

MaxDiffusion 是一系列參考實作項目，以 Python 和 Jax 撰寫，適用於各種潛在擴散模型，可在 XLA 裝置 (包括 TPU 和 GPU) 上執行。MaxDiffusion 是研究和生產的 Diffusion 專案起點。

詳情請參閱 MaxDiffusion 存放區。

目標

本教學課程適用於使用 JAX 的生成式 AI 客戶、SDXL 的新使用者或現有使用者，以及對使用 Kubernetes 容器協調功能提供 LLM 服務感興趣的任何機器學習工程師、MLOps (DevOps) 工程師或平台管理員。

本教學課程包含下列步驟：

1. 根據模型特性，建立具有建議 TPU 拓撲的 GKE Autopilot 或 Standard 叢集。
2. 建構 SDXL 推論容器映像檔。
3. 在 GKE 上部署 SDXL 推論伺服器。
4. 透過網頁應用程式提供模型並與模型互動。

架構

本節說明本教學課程使用的 GKE 架構。此架構包含 GKE Autopilot 或標準叢集，可佈建 TPU 並代管 MaxDiffusion 元件。GKE 會使用這些元件部署及提供模型。

下圖顯示這個架構的元件：

這個架構包含下列元件：

• GKE Autopilot 或 Standard 區域叢集。
• 一個單一主機 TPU 配量節點集區，用於在 MaxDiffusion 部署作業中代管 SDXL 模型。
• 服務元件，負載平衡器類型為 ClusterIP。這項服務會將傳入流量分配給所有 MaxDiffusion HTTP 副本。
• 具有外部 LoadBalancer 服務的 WebApp HTTP 伺服器，可分配傳入流量，並將模型服務流量重新導向至 ClusterIP 服務。

事前準備

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

Verify 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

Verify 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. 在「New principals」(新增主體) 欄位中，輸入您的使用者 ID。 這通常是 Google 帳戶的電子郵件地址。

  5. 在「Select a role」(選取角色) 清單中，選取角色。
  6. 如要授予其他角色，請按一下 「新增其他角色」，然後新增每個其他角色。
  7. 按一下 [Save]。
  • 確認您有足夠的 TPU v5e PodSlice Lite 晶片配額。在本教學課程中，您將使用隨選執行個體。

  準備環境

  在本教學課程中，您將使用 Cloud Shell 管理託管於Google Cloud的資源。Cloud Shell 已預先安裝本教學課程所需的軟體，包括 kubectl 和 gcloud CLI。

  如要使用 Cloud Shell 設定環境，請按照下列步驟操作：

  1. 在 Google Cloud 控制台中，按一下Google Cloud 控制台中的「啟用 Cloud Shell」，啟動 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 或 Standard 叢集的 TPU 上提供 SDXL 服務。建議您使用 Autopilot 叢集，享受全代管 Kubernetes 體驗。如要為工作負載選擇最合適的 GKE 作業模式，請參閱「選擇 GKE 作業模式」。

  Autopilot

  1. 在 Cloud Shell 中執行下列指令：

     gcloud container clusters create-auto ${CLUSTER_NAME} \
       --project=${PROJECT_ID} \
       --location=${REGION} \
       --release-channel=rapid \
       --cluster-version=1.29
     

     GKE 會根據部署的工作負載要求，建立含有 CPU 和 TPU 節點的 Autopilot 叢集。

  2. 設定 kubectl 與叢集通訊：

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

  標準

  1. 建立使用 Workload Identity Federation for GKE 的地區 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 \
       --location=${REGION} \
       --node-locations=${ZONE} \
       --spot
     

     GKE 會建立具有 1x1 拓撲和一個節點的 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 是 Kubernetes API 物件，可讓您執行多個 Pod 副本，並將這些副本分散到叢集的節點中。

  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
     

     部署完成後，輸出內容會類似如下：

     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 模型。

  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. 透過網路瀏覽器存取下列網址：

     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 的網址 (例如 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} --location=${REGION}
  

  後續步驟

  • 使用其他 TPU 拓撲設定本教學課程。如要進一步瞭解其他 TPU 拓撲，請參閱「規劃 TPU 設定」。
  • 在本教學課程中複製的範例存放區中，探索 MaxDiffusion 推論伺服器範例程式碼。
  • 進一步瞭解 GKE 中的 TPU。
  • 探索 JetStream GitHub 存放區。
  • 探索 Vertex AI Model Garden。