Menyajikan Stable Diffusion XL (SDXL) menggunakan TPU di GKE dengan MaxDiffusion

Tutorial ini menunjukkan cara menayangkan model pembuatan gambar SDXL menggunakan Unit Pemrosesan Tensor (TPU) di Google Kubernetes Engine (GKE) dengan MaxDiffusion. Dalam tutorial ini, Anda akan mendownload model dari Hugging Face dan men-deploy-nya di cluster Autopilot atau Standard menggunakan container yang menjalankan MaxDiffusion.

Panduan ini adalah titik awal yang baik jika Anda memerlukan kontrol terperinci, penyesuaian, skalabilitas, ketahanan, portabilitas, dan efektivitas biaya Kubernetes terkelola saat men-deploy dan menyajikan workload AI/ML Anda. Jika Anda memerlukan platform AI terkelola terpadu untuk membangun dan menayangkan model ML dengan cepat dan efektif biaya, sebaiknya coba solusi deployment Vertex AI kami.

Latar belakang

Dengan menayangkan SDXL menggunakan TPU di GKE dengan MaxDiffusion, Anda dapat membangun solusi penayangan yang tangguh dan siap produksi dengan semua manfaat Kubernetes terkelola, termasuk efisiensi biaya, skalabilitas, dan ketersediaan yang lebih tinggi. Bagian ini menjelaskan teknologi utama yang digunakan dalam tutorial ini.

Stable Diffusion XL (SDXL)

Stable Diffusion XL (SDXL) adalah jenis model difusi laten (LDM) yang didukung oleh MaxDiffusion untuk inferensi. Untuk AI generatif, Anda dapat menggunakan LDM untuk menghasilkan gambar berkualitas tinggi dari deskripsi teks. LDM berguna untuk aplikasi seperti penelusuran gambar dan pemberian teks gambar.

SDXL mendukung inferensi host tunggal atau multi-host dengan anotasi sharding. Hal ini memungkinkan SDXL dilatih dan dijalankan di beberapa mesin, yang dapat meningkatkan efisiensi.

Untuk mempelajari lebih lanjut, lihat repositori Model Generatif oleh Stability AI dan artikel SDXL.

TPU

TPU adalah sirkuit terintegrasi khusus aplikasi (ASIC) yang dikembangkan khusus oleh Google dan digunakan untuk mempercepat model machine learning dan AI yang dibangun menggunakan framework seperti TensorFlow, PyTorch, dan JAX.

Sebelum menggunakan TPU di GKE, sebaiknya selesaikan jalur pembelajaran berikut:

  1. Pelajari ketersediaan versi TPU saat ini dengan arsitektur sistem Cloud TPU.
  2. Pelajari TPU di GKE.

Tutorial ini membahas cara menyajikan model SDXL. GKE men-deploy model pada node TPU v5e host tunggal dengan topologi TPU yang dikonfigurasi berdasarkan persyaratan model untuk menyajikan perintah dengan latensi rendah. Dalam panduan ini, model menggunakan chip TPU v5e dengan topologi 1x1.

MaxDiffusion

MaxDiffusion adalah kumpulan implementasi referensi, yang ditulis dalam Python dan Jax, dari berbagai model difusi laten yang berjalan di perangkat XLA, termasuk TPU dan GPU. MaxDiffusion adalah titik awal untuk project Difusi untuk riset dan produksi.

Untuk mempelajari lebih lanjut, lihat repositori MaxDiffusion.

Tujuan

Tutorial ini ditujukan bagi pelanggan AI generatif yang menggunakan JAX, pengguna SDXL baru atau lama, dan Engineer ML, engineer MLOps (DevOps), atau administrator platform yang tertarik menggunakan kemampuan orkestrasi container Kubernetes untuk menayangkan LLM.

Tutorial ini membahas langkah-langkah berikut:

  1. Buat cluster GKE Autopilot atau Standard dengan topologi TPU yang direkomendasikan, berdasarkan karakteristik model.
  2. Buat image container inferensi SDXL.
  3. Deploy server inferensi SDXL di GKE.
  4. Menyajikan dan berinteraksi dengan model melalui aplikasi web.

Arsitektur

Bagian ini menjelaskan arsitektur GKE yang digunakan dalam tutorial ini. Arsitektur ini terdiri dari cluster GKE Autopilot atau Standard yang menyediakan TPU dan menghosting komponen MaxDiffusion. GKE menggunakan komponen ini untuk men-deploy dan menyajikan model.

Diagram berikut menunjukkan komponen arsitektur ini:

Contoh arsitektur untuk menayangkan MaxDiffusion dengan TPU v5e di GKE.

Arsitektur ini mencakup komponen berikut:

  • Cluster regional GKE Autopilot atau Standard.
  • Satu node pool slice TPU host tunggal yang menghosting model SDXL pada deployment MaxDiffusion.
  • Komponen Service dengan load balancer jenis ClusterIP. Layanan ini mendistribusikan traffic masuk ke semua replika MaxDiffusion HTTP.
  • Server WebApp HTTP dengan Layanan LoadBalancer eksternal yang mendistribusikan traffic masuk dan mengalihkan traffic penayangan model ke Layanan ClusterIP.

Sebelum memulai

  • Login ke akun Google Cloud Anda. Jika Anda baru menggunakan Google Cloud, buat akun untuk mengevaluasi performa produk kami dalam skenario dunia nyata. Pelanggan baru juga mendapatkan kredit gratis senilai $300 untuk menjalankan, menguji, dan men-deploy workload.
  • 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

  • Pastikan Anda memiliki peran berikut di project: roles/container.admin, roles/iam.serviceAccountAdmin, roles/artifactregistry.admin, roles/cloudbuild.builds.editor

    Memeriksa peran

    1. Di konsol Google Cloud , buka halaman IAM.

      Buka IAM
    2. Pilih project.
    3. Di kolom Akun utama, temukan semua baris yang mengidentifikasi Anda atau grup yang Anda ikuti. Untuk mengetahui grup mana saja yang Anda ikuti, hubungi administrator Anda.

    4. Untuk semua baris yang menentukan atau menyertakan Anda, periksa kolom Peran untuk melihat apakah daftar peran menyertakan peran yang diperlukan.

    Memberikan peran

    1. Di konsol Google Cloud , buka halaman IAM.

      Buka IAM
    2. Pilih project.
    3. Klik Grant access.
    4. Di kolom New principals, masukkan ID pengguna Anda. Biasanya, ini adalah alamat email untuk Akun Google.

    5. Klik Pilih peran, lalu telusuri peran.
    6. Untuk memberikan peran tambahan, klik Add another role, lalu tambahkan tiap peran tambahan.
    7. Klik Simpan.
  • Pastikan Anda memiliki kuota yang cukup untuk chip TPU v5e PodSlice Lite. Dalam tutorial ini, Anda akan menggunakan instance sesuai permintaan.

Menyiapkan lingkungan

Dalam tutorial ini, Anda akan menggunakan Cloud Shell untuk mengelola resource yang dihosting di Google Cloud. Cloud Shell telah diinstal dengan software yang akan Anda perlukan untuk tutorial ini, termasuk kubectl dan gcloud CLI.

Untuk menyiapkan lingkungan Anda dengan Cloud Shell, ikuti langkah-langkah berikut:

  1. Di konsol Google Cloud , luncurkan sesi Cloud Shell dengan mengklik Ikon aktivasi Cloud Shell Aktifkan Cloud Shell di konsolGoogle Cloud . Tindakan ini akan meluncurkan sesi di panel bawah konsol Google Cloud .

  2. Tetapkan variabel lingkungan default:

    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
    

    Ganti nilai berikut:

    • PROJECT_ID: Google Cloud Project ID Anda.
    • CLUSTER_NAME: nama cluster GKE Anda.
    • CLUSTER_VERSION : versi GKE. Anda harus menentukan versi GKE yang mendukung TPU Trillium (v6e). Untuk mengetahui informasi selengkapnya, lihat Memvalidasi ketersediaan TPU di GKE.
    • REGION_NAME: region tempat cluster GKE, bucket Cloud Storage, dan node TPU Anda berada. Wilayah ini berisi zona tempat jenis mesin TPU v5e tersedia (misalnya, us-west1, us-west4, us-central1, us-east1, us-east5, atau europe-west4).
    • (Khusus cluster Standard) ZONE: zona tempat resource TPU tersedia (misalnya, us-west4-a). Untuk cluster Autopilot, Anda tidak perlu menentukan zona, hanya region.
  3. Clone repositori contoh dan buka direktori tutorial:

    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
    

Membuat dan mengonfigurasi resource Google Cloud

Ikuti petunjuk berikut untuk membuat resource yang diperlukan.

Membuat cluster GKE

Anda dapat menayangkan SDXL di TPU dalam cluster GKE Autopilot atau Standard. Sebaiknya gunakan cluster Autopilot untuk pengalaman Kubernetes yang terkelola sepenuhnya. Untuk memilih mode operasi GKE yang paling sesuai untuk workload Anda, lihat Memilih mode operasi GKE.

Autopilot

  1. Jalankan perintah berikut di Cloud Shell:

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

    GKE membuat cluster Autopilot dengan node CPU dan TPU seperti yang diminta oleh workload yang di-deploy.

  2. Konfigurasi kubectl untuk berkomunikasi dengan cluster Anda:

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

Standar

  1. Buat cluster Standar GKE regional yang menggunakan Workload Identity Federation for GKE.

    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}
    

    Pembuatan cluster mungkin memerlukan waktu beberapa menit.

  2. Jalankan perintah berikut untuk membuat node pool untuk cluster Anda:

    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 membuat node pool TPU v5e dengan topologi 1x1 dan satu node.

    Untuk membuat node pool dengan topologi yang berbeda, pelajari cara Merencanakan konfigurasi TPU. Pastikan Anda memperbarui nilai contoh dalam tutorial ini, seperti cloud.google.com/gke-tpu-topology dan google.com/tpu.

  3. Konfigurasi kubectl untuk berkomunikasi dengan cluster Anda:

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

Membangun container inferensi SDXL

Ikuti petunjuk ini untuk membuat image container bagi server inferensi SDXL.

  1. Buka manifes 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. Jalankan build dan buat image container inferensi.

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

    Output berisi jalur image container.

Men-deploy server inferensi SDXL

Di bagian ini, Anda akan men-deploy server inferensi SDXL. Untuk men-deploy server, tutorial ini menggunakan Deployment Kubernetes. Deployment adalah objek Kubernetes API yang memungkinkan Anda menjalankan beberapa replika Pod yang didistribusikan di antara node dalam cluster.

  1. Pelajari manifes 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. Perbarui project ID di manifes.

    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. Terapkan manifes:

    kubectl apply -f serve_sdxl_v5e.yaml
    

    Outputnya mirip dengan hal berikut ini:

    deployment.apps/max-diffusion-server created
    
  4. Verifikasi status model:

    kubectl get deploy --watch
    

    Outputnya mirip dengan hal berikut ini:

    NAME                          READY   UP-TO-DATE   AVAILABLE   AGE
    stable-diffusion-deployment   1/1     1            1           8m21s
    
  5. Ambil alamat ClusterIP:

    kubectl get service max-diffusion-server
    

    Output berisi kolom ClusterIP. Catat nilai CLUSTER-IP.

  6. Memvalidasi Deployment:

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

    Ganti CLUSTER_IP dengan nilai CLUSTER-IP yang Anda catat sebelumnya. Outputnya mirip dengan hal berikut ini:

    {"message":"Hello world! From FastAPI running on Uvicorn with Gunicorn."}
    pod "curl" deleted
    
  7. Lihat log dari Deployment:

    kubectl logs -l app=max-diffusion-server
    

    Setelah Deployment selesai, output-nya akan mirip dengan berikut ini:

    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.
    

Men-deploy klien aplikasi web

Di bagian ini, Anda akan men-deploy klien aplikasi web untuk menayangkan model SDXL.

  1. Pelajari manifes 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. Jalankan build dan buat image container klien di direktori build/webapp.

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

    Output berisi jalur image container.

  3. Buka manifes 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. Edit project ID dalam manifes:

    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. Terapkan manifes:

    kubectl apply -f serve_sdxl_client.yaml
    
  6. Ambil alamat IP LoadBalancer:

    kubectl get service max-diffusion-client-service
    

    Output berisi kolom LoadBalancer. Catat nilai EXTERNAL-IP.

Berinteraksi dengan model menggunakan halaman web

  1. Akses URL berikut dari browser web:

    http://EXTERNAL_IP:8080
    

    Ganti EXTERNAL_IP dengan nilai EXTERNAL_IP yang Anda catat sebelumnya.

  2. Berinteraksi dengan SDXL menggunakan antarmuka chat. Tambahkan perintah, lalu klik Kirim. Contoh:

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

Outputnya adalah gambar yang dihasilkan model dan mirip dengan contoh berikut:

Gambar yang dihasilkan SDXL

Pembersihan

Agar tidak perlu membayar biaya pada akun Google Cloud Anda untuk resource yang digunakan dalam tutorial ini, hapus project yang berisi resource tersebut, atau simpan project dan hapus setiap resource.

Menghapus project

  1. Di Konsol Google Cloud , buka halaman Manage resources.

    Buka Kelola resource

  2. Pada daftar project, pilih project yang ingin Anda hapus, lalu klik Delete.
  3. Pada dialog, ketik project ID, lalu klik Shut down untuk menghapus project.

Menghapus resource satu per satu

Pertahankan project dan hapus resource individual, seperti yang dijelaskan di bagian berikut. Jalankan perintah berikut dan ikuti perintahnya:

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

Langkah berikutnya