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.

  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

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

  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

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

     Ir a IAM
  2. Selecciona el proyecto.
  3. Haz clic en Grant access.

  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 un rol.
  6. Para otorgar funciones adicionales, haz clic en Agregar otro rol y agrega cada rol adicional.
  7. Haz clic en Guardar.

• Asegúrate de tener 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 en Google 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 consola de Google Cloud, haz clic en Activar Cloud Shell en la consola de Google Cloud 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
   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: El ID del proyecto de Google Cloud.
   • 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} \
     --region=${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 \
     --region=${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 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

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
   sed -i "s|PROJECT_ID|$PROJECT_ID|g" serve_sdxl_v5e.yaml
   sed -i "s|REGION|$REGION|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:

   watch kubectl get deploy
   

   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
   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. 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 en el proyecto. 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 ID personalizados de proyectos. Cuando creaste este proyecto, es posible que hayas creado un ID del proyecto personalizado que desees utilizar en el futuro. Para conservar las URL que utilizan el ID del proyecto, como una URL appspot.com, borra los recursos seleccionados dentro del proyecto en lugar de borrar todo el proyecto.

Si planeas explorar varias infraestructuras, instructivos y guías de inicio rápido la reutilización de proyectos puede ayudarte a evitar exceder los límites de las cuotas del proyecto.

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} --region=${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.