Entrega Stable Diffusion XL (SDXL) con TPU en GKE con MaxDiffusion

Autopilot estándar

En este instructivo, se muestra cómo entregar un modelo de generación de imágenes de SDXL con unidades de procesamiento tensorial (TPU) en Google Kubernetes Engine (GKE) con MaxDiffusion. En este instructivo, descargarás el modelo de Hugging Face y, luego, lo implementarás en un clúster de Autopilot o Standard con un contenedor que ejecute MaxDiffusion.

Esta guía es un buen punto de partida si necesitas el control detallado, la personalización, la escalabilidad, la resiliencia, la portabilidad y la rentabilidad de Kubernetes administrado cuando implementas y entregas tus cargas de trabajo de IA/AA. Si necesitas una plataforma de IA administrada unificada para compilar y entregar modelos de AA con rapidez de forma rentable, te recomendamos que pruebes nuestra solución de implementación de Vertex AI.

Fondo

La entrega de SDXL con TPU en GKE con MaxDiffusion te permite compilar una solución de entrega sólida y lista para la producción con todos los beneficios de Kubernetes administrado, incluida la rentabilidad, la escalabilidad y una mayor disponibilidad. En esta sección, se describen las tecnologías clave que se usan en este instructivo.

Stable Diffusion XL (SDXL)

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

SDXL admite inferencia de uno o 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 documento de SDXL.

TPU

Las TPU son circuitos integrados personalizados específicos de aplicaciones (ASIC) de Google que se usan para acelerar el aprendizaje automático y los modelos de IA compilados con frameworks como el siguiente:TensorFlow, PyTorch yJAX.

Antes de usar las TPU en GKE, te recomendamos que completes la siguiente ruta de aprendizaje:

1. Obtén información sobre la disponibilidad actual de la versión de TPU con la arquitectura del sistema de Cloud TPU.
2. Obtén información sobre las TPU en GKE.

En este instructivo, se aborda la entrega del modelo SDXL. GKE implementa el modelo en los nodos TPU v5e de host único con topologías de TPU configuradas según los requisitos del modelo para entregar mensajes con baja latencia. En esta guía, el modelo usa un chip 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 latentes que se ejecutan en dispositivos XLA, incluidas TPU y GPU. MaxDiffusion es un punto de partida de los proyectos de difusión para la investigación y la producción.

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

Objetivos

Este instructivo está dirigido a clientes de IA generativa que usan JAX, usuarios nuevos o existentes de SDXL e ingenieros de AA, ingenieros de MLOps (DevOps) o administradores de plataformas interesados en usar las funciones de organización de contenedores de Kubernetes para entrega de LLM.

En este instructivo, se abarcan los siguientes pasos:

1. Crea un clúster de GKE en modo Autopilot o Standard con la topología de TPU recomendada según las características del modelo.
2. Compilar una imagen de contenedor de inferencia de SDXL.
3. Implementar el servidor de inferencia de SDXL en GKE.
4. Entregar una interacción con el modelo a través de una app web

Arquitectura

En esta sección, se describe la arquitectura de GKE que se usa en este instructivo. La arquitectura consiste en un clúster de GKE Autopilot o Standard que aprovisiona TPU y aloja componentes de MaxDiffusion. GKE usa estos componentes para implementar y entregar los modelos.

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

Esta arquitectura incluye los siguientes componentes:

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

Antes de comenzar

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 role (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 role (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 Otorgar acceso.

  4. En el campo Principales nuevas, ingresa tu identificador de usuario. Esta suele ser la dirección de correo electrónico de una Cuenta de Google.

  5. En la lista Seleccionar un rol, elige uno.
  6. Para otorgar roles adicionales, haz clic en Agregar otro rol y agrega uno más.
  7. Haz clic en Guardar.
  • Verifica que tengas suficiente cuota para los chips TPU v5e PodSlice Lite. En este instructivo, usarás instancias bajo demanda.

  Prepare el entorno

  En este instructivo, usarás Cloud Shell para administrar recursos alojados enGoogle Cloud. Cloud Shell tiene preinstalado el software que necesitarás para este instructivo, incluidos kubectl y la CLI de gcloud.

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

  1. En la Google Cloud consola, haz clic en Activar Cloud Shell en la Google Cloud consola para iniciar una sesión de Cloud Shell. Esto inicia una sesión en el panel inferior de la consola de Google Cloud .

  2. Configura 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
     

     Reemplaza los siguientes valores:

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

  3. Clona el repositorio de ejemplo y abre el directorio del instructivo:

     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
     

  Crea y configura recursos de Google Cloud

  Sigue estas instrucciones para crear los recursos necesarios.

  Nota: Es posible que debas crear una reserva de capacidad para el uso de algunos aceleradores. Para obtener información sobre cómo reservar y consumir recursos reservados, consulta Consume recursos zonales reservados.

  Cree un clúster de GKE

  Puedes entregar SDXL en TPU en un clúster de GKE Autopilot o Standard. Te recomendamos que uses un clúster de Autopilot para una experiencia de Kubernetes completamente administrada. Para elegir el modo de operación de GKE que se adapte mejor a tus cargas de trabajo, consulta Elige un modo de operación de GKE.

  Autopilot

  1. En Cloud Shell, ejecute 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 en modo Autopilot con nodos de CPU y TPU según lo solicitan las cargas de trabajo implementadas.

  2. Configura kubectl para comunicarse con tu clúster:

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

  Estándar

  1. Crea un clúster de GKE Estándar regional que use la federación de identidades para cargas de trabajo en 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 tomar varios minutos.

  2. Ejecuta el siguiente comando para crear un grupo de nodos para el 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 TPU v5e con una topología 1x1 y un nodo.

     Para crear grupos de nodos con diferentes topologías, aprende a planificar tu configuración de TPU. Asegúrate de actualizar los valores de muestra de este tutor, como cloud.google.com/gke-tpu-topology y google.com/tpu.

  3. Configura kubectl para comunicarse con tu clúster:

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

  Compila el contenedor de inferencia de SDXL

  Sigue estas instrucciones a fin de compilar 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 una imagen de contenedor de inferencia.

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

     El resultado contiene la ruta de la imagen de contenedor.

  Implementa el servidor de inferencia de SDXL

  En esta sección, implementarás el servidor de inferencia de SDXL. Para implementar el servidor, en este instructivo se usa una implementación de Kubernetes. Un Deployment es un objeto de la API de Kubernetes que te permite ejecutar varias réplicas de Pods que se distribuyen entre los nodos de un clúster.

  1. Explora el manifiesto 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 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 manifiesto

     kubectl apply -f serve_sdxl_v5e.yaml
     

     El resultado es similar a este:

     deployment.apps/max-diffusion-server created
     

  4. Verifica el estado del modelo:

     kubectl get deploy --watch
     

     El resultado es similar a este:

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

  5. Recupera la dirección ClusterIP:

     kubectl get service max-diffusion-server
     

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

  6. Valida el Deployment:

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

     Reemplaza CLUSTER_IP por el valor CLUSTER-IP que anotaste antes. El resultado es similar a este:

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

  7. Visualiza los registros desde el Deployment:

     kubectl logs -l app=max-diffusion-server
     

     Cuando el Deployment finaliza, el resultado es 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.
     

  Implementa el cliente de webapp

  En esta sección, implementarás el cliente de webapp para entregar el modelo SDXL.

  1. Explora el manifiesto 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 de contenedor de cliente en el directorio build/webapp.

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

     El resultado 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 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 manifiesto

     kubectl apply -f serve_sdxl_client.yaml
     

  6. Recupera la dirección IP LoadBalancer:

     kubectl get service max-diffusion-client-service
     

     El resultado contiene un campo LoadBalancer. Toma nota del valor de EXTERNAL-IP.

  Completado correctamente: Implementaste con éxito el modelo SDXL con MaxDiffusion y TPU v5e en GKE.

  Interactúa con el modelo mediante la página web

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

     http://EXTERNAL_IP:8080
     

     Reemplaza EXTERNAL_IP por el valor de EXTERNAL_IP que anotaste antes.

  2. Interactuar con SDXL mediante la interfaz de chat. Agrega una instrucció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 un modelo similar al siguiente ejemplo:

  

  Limpia

  Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos usados en este instructivo, borra el proyecto que contiene los recursos o conserva el proyecto y borra los recursos individuales.

  Borra el proyecto

  Precaución: Borrar un proyecto tiene las siguientes consecuencias:

  • Se borra todo el contenido. Si usaste un proyecto existente para las tareas de este documento, cuando lo borres, también se borrará cualquier otro trabajo que hayas realizado en el proyecto.
  • Se pierden los IDs personalizados de los proyectos. Cuando creaste este proyecto, es posible que hayas creado un ID personalizado del proyecto que desees utilizar en el futuro. Para conservar las URLs que usan el ID del proyecto, como una URL appspot.com, borra los recursos seleccionados dentro del proyecto, en lugar de todo el proyecto.

  Si planeas explorar varios instructivos, arquitecturas o guías de inicio rápido, reutilizar proyectos puede ser útil para no exceder los límites de cuota de los 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.

  Borra los recursos individuales

  Conserva el proyecto y borra los recursos individuales, como se describe en la siguiente sección. Ejecuta los siguientes comandos y sigue las instrucciones:

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

  ¿Qué sigue?

  • Configura el instructivo con otras topologías de TPU. Para obtener más información sobre otras topologías de TPU, consulta Planifica tu configuración de TPU.
  • Explora el código de muestra del servidor de inferencia de MaxDiffusion en el repositorio de muestra que clonaste en este instructivo.
  • Obtén más información sobre las TPU en GKE.
  • Explora el repositorio de GitHub de JetStream.
  • Explora Vertex AI Model Garden.