Stable Diffusion XL (SDXL) mit TPUs in GKE mit MaxDiffusion bereitstellen

Autopilot Standard

---

In dieser Anleitung erfahren Sie, wie Sie ein SDXL Bildgenerierungsmodell mit Tensor Processing Units (TPUs) in Google Kubernetes Engine (GKE) mitMaxDiffusion bereitstellen. In dieser Anleitung laden Sie das Modell aus Hugging Face herunter und stellen es in einem Autopilot- oder Standard-Cluster bereit. Dazu verwenden Sie einen Container, in dem MaxDiffusion ausgeführt wird.

Dieser Leitfaden ist ein guter Ausgangspunkt, wenn Sie bei der Bereitstellung und Zugänglichmachung Ihrer KI/ML-Arbeitslasten die detaillierte Kontrolle, Anpassung, Skalierbarkeit, Robustheit, Übertragbarkeit und Kosteneffizienz von verwaltetem Kubernetes benötigen. Wenn Sie eine einheitliche verwaltete KI-Plattform benötigen, um ML-Modelle schnell und kostengünstig zu erstellen und bereitzustellen, empfehlen wir Ihnen, unsere Bereitstellungslösung Vertex AI zu testen.

Hintergrund

Wenn Sie SDXL mit TPUs in GKE mit MaxDiffusion bereitstellen, können Sie eine robuste, produktionsreife Bereitstellungslösung mit allen Vorteilen von verwaltetem Kubernetes erstellen, einschließlich Kosteneffizienz. Skalierbarkeit und höhere Verfügbarkeit. In diesem Abschnitt werden die in dieser Anleitung verwendeten Schlüsseltechnologien beschrieben.

Stabile Diffusion XL (SDXL)

Stable Diffusion XL (SDXL) ist ein latentes Diffusionsmodell (LDM), das von MaxDiffusion für die Inferenz unterstützt wird. Für generative KI können Sie LDMs verwenden, um hochwertige Bilder aus Textbeschreibungen zu generieren. LDMs sind für Anwendungen wie die Bildersuche und die Bilduntertitelung nützlich.

SDXL unterstützt Inferenz mit einzelnen oder mehreren Hosts mit Fragmentierungsannotationen. So kann SDXL trainiert und auf mehreren Rechnern ausgeführt werden, was die Effizienz verbessern kann.

Weitere Informationen finden Sie im Generative Models by Stability AI-Repository und im SDXL-Artikel.

TPUs

TPUs sind von Google speziell entwickelte anwendungsspezifische integrierte Schaltungen (Application-Specific Integrated Circuits, ASICs), die verwendet werden, um das maschinelle Lernen und die KI-Modelle zu beschleunigen, die mit Frameworks wie folgenden erstellt wurden:TensorFlow, PyTorch undJAX.

Bevor Sie TPUs in GKE verwenden, sollten Sie den folgenden Lernpfad durcharbeiten:

1. Lernen Sie mehr über die aktuelle Verfügbarkeit von TPU-Versionen unter Cloud TPU-Systemarchitektur.
2. TPUs in GKE

In dieser Anleitung wird das Bereitstellen des SDXL-Modells beschrieben. GKE stellt das Modell auf TPUv5e-Knoten mit einem einzelnen Host bereit. Dabei werden TPU-Topologien basierend auf den Modellanforderungen für die Bereitstellung von Eingabeaufforderungen mit niedriger Latenz konfiguriert. In dieser Anleitung verwendet das Modell einen TPU v5e-Chip mit einer 1x1-Topologie.

MaxDiffusion

MaxDiffusion ist eine Sammlung von Referenzimplementierungen, geschrieben in Python und Jax, von verschiedenen latenten Diffusionsmodellen, die auf XLA-Geräten ausgeführt werden, einschließlich TPUs und GPUs. MaxDiffusion ist ein Ausgangspunkt für Diffusion-Projekte für Forschung und Produktion.

Weitere Informationen finden Sie im MaxDiffusion-Repository.

Lernziele

Diese Anleitung richtet sich an Kunden von generativer KI, die JAX verwenden, neue oder bestehende Nutzer von SDXL und alle ML-Entwickler, MLOps-Entwickler (DevOps) oder Plattformadministratoren, die daran interessiert sind, Kubernetes-Container-Orchestrierungsfunktionen zu nutzen, um LLMs bereitzustellen.

Diese Anleitung umfasst die folgenden Schritte:

1. Erstellen Sie einen GKE Autopilot- oder Standardcluster mit der empfohlenen TPU-Topologie anhand der Modelleigenschaften.
2. Erstellen Sie ein SDXL-Inferenz-Container-Image.
3. SDXL-Inferenzserver in GKE bereitstellen
4. Interaktion mit dem Modell über eine Webanwendung bereitstellen

Architektur

In diesem Abschnitt wird die in dieser Anleitung verwendete GKE-Architektur beschrieben. Die Architektur besteht aus einem GKE Autopilot- oder Standardcluster, der TPUs bereitstellt und MaxDiffusion-Komponenten hostet. GKE verwendet diese Komponenten, um die Modelle bereitzustellen.

Das folgende Diagramm zeigt die Komponenten dieser Architektur:

Diese Architektur umfasst die folgenden Komponenten:

• Regionaler GKE-Autopilot- oder Standard-Cluster.
• Ein TPU-Slice-Knotenpool mit einem einzelnen Host, der das SDXL-Modell in der MaxDiffusion-Bereitstellung hostet.
• Die Dienstkomponente mit einem Load-Balancer vom Typ ClusterIP. Dieser Service verteilt eingehenden Traffic an alle MaxDiffusion HTTP-Replikate.
• Der Server WebApp HTTP mit einem externen LoadBalancer-Dienst, der eingehenden Traffic verteilt und Traffic der Modellbereitstellung an den ClusterIP-Dienst weiterleitet.

Hinweise

• 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

• Die Abrechnung für das Google Cloud-Projekt muss aktiviert sein.

• 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

• Die Abrechnung für das Google Cloud-Projekt muss aktiviert sein.

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

     Zu IAM
  2. Wählen Sie das Projekt aus.
  3. Klicken Sie auf Zugriff erlauben.

  4. Geben Sie im Feld Neue Hauptkonten Ihre Nutzer-ID ein. Dies ist in der Regel die E-Mail-Adresse eines Google-Kontos.

  5. Wählen Sie in der Liste Rolle auswählen eine Rolle aus.
  6. Wenn Sie weitere Rollen hinzufügen möchten, klicken Sie auf Weitere Rolle hinzufügen und fügen Sie weitere Rollen hinzu.
  7. Klicken Sie auf Speichern.

• Achten Sie darauf, dass Ihr Kontingent für TPU v5e PodSlice Lite-Chips ausreicht. In dieser Anleitung verwenden Sie On-Demand-Instanzen.

Umgebung vorbereiten

In dieser Anleitung verwenden Sie Cloud Shell zum Verwalten von Ressourcen, die in Google Cloud gehostet werden. Die Software, die Sie für diese Anleitung benötigen, ist in Cloud Shell vorinstalliert, einschließlich kubectl und gcloud CLI.

So richten Sie Ihre Umgebung mit Cloud Shell ein:

1. Starten Sie in der Google Cloud Console eine Cloud Shell-Sitzung. Klicken Sie dazu in der Google Cloud Console auf Cloud Shell aktivieren. Dadurch wird im unteren Bereich der Google Cloud Console eine Sitzung gestartet.

2. Legen Sie die Standardumgebungsvariablen fest:

   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
   

   Ersetzen Sie die folgenden Werte:

   • PROJECT_ID: Ihre Google Cloud-Projekt-ID.
   • CLUSTER_NAME: der Name Ihres GKE-Clusters.
   • REGION_NAME: Die Region, in der sich der GKE-Cluster, der Cloud Storage-Bucket und die TPU-Knoten befinden. Die Region enthält Zonen, in denen TPU v5e-Maschinentypen verfügbar sind, z. B. us-west1, us-west4, us-central1, us-east1, us-east5 oder europe-west4.
   • (Nur Standardcluster) ZONE: Die Zone, in der die TPU-Ressourcen verfügbar sind (z. B. us-west4-a). Bei Autopilot-Clustern müssen Sie nicht die Zone angeben, sondern nur die Region.

3. Klonen Sie das Beispiel-Repository und öffnen Sie das Anleitungsverzeichnis:

   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-Ressourcen erstellen und konfigurieren

Folgen Sie dieser Anleitung, um die erforderlichen Ressourcen zu erstellen.

Hinweis: Für die Nutzung einiger Beschleuniger müssen Sie möglicherweise eine Kapazitätsreservierung erstellen. Informationen zum Reservieren und Verbrauchen reservierter Ressourcen finden Sie unter Reservierte zonale Ressourcen aufnehmen.

GKE-Cluster erstellen

Sie können SDXL auf GPUs in einem TPUs-Cluster im Autopilot- oder Standardmodus bereitstellen. Für eine vollständig verwaltete Kubernetes-Umgebung empfehlen wir die Verwendung eines Autopilot-Clusters. Informationen zum Auswählen des GKE-Betriebsmodus, der für Ihre Arbeitslasten am besten geeignet ist, finden Sie unter GKE-Betriebsmodus auswählen.

Autopilot

1. Führen Sie in Cloud Shell den folgenden Befehl aus:

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

   GKE erstellt einen Autopilot-Cluster mit CPU- und TPU-Knoten, wie von den bereitgestellten Arbeitslasten angefordert.

2. Konfigurieren Sie kubectl für die Kommunikation mit Ihrem Cluster:

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

Standard

1. Einen regionalen GKE Standard-Cluster erstellen, der Workload Identity-Föderation für GKE verwendet.

   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}
   

   Die Erstellung eines Clusters kann einige Minuten dauern.

2. Führen Sie den folgenden Befehl aus, um einen Knotenpool für Ihren Cluster zu erstellen:

   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 erstellt einen TPU v5e-Knotenpool mit einer 1x1-Topologie und einem Knoten.

   Informationen zum Erstellen von Knotenpools mit unterschiedlichen Topologien finden Sie unter TPU-Konfiguration planen. Achten Sie darauf, dass Sie die Beispielwerte in dieser Anleitung aktualisieren, z. B. cloud.google.com/gke-tpu-topology und google.com/tpu.

3. Konfigurieren Sie kubectl für die Kommunikation mit Ihrem Cluster:

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

SDXL-Inferenzcontainer erstellen

Folgen Sie dieser Anleitung, um ein Container-Image für den SDXL-Inferenzserver zu erstellen.

1. Öffnen Sie das Manifest 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. Führen Sie den Build aus und erstellen Sie das Inferenz-Container-Image.

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

   Die Ausgabe enthält den Pfad des Container-Images.

SDXL-Inferenzserver bereitstellen

1. Sehen Sie sich das serve_sdxl_v5e.yaml-Manifest an.

   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. Aktualisieren Sie die Projekt-ID im Manifest.

   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
   

3. Wenden Sie das Manifest an:

   kubectl apply -f serve_sdxl_v5e.yaml
   

   Die Ausgabe sieht in etwa so aus:

   deployment.apps/max-diffusion-server created
   

4. Prüfen Sie den Status des Modells:

   watch kubectl get deploy
   

   Die Ausgabe sieht in etwa so aus:

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

5. Rufen Sie die ClusterIP-Adresse ab:

   kubectl get service max-diffusion-server
   

   Die Ausgabe enthält das Feld ClusterIP. Notieren Sie sich den Wert von CLUSTER-IP.

6. Deployment validieren:

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

   Ersetzen Sie CLUSTER_IP durch den Wert CLUSTER-IP, den Sie zuvor notiert haben. Die Ausgabe sieht in etwa so aus:

   {"message":"Hello world! From FastAPI running on Uvicorn with Gunicorn."}
   pod "curl" deleted
   

7. So rufen Sie die Logs des Deployments auf:

   kubectl logs -l app=max-diffusion-server
   

   Wenn das Deployment abgeschlossen ist, sieht die Ausgabe in etwa so aus:

   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.
   

webapp-Client bereitstellen

In diesem Abschnitt stellen Sie den Webapp-Client für das SDXL-Modell bereit.

1. Sehen Sie sich das build/webapp/cloudbuild.yaml-Manifest an.

   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. Führen Sie den Build aus und erstellen Sie das Client-Container-Image im Verzeichnis build/webapp.

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

   Die Ausgabe enthält den Pfad des Container-Images.

3. Öffnen Sie das Manifest 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. Bearbeiten Sie die Projekt-ID im Manifest:

   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
   

5. Wenden Sie das Manifest an:

   kubectl apply -f serve_sdxl_client.yaml
   

6. Rufen Sie die IP-Adresse LoadBalancer ab:

   kubectl get service max-diffusion-client-service
   

   Die Ausgabe enthält das Feld LoadBalancer. Notieren Sie sich den externen IP-Wert.

Erfolg: Sie haben das SDXL-Modell mit MaxDiffusion und TPU v5e in GKE erfolgreich bereitgestellt.

Über die Webseite mit dem Modell interagieren

1. Zugriff auf die folgende URL aus einem Webbrowser:

   http://EXTERNAL_IP:8080
   

   Ersetzen Sie dabei EXTERNAL_IP durch den zuvor notierten EXTERNAL_IP-Wert.

2. Interagieren Sie über die Chatoberfläche mit SDXL. Fügen Sie einen Prompt hinzu und klicken Sie auf Senden. Beispiel:

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

Die Ausgabe ist ein vom Modell generiertes Bild ähnlich dem folgenden Beispiel:

Bereinigen

Damit Ihrem Google Cloud-Konto die in dieser Anleitung verwendeten Ressourcen nicht in Rechnung gestellt werden, löschen Sie entweder das Projekt, das die Ressourcen enthält, oder Sie behalten das Projekt und löschen die einzelnen Ressourcen.

Projekt löschen

Achtung: Das Löschen von Projekten hat folgende Auswirkungen:

• Alle Inhalte des Projekts werden gelöscht. Wenn Sie für die Aufgaben in diesem Dokument ein bereits bestehendes Projekt verwendet haben und dieses löschen, werden auch alle anderen im Rahmen des Projekts erstellten Daten gelöscht.
• Benutzerdefinierte Projekt-IDs gehen verloren. Beim Erstellen dieses Projekts haben Sie möglicherweise eine benutzerdefinierte Projekt-ID erstellt, die Sie weiterhin verwenden möchten. Damit die URLs, die die Projekt-ID nutzen, z. B. eine appspot.com-URL, erhalten bleiben, sollten Sie ausgewählte Ressourcen innerhalb des Projekts löschen, statt das gesamte Projekt.

Wenn Sie mehrere Architekturen, Anleitungen und Kurzanleitungen durcharbeiten möchten, können Sie die Überschreitung von Projektkontingenten verhindern, indem Sie Projekte wiederverwenden.

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.

Einzelne Ressourcen löschen

Behalten Sie das Projekt bei und löschen Sie die einzelnen Ressourcen, wie im folgenden Abschnitt beschrieben. Führen Sie folgende Befehle aus und folgen Sie den Eingabeaufforderungen:

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

Nächste Schritte

• Konfigurieren Sie die Anleitung mit anderen TPU-Topologien. Weitere Informationen zu anderen TPU-Topologien finden Sie unter TPU-Konfiguration planen.
• Sehen Sie sich den Beispielcode des MaxDiffusion-Inferenzservers im Beispiel-Repository an, das Sie in dieser Anleitung geklont haben.
• Mehr über TPUs in GKE erfahren
• GitHub-Repository von JetStream ansehen
• Erfahren Sie mehr über den Vertex AI Model Garden.