Servir Stable Diffusion XL (SDXL) mediante TPUs en GKE con MaxDiffusion

Autopilot Standard

---

En este tutorial se muestra cómo servir un modelo de generación de imágenes SDXL mediante unidades de procesamiento de tensor (TPUs) en Google Kubernetes Engine (GKE) con MaxDiffusion. En este tutorial, descargarás el modelo de Hugging Face y lo desplegarás en un clúster Autopilot o Standard con un contenedor que ejecute MaxDiffusion.

Esta guía es un buen punto de partida si necesitas el control granular, la personalización, la escalabilidad, la resiliencia, la portabilidad y la rentabilidad de Kubernetes gestionado al desplegar y servir tus cargas de trabajo de IA o aprendizaje automático. Si necesitas una plataforma de IA gestionada unificada para crear y ofrecer modelos de aprendizaje automático rápidamente y de forma rentable, te recomendamos que pruebes nuestra solución de despliegue Vertex AI.

Fondo

Si sirves SDXL mediante TPUs en GKE con MaxDiffusion, puedes crear una solución de servicio estable y lista para producción con todas las ventajas de Kubernetes gestionado, como la rentabilidad, la escalabilidad y una mayor disponibilidad. En esta sección se describen las tecnologías clave que se usan en este tutorial.

Stable Diffusion XL (SDXL)

Stable Diffusion XL (SDXL) es un tipo de modelo de difusión latente (LDM) compatible con MaxDiffusion para la inferencia. En el caso de la IA generativa, puedes usar los modelos de difusión latente para generar imágenes de alta calidad a partir de descripciones de texto. Los modelos de difusión latente son útiles para aplicaciones como la búsqueda de imágenes y la generación de títulos de imágenes.

SDXL admite la inferencia de un solo host o de varios hosts con anotaciones de fragmentación. Esto permite que SDXL se entrene y se ejecute en varias máquinas, lo que puede mejorar la eficiencia.

Para obtener más información, consulta el repositorio de modelos generativos de Stability AI y el artículo sobre SDXL.

TPUs

Las TPUs son circuitos integrados para aplicaciones específicas (ASIC) desarrollados a medida por Google que se utilizan para acelerar los modelos de aprendizaje automático y de IA creados con frameworks como TensorFlow, PyTorch y JAX.

Antes de usar las TPUs en GKE, te recomendamos que completes el siguiente plan de formación:

1. Consulta la arquitectura del sistema de las TPU de Cloud para obtener información sobre la disponibilidad de las versiones actuales de las TPU.
2. Consulta información sobre las TPUs en GKE.

En este tutorial se explica cómo servir el modelo SDXL. GKE implementa el modelo en nodos de TPU v5e de un solo host con topologías de TPU configuradas en función de los requisitos del modelo para servir peticiones con baja latencia. En esta guía, el modelo usa un chip de TPU v5e con una topología 1x1.

MaxDiffusion

MaxDiffusion es una colección de implementaciones de referencia, escritas en Python y Jax, de varios modelos de difusión latente que se ejecutan en dispositivos XLA, incluidas las TPUs y las GPUs. MaxDiffusion es un punto de partida para proyectos de difusión, tanto de investigación como de producción.

Para obtener más información, consulta el repositorio MaxDiffusion.

Objetivos

Este tutorial está dirigido a clientes de IA generativa que usen JAX, a usuarios nuevos o actuales de SDXL, y a ingenieros de aprendizaje automático, ingenieros de MLOps (DevOps) o administradores de plataformas que quieran usar las funciones de orquestación de contenedores de Kubernetes para ofrecer LLMs.

Este tutorial abarca los siguientes pasos:

1. Crea un clúster de Autopilot o Estándar de GKE con la topología de TPU recomendada en función de las características del modelo.
2. Crea una imagen de contenedor de inferencia de SDXL.
3. Despliega el servidor de inferencia de SDXL en GKE.
4. Servir una interacción con el modelo a través de una aplicación web.

Arquitectura

En esta sección se describe la arquitectura de GKE que se usa en este tutorial. La arquitectura consta de un clúster de GKE Autopilot o Standard que aprovisiona TPUs y aloja componentes de MaxDiffusion. GKE usa estos componentes para desplegar y servir los modelos.

En el siguiente diagrama se muestran los componentes de esta arquitectura:

Esta arquitectura incluye los siguientes componentes:

• Un clúster regional de Autopilot o Estándar de GKE.
• Un único grupo de nodos de segmento de TPU de un solo host que aloja el modelo SDXL en la implementación de MaxDiffusion.
• El componente Service con un balanceador de carga de tipo ClusterIP. Este servicio distribuye el tráfico entrante a todas las réplicas de MaxDiffusion HTTP.
• El servidor WebApp HTTP con un servicio LoadBalancer externo que distribuye el tráfico entrante y redirige el tráfico de servicio del modelo al servicio ClusterIP.

Antes de empezar

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 (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

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

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 (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

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

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.

     Ir a IAM
  2. Selecciona el proyecto.
  3. Haz clic en Conceder acceso.

  4. En el campo Nuevos principales, introduce tu identificador de usuario. Normalmente, se trata de la dirección de correo de una cuenta de Google.

  5. En la lista Selecciona un rol, elige un rol.
  6. Para conceder más roles, haz clic en Añadir otro rol y añade cada rol adicional.
  7. Haz clic en Guardar.
  • Comprueba que tengas suficiente cuota para los chips TPU v5e PodSlice Lite. En este tutorial, usarás instancias bajo demanda.

  Preparar el entorno

  En este tutorial, usarás Cloud Shell para gestionar los recursos alojados enGoogle Cloud. Cloud Shell tiene preinstalado el software que necesitarás para este tutorial, como kubectl y la CLI de gcloud.

  Para configurar tu entorno con Cloud Shell, sigue estos pasos:

  1. En la Google Cloud consola, inicia una sesión de Cloud Shell haciendo clic en Activar Cloud Shell en la Google Cloud consola. Se iniciará una sesión en el panel inferior de la consola Google Cloud .

  2. Define las variables de entorno predeterminadas:

     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
     

     Sustituye los siguientes valores:

     • PROJECT_ID: tu Google Cloud ID de proyecto.
     • CLUSTER_NAME: el nombre de tu clúster de GKE.
     • REGION_NAME: la región en la que se encuentran tu clúster de GKE, tu segmento de Cloud Storage y tus nodos de TPU. La región contiene zonas en las que están disponibles los tipos de máquinas de TPU v5e (por ejemplo, us-west1, us-west4, us-central1, us-east1, us-east5 o europe-west4).
     • (Solo en clústeres estándar) ZONE: la zona en la que están disponibles los recursos de TPU (por ejemplo, us-west4-a). En el caso de los clústeres de Autopilot, no es necesario especificar la zona, solo la región.

  3. Clona el repositorio de ejemplo y abre el directorio del 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
     

  Crear y configurar Google Cloud recursos

  Sigue estas instrucciones para crear los recursos necesarios.

  Nota: Es posible que tengas que crear una reserva de capacidad para usar algunos aceleradores. Para saber cómo reservar y consumir recursos reservados, consulta el artículo Consumir recursos de zona reservados.

  Crear un clúster de GKE

  Puedes servir SDXL en TPUs en un clúster de Autopilot o Estándar de GKE. Te recomendamos que uses un clúster de Autopilot para disfrutar de una experiencia de Kubernetes totalmente gestionada. Para elegir el modo de funcionamiento de GKE que mejor se adapte a tus cargas de trabajo, consulta Elegir un modo de funcionamiento de GKE.

  Autopilot

  1. En Cloud Shell, ejecuta el siguiente comando:

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

     GKE crea un clúster de Autopilot con nodos de CPU y TPU según lo soliciten las cargas de trabajo desplegadas.

  2. Configura kubectl para que se comunique con tu clúster:

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

  Estándar

  1. Crea un clúster estándar de GKE regional que use Workload Identity Federation para 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 creación del clúster puede tardar varios minutos.

  2. Ejecuta el siguiente comando para crear un grupo de nodos para tu clúster:

     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 crea un grupo de nodos de TPU v5e con una topología 1x1 y un nodo.

     Para crear grupos de nodos con diferentes topologías, consulta cómo planificar la configuración de las TPUs. Asegúrate de actualizar los valores de ejemplo de este tutorial, como cloud.google.com/gke-tpu-topology y google.com/tpu.

  3. Configura kubectl para que se comunique con tu clúster:

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

  Compilar el contenedor de inferencia de SDXL

  Sigue estas instrucciones para crear una imagen de contenedor para el servidor de inferencia de SDXL.

  1. Abre el manifiesto 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. Ejecuta la compilación y crea la imagen del contenedor de inferencia.

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

     La salida contiene la ruta de la imagen de contenedor.

  Implementar el servidor de inferencia de SDXL

  En esta sección, implementará el servidor de inferencia de SDXL. Para desplegar el servidor, en este tutorial se usa un despliegue de Kubernetes. Un Deployment es un objeto de la API de Kubernetes que te permite ejecutar varias réplicas de pods distribuidas entre los nodos de un clúster.

  1. Consulta el archivo de manifiesto de 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. Actualiza el ID del proyecto en el archivo de manifiesto.

     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. Aplica el archivo de manifiesto:

     kubectl apply -f serve_sdxl_v5e.yaml
     

     El resultado debería ser similar al siguiente:

     deployment.apps/max-diffusion-server created
     

  4. Verifica el estado del modelo:

     kubectl get deploy --watch
     

     El resultado debería ser similar al siguiente:

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

  5. Obtén la dirección ClusterIP:

     kubectl get service max-diffusion-server
     

     El resultado contiene un campo ClusterIP. Anota el valor de CLUSTER-IP.

  6. Valida la implementación:

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

     Sustituye CLUSTER_IP por el valor CLUSTER-IP que anotaste anteriormente. El resultado debería ser similar al siguiente:

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

  7. Consulta los registros de la implementación:

     kubectl logs -l app=max-diffusion-server
     

     Cuando finalice la implementación, el resultado será similar al siguiente:

     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.
     

  Desplegar el cliente de la aplicación web

  En esta sección, implementará el cliente de la aplicación web para servir el modelo SDXL.

  1. Consulta el archivo de manifiesto de 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. Ejecuta la compilación y crea la imagen del contenedor de cliente en el directorio build/webapp.

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

     La salida contiene la ruta de la imagen de contenedor.

  3. Abre el manifiesto 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. Edita el ID del proyecto en el archivo de manifiesto:

     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. Aplica el archivo de manifiesto:

     kubectl apply -f serve_sdxl_client.yaml
     

  6. Recupera la dirección IP de LoadBalancer:

     kubectl get service max-diffusion-client-service
     

     El resultado contiene un campo LoadBalancer. Anota el valor de EXTERNAL-IP.

  Éxito: has desplegado correctamente el modelo SDXL con MaxDiffusion y la TPU v5e en GKE.

  Interactuar con el modelo mediante la página web

  1. Acceso a la siguiente URL desde un navegador web:

     http://EXTERNAL_IP:8080
     

     Sustituye EXTERNAL_IP por el valor de EXTERNAL_IP que anotaste anteriormente.

  2. Interactúa con SDXL mediante la interfaz de chat. Añade una petición y haz clic en Enviar. Por ejemplo:

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

  El resultado es una imagen generada por el modelo similar al siguiente ejemplo:

  

  Limpieza

  Para evitar que los recursos utilizados en este tutorial se cobren en tu cuenta de Google Cloud, elimina el proyecto que contiene los recursos o conserva el proyecto y elimina los recursos.

  Eliminar el proyecto

  Precaución: Eliminar un proyecto tiene los siguientes efectos:

  • Se elimina todo el contenido del proyecto. Si has usado un proyecto que ya existía para las tareas de este documento, cuando lo elimines, también se eliminará cualquier otro trabajo que hayas realizado en él.
  • Se pierden los IDs de proyecto personalizados. Cuando creaste este proyecto, es posible que hayas creado un ID de proyecto personalizado que quieras usar en el futuro. Para conservar las URLs que usan el ID del proyecto, como una URL appspot.com, elimina los recursos seleccionados dentro del proyecto en lugar de eliminar todo el proyecto.

  Si tienes previsto consultar varias arquitecturas, tutoriales o guías de inicio rápido, reutilizar los proyectos puede ayudarte a no superar los límites de cuota de proyectos.

  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.

  Eliminar los recursos concretos

  Conserva el proyecto y elimina los recursos concretos, tal como se describe en la sección siguiente. Ejecuta los siguientes comandos y sigue las indicaciones:

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

  Siguientes pasos

  • Configura el tutorial con otras topologías de TPU. Para obtener más información sobre otras topologías de TPU, consulta Planificar la configuración de la TPU.
  • Consulta el código de ejemplo del servidor de inferencia de MaxDiffusion en el repositorio de ejemplo que has clonado en este tutorial.
  • Consulta más información sobre las TPUs en GKE.
  • Consulta el repositorio de GitHub de JetStream.
  • Consulta Vertex AI Model Garden.