透過 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 是擴散專案的起點，適用於研究和生產。

詳情請參閱 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 上使用 TPU v5e 供應 MaxDiffusion 的架構範例。

這個架構包含下列元件：

  • 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.

    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 API.

    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 API

  • 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 API.

    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 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. 按一下「Grant access」(授予存取權)。

    4. 在「New principals」(新增主體) 欄位中，輸入您的使用者 ID。 這通常是指 Google 帳戶的電子郵件地址。

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

      準備環境

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

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

      1. 在 Google Cloud 控制台中，按一下Cloud Shell 啟用圖示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 值。

      使用網頁與模型互動

      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

      輸出內容是模型生成的圖片，類似於下列範例：

      SDXL 生成的圖片

      清除所用資源

      為避免因為本教學課程所用資源，導致系統向 Google Cloud 收取費用，請刪除含有相關資源的專案，或者保留專案但刪除個別資源。

      刪除專案

      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}

      後續步驟