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. Entrega el modelo y, luego, interactúa con él 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

Accede a tu cuenta de Google Cloud . Si eres nuevo en Google Cloud, crea una cuenta para evaluar el rendimiento de nuestros productos en situaciones reales. Los clientes nuevos también obtienen $300 en créditos gratuitos para ejecutar, probar y, además, implementar cargas de trabajo.

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

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

• Asegúrate de tener los siguientes roles en el proyecto: roles/container.admin, roles/iam.serviceAccountAdmin, roles/artifactregistry.admin, roles/cloudbuild.builds.editor

  Verifica los roles

  1. En la consola de Google Cloud , dirígete a la página IAM.

     Ir a IAM
  2. Selecciona el proyecto.

  3. En la columna Principal, busca todas las filas que te identifiquen a ti o a un grupo en el que se te incluya. Para saber en qué grupos estás incluido, comunícate con tu administrador.

  4. Para todas las filas en las que se te especifique o se te incluya, verifica la columna Rol para ver si la lista de roles incluye los roles necesarios.

  Otorga los roles

  1. En la consola de Google Cloud , dirígete a la página IAM.

     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. Haz clic en Seleccionar un rol y, luego, busca el rol.
  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 CLUSTER_VERSION=CLUSTER_VERSION
   export REGION=REGION_NAME
   export ZONE=ZONE
   

   Reemplaza los siguientes valores:

   • PROJECT_ID: Es el Google Cloud ID del proyecto.
   • CLUSTER_NAME: Es el nombre del clúster de GKE.
   • CLUSTER_VERSION : Es la versión de GKE. Debes especificar una versión de GKE que admita la TPU Trillium (v6e). Para obtener más información, consulta Cómo validar la disponibilidad de TPU en 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, ejecuta el siguiente comando:

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

   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 instructivo, 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 su 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. En la Google Cloud consola, ve a la página Administrar recursos.

   Ir a Administrar recursos

2. En la lista de proyectos, elige el proyecto que quieres borrar y haz clic en Borrar.
3. En el diálogo, escribe el ID del proyecto y, luego, haz clic en Cerrar para borrar el proyecto.

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.