Diffuser Stable Diffusion XL (SDXL) à l'aide de TPU sur GKE avec MaxDiffusion

Ce tutoriel vous montre comment diffuser un modèle de génération d'images SDXL à l'aide des TPU (Tensor Processing Unit) sur Google Kubernetes Engine (GKE) avec MaxDiffusion. Dans ce tutoriel, vous allez télécharger le modèle depuis Hugging Face et le déployer sur un cluster Autopilot ou Standard à l'aide d'un conteneur exécutant MaxDiffusion.

Ce guide est un bon point de départ si vous avez besoin du contrôle précis, de la personnalisation, de l'évolutivité, de la résilience, de la portabilité et de la rentabilité des services Kubernetes gérés lors du déploiement et de la diffusion de vos charges de travail d'IA/de ML. Si vous avez besoin d'une plate-forme d'IA gérée unifiée pour créer et diffuser rapidement des modèles de ML à moindre coût, nous vous recommandons d'essayer notre solution de déploiement Vertex AI.

Contexte

En diffusant SDXL à l'aide de TPU sur GKE avec MaxDiffusion, vous pouvez créer une solution de diffusion robuste et prête pour la production avec tous les avantages de la plate-forme Kubernetes gérée, y compris la rentabilité, l'évolutivité et la plus grande disponibilité. Cette section décrit les principales technologies utilisées dans ce tutoriel.

Stable Diffusion XL (SDXL)

Stable Diffusion XL (SDXL) est un type de modèle de diffusion latente (LDM) compatible avec MaxDiffusion pour l'inférence. Pour l'IA générative, vous pouvez utiliser des LDM pour générer des images de haute qualité à partir de descriptions textuelles. Les LDM sont utiles pour des applications telles que la recherche d'images et le sous-titrage d'images.

SDXL est compatible avec l'inférence à hôte unique ou multi-hôtes avec des annotations de partitionnement. Cela permet à SDXL d'être entraîné et exécuté sur plusieurs machines, ce qui peut améliorer l'efficacité.

Pour en savoir plus, consultez le dépôt de modèles génératifs de Stability AI et le SDXL.

TPU

Les TPU sont des circuits intégrés propres aux applications (Application-Specific Integrated Circuit ou ASIC), développés spécifiquement par Google et permettant d'accélérer le machine learning et les modèles d'IA créés à l'aide de frameworks tels que TensorFlow, PyTorch et JAX.

Avant d'utiliser des TPU dans GKE, nous vous recommandons de suivre le parcours de formation suivant :

  1. Découvrez la disponibilité actuelle des versions de TPU avec l'architecture système de Cloud TPU.
  2. Apprenez-en plus sur les TPU dans GKE.

Ce tutoriel explique comment diffuser le modèle SDXL. GKE déploie le modèle sur des nœuds TPU v5e à hôte unique avec des topologies TPU configurées en fonction des exigences du modèle pour diffuser des requêtes avec une faible latence. Dans ce guide, le modèle utilise une puce TPU v5e avec une topologie 1x1.

MaxDiffusion

MaxDiffusion est une collection d'implémentations de référence, écrites en Python et Jax, de divers modèles de diffusion latente qui s'exécutent sur des appareils XLA, y compris les TPU et les GPU. MaxDiffusion est un point de départ pour les projets de diffusion, à la fois pour la recherche et la production.

Pour en savoir plus, consultez le dépôt MaxDiffusion.

Objectifs

Ce tutoriel est destiné aux clients d'IA générative qui utilisent JAX, aux utilisateurs nouveaux ou existants de SDXL, aux ingénieurs en ML, aux ingénieurs MLOps (DevOps) ou aux administrateurs de plate-forme qui souhaitent utiliser les fonctionnalités d'orchestration de conteneurs Kubernetes pour diffuser des LLM.

Ce tutoriel couvre les étapes suivantes :

  1. Créer un cluster GKE Autopilot ou standard avec la topologie TPU recommandée en fonction des caractéristiques du modèle.
  2. Créez une image de conteneur d'inférence SDXL.
  3. Déployer le serveur d'inférence SDXL sur GKE.
  4. Diffuser le modèle et interagir avec lui via une application Web.

Architecture

Cette section décrit l'architecture GKE utilisée dans ce tutoriel. L'architecture se compose d'un cluster GKE Autopilot ou GKE Standard qui provisionne des TPU et héberge des composants MaxDiffusion. GKE utilise ces composants pour déployer et diffuser les modèles.

Le schéma suivant montre les composants de cette architecture:

Exemple d'architecture de diffusion de MaxDiffusion avec des TPU v5e sur GKE.

Cette architecture comprend les composants suivants :

  • Un cluster régional GKE Autopilot ou Standard.
  • Un pool de nœuds de tranche de TPU à hôte unique qui héberge le modèle SDXL sur le déploiement MaxDiffusion.
  • Composant Service avec un équilibreur de charge de type ClusterIP. Ce service distribue le trafic entrant à toutes les répliques MaxDiffusion HTTP.
  • Le serveur WebApp HTTP avec un service LoadBalancer externe qui distribue le trafic entrant et redirige le trafic de diffusion du modèle vers le service ClusterIP.

Avant de commencer

  • Connectez-vous à votre compte Google Cloud . Si vous débutez sur Google Cloud, créez un compte pour évaluer les performances de nos produits en conditions réelles. Les nouveaux clients bénéficient également de 300 $de crédits sans frais pour exécuter, tester et déployer des charges de travail.

  • 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

    •

  • Assurez-vous de disposer des rôles suivants sur le projet : roles/container.admin, roles/iam.serviceAccountAdmin, roles/artifactregistry.admin, roles/cloudbuild.builds.editor

    Vérifier les rôles

    1. Dans la console Google Cloud , accédez à la page IAM.

      Accéder à IAM
    2. Sélectionnez le projet.

    3. Dans la colonne Compte principal, recherchez toutes les lignes qui vous identifient ou identifient un groupe dont vous faites partie. Pour savoir à quels groupes vous appartenez, contactez votre administrateur.

    4. Pour toutes les lignes qui vous spécifient ou vous incluent, consultez la colonne Rôle pour vous assurer que la liste inclut les rôles requis.

    Attribuer les rôles

    1. Dans la console Google Cloud , accédez à la page IAM.

      Accéder à IAM
    2. Sélectionnez le projet.
    3. Cliquez sur  Accorder l'accès.

    4. Dans le champ Nouveaux comptes principaux, saisissez votre identifiant utilisateur. Il s'agit généralement de l'adresse e-mail d'un compte Google.

    5. Cliquez sur Sélectionner un rôle, puis recherchez le rôle.
    6. Pour attribuer des rôles supplémentaires, cliquez sur  Ajouter un autre rôle et ajoutez tous les rôles supplémentaires.
    7. Cliquez sur Enregistrer.
  • Vérifiez que vous disposez d'un quota suffisant pour les puces TPU v5e PodSlice Lite. Dans ce tutoriel, vous utilisez des instances à la demande.

Préparer l'environnement

Dans ce tutoriel, vous utilisez Cloud Shell pour gérer les ressources hébergées surGoogle Cloud. Cloud Shell est préinstallé avec les logiciels dont vous avez besoin pour ce tutoriel, y compris kubectl et gcloud CLI.

Pour configurer votre environnement avec Cloud Shell, procédez comme suit :

  1. Dans la console Google Cloud , lancez une session Cloud Shell en cliquant sur Icône d'activation Cloud Shell Activer Cloud Shell dans la consoleGoogle Cloud . Une session s'ouvre dans le volet inférieur de la console Google Cloud .

  2. Définissez les variables d'environnement par défaut :

    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

    Remplacez les valeurs suivantes :

    • PROJECT_ID : ID de votre projet Google Cloud .
    • CLUSTER_NAME : nom de votre cluster GKE.
    • CLUSTER_VERSION : version de GKE. Vous devez spécifier une version de GKE compatible avec les TPU Trillium (v6e). Pour en savoir plus, consultez Valider la disponibilité des TPU dans GKE.
    • REGION_NAME: région où se trouvent le cluster GKE, le bucket Cloud Storage et les nœuds TPU. La région contient des zones où les types de machines TPU v5e sont disponibles (par exemple, us-west1, us-west4, us-central1, us-east1, us-east5 ou europe-west4).
    • (Cluster standard uniquement) ZONE : zone où les ressources TPU sont disponibles (par exemple, us-west4-a). Pour les clusters Autopilot, il n'est pas nécessaire de spécifier la zone, mais uniquement la région.

  3. Clonez l'exemple de dépôt et ouvrez le répertoire du tutoriel :

    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

Créer et configurer des ressources Google Cloud

Suivez les instructions ci-dessous pour créer les ressources requises.

Créer un cluster GKE

Vous pouvez diffuser les modèles SDLX sur des TPU dans un cluster GKE Autopilot ou GKE Standard. Nous vous recommandons d'utiliser un cluster GKE Autopilot pour une expérience Kubernetes entièrement gérée. Pour choisir le mode de fonctionnement GKE le mieux adapté à vos charges de travail, consultez la section Choisir un mode de fonctionnement GKE.

Autopilot

  1. Dans Cloud Shell, exécutez la commande suivante :

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

    GKE crée un cluster Autopilot avec des nœuds de processeur et de GPU, à la demande des charges de travail déployées.

  2. Configurez kubectl de manière à communiquer avec votre cluster :

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

Standard

  1. Créez un cluster GKE Standard régional qui utilise la fédération d'identité de charge de travail pour 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}

    La création du cluster peut prendre plusieurs minutes.

  2. Exécutez la commande suivante pour créer un pool de nœuds pour votre cluster :

    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 crée un pool de nœuds TPU v5e avec une topologie 1x1 et un nœud.

    Pour créer des pools de nœuds avec différentes topologies, découvrez comment planifier votre configuration de TPU. Assurez-vous de mettre à jour les exemples de valeurs de ce tutoriel, comme cloud.google.com/gke-tpu-topology et google.com/tpu.

  3. Configurez kubectl de manière à communiquer avec votre cluster :

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

Créer le conteneur d'inférence SDXL

Suivez ces instructions pour créer une image de conteneur pour le serveur d'inférence SDXL.

  1. Ouvrez le fichier manifeste 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. Exécutez la compilation et créez une image de conteneur d'inférence.

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

    La sortie contient le chemin d'accès à l'image du conteneur.

Déployer le serveur d'inférence SDXL

Dans cette section, vous allez déployer le serveur d'inférence SDXL. Pour déployer le serveur, ce tutoriel utilise un déploiement Kubernetes. Un déploiement est un objet de l'API Kubernetes qui vous permet d'exécuter plusieurs instances dupliquées de pods répartis entre les nœuds d'un cluster.

  1. Explorez le fichier manifeste 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. Mettez à jour l'ID du projet dans le fichier manifeste.

    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. Appliquez le fichier manifeste :

    kubectl apply -f serve_sdxl_v5e.yaml

    Le résultat ressemble à ce qui suit :

    deployment.apps/max-diffusion-server created

  4. Vérifiez l'état du modèle :

    kubectl get deploy --watch

    Le résultat ressemble à ce qui suit :

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

  5. Récupérez l'adresse ClusterIP :

    kubectl get service max-diffusion-server

    Le résultat contient un champ ClusterIP. Notez la valeur de CLUSTER-IP.

  6. Validez le déploiement.

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

    Remplacez CLUSTER_IP par la valeur CLUSTER-IP que vous avez notée précédemment. Le résultat ressemble à ce qui suit :

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

  7. Affichez les journaux du déploiement :

    kubectl logs -l app=max-diffusion-server

    Une fois le déploiement terminé, le résultat ressemble à ce qui suit :

    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.

Déployer le client de l'application Web

Dans cette section, vous allez déployer le client de l'application Web pour diffuser le modèle SDXL.

  1. Explorez le fichier manifeste 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. Exécutez la compilation et créez l'image de conteneur client dans le répertoire build/webapp.

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

    La sortie contient le chemin d'accès à l'image du conteneur.

  3. Ouvrez le fichier manifeste 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. Modifiez l'ID du projet dans le fichier manifeste :

    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. Appliquez le fichier manifeste :

    kubectl apply -f serve_sdxl_client.yaml

  6. Récupérez l'adresse IP de LoadBalancer :

    kubectl get service max-diffusion-client-service

    Le résultat contient un champ LoadBalancer. Notez la valeur EXTERNAL-IP.

Interagir avec le modèle à l'aide de la page Web

  1. Accès à l'URL suivante depuis un navigateur Web :

    http://EXTERNAL_IP:8080

    Remplacez EXTERNAL_IP par la valeur EXTERNAL_IP que vous avez notée précédemment.

  2. Interagissez avec SDXL à l'aide de l'interface de chat. Ajoutez une requête et cliquez sur Envoyer. Exemple :

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

Le résultat est une image générée par le modèle, semblable à l'exemple suivant :

Image générée par SDXL

Effectuer un nettoyage

Pour éviter que les ressources utilisées lors de ce tutoriel soient facturées sur votre compte Google Cloud, supprimez le projet contenant les ressources, ou conservez le projet et supprimez les ressources individuelles.

Supprimer le projet

  1. Dans la console Google Cloud , accédez à la page Gérer les ressources.

    Accéder à la page "Gérer les ressources"

  2. Dans la liste des projets, sélectionnez le projet que vous souhaitez supprimer, puis cliquez sur Supprimer.
  3. Dans la boîte de dialogue, saisissez l'ID du projet, puis cliquez sur Arrêter pour supprimer le projet.

Supprimer les ressources individuelles

Conservez le projet et supprimez les ressources individuelles, comme décrit dans la section suivante. Exécutez les commandes suivantes et suivez les instructions :

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

Étape suivante