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


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:

Beispielarchitektur für die Bereitstellung von MaxDiffusion mit TPU v5e in GKE.

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.

    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.

      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 Symbol für die Cloud Shell-Aktivierung 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.

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.

Ü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:

Mit SDXL generiertes Bild

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

  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