Servir modelos de código abierto mediante TPUs en GKE con Optimum TPU

Estándar

En este tutorial se muestra cómo servir modelos de código abierto de modelos de lenguaje grandes (LLMs) mediante unidades de procesamiento tensorial (TPUs) en Google Kubernetes Engine (GKE) con el framework de servicio Optimum TPU de Hugging Face. En este tutorial, descargarás modelos de código abierto de Hugging Face y los desplegarás en un clúster Standard de GKE mediante un contenedor que ejecute Optimum TPU.

Esta guía proporciona un punto de partida si necesitas el control granular, 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.

Este tutorial está dirigido a clientes de IA generativa del ecosistema de Hugging Face, usuarios nuevos o actuales de GKE, ingenieros de aprendizaje automático, ingenieros de MLOps (DevOps) o administradores de plataformas que estén interesados en usar las funciones de orquestación de contenedores de Kubernetes para ofrecer LLMs.

Te recordamos que tienes varias opciones para la inferencia de LLMs en Google Cloud, que abarcan ofertas como Vertex AI, GKE y Google Compute Engine, donde puedes incorporar bibliotecas de servicio como JetStream, vLLM y otras ofertas de partners. Por ejemplo, puedes usar JetStream para obtener las optimizaciones más recientes del proyecto. Si prefieres las opciones de Hugging Face, puedes usar Optimum TPU.

Optimum TPU admite las siguientes funciones:

• Procesamiento por lotes continuo
• Streaming de tokens
• Búsqueda voraz y muestreo multinomial con transformadores.

Objetivos

1. Prepara un clúster de GKE Standard con la topología de TPU recomendada en función de las características del modelo.
2. Despliega Optimum TPU en GKE.
3. Usa Optimum TPU para servir los modelos admitidos a través de curl.

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.
  • Crea una cuenta de Hugging Face si aún no tienes una.
  • Asegúrate de que tu proyecto tiene suficiente cuota para usar Cloud TPU en GKE.

  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
     export PROJECT_ID=$(gcloud config get project)
     export CLUSTER_NAME=CLUSTER_NAME
     export REGION=REGION_NAME
     export ZONE=ZONE
     export HF_TOKEN=HF_TOKEN
     

     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.
     • HF_TOKEN: tu token de Hugging Face.

  3. Clona el repositorio de Optimum TPU:

     git clone https://github.com/huggingface/optimum-tpu.git
     

  Acceder al modelo

  Puedes usar los modelos Gemma 2B o Llama 3 8B. En este tutorial nos centraremos en estos dos modelos, pero Optimum TPU admite más modelos.

  Gemma 2B

  Para acceder a los modelos de Gemma e implementarlos en GKE, primero debes firmar el acuerdo de consentimiento de licencia y, a continuación, generar un token de acceso de Hugging Face.

  Firmar el contrato de consentimiento de licencia

  Debes firmar el acuerdo de consentimiento para usar Gemma. Te las indicamos a continuación:

  1. Accede a la página de consentimiento del modelo.
  2. Verifica el consentimiento con tu cuenta de Hugging Face.
  3. Acepta los términos del modelo.

  Generar un token de acceso

  Genera un nuevo token de Hugging Face si aún no tienes uno:

  1. Haz clic en Tu perfil > Configuración > Tokens de acceso.
  2. Haz clic en New Token (Nuevo token).
  3. Especifica el nombre que quieras y un rol de al menos Read.
  4. Haz clic en Generar un token.
  5. Copia el token generado en el portapapeles.

  Llama3 8B

  Debes firmar el acuerdo de consentimiento para usar Llama 3 8b en el repositorio de Hugging Face.

  Generar un token de acceso

  Genera un nuevo token de Hugging Face si aún no tienes uno:

  1. Haz clic en Tu perfil > Configuración > Tokens de acceso.
  2. Selecciona New Token (Nuevo token).
  3. Especifica el nombre que quieras y un rol de al menos Read.
  4. Selecciona Generar un token.
  5. Copia el token generado en el portapapeles.

  Crear un clúster de GKE

  Crea un clúster de GKE Standard con un nodo de 1 CPU:

  gcloud container clusters create CLUSTER_NAME \
      --project=PROJECT_ID \
      --num-nodes=1 \
      --location=ZONE
  

  Crear un grupo de nodos de TPU

  Crea un grupo de nodos de TPU v5e con 1 nodo y 8 chips:

  gcloud container node-pools create tpunodepool \
      --location=ZONE \
      --num-nodes=1 \
      --machine-type=ct5lp-hightpu-8t \
      --cluster=CLUSTER_NAME
  

  Si hay recursos de TPU disponibles, GKE aprovisiona el grupo de nodos. Si los recursos de TPU no están disponibles temporalmente, el resultado muestra un mensaje de error GCE_STOCKOUT. Para solucionar los errores de falta de stock de TPUs, consulta el artículo No hay suficientes recursos de TPU para satisfacer la solicitud de TPU.

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

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

  Compila el contenedor

  Ejecuta el comando make para compilar la imagen

  cd optimum-tpu && make tpu-tgi
  

  Enviar la imagen a Artifact Registry

  gcloud artifacts repositories create optimum-tpu --repository-format=docker --location=REGION_NAME && \
  gcloud auth configure-docker REGION_NAME-docker.pkg.dev && \
  docker image tag huggingface/optimum-tpu REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest && \
  docker push REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
  

  Crear un secreto de Kubernetes para las credenciales de Hugging Face

  Crea un secreto de Kubernetes que contenga el token de Hugging Face:

  kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  

  Implementar Optimum TPU

  Para desplegar Optimum TPU, en este tutorial se usa un Deployment 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.

  Gemma 2B

  1. Guarda el siguiente archivo de manifiesto de implementación como optimum-tpu-gemma-2b-2x4.yaml:

     apiVersion: apps/v1
     kind: Deployment
     metadata:
       name: tgi-tpu
     spec:
       replicas: 1
       selector:
         matchLabels:
           app: tgi-tpu
       template:
         metadata:
           labels:
             app: tgi-tpu
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-topology: 2x4
             cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
           containers:
           - name: tgi-tpu
             image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
             args:
             - --model-id=google/gemma-2b
             - --max-concurrent-requests=4
             - --max-input-length=8191
             - --max-total-tokens=8192
             - --max-batch-prefill-tokens=32768
             - --max-batch-size=16
             securityContext:
                 privileged: true
             env:
               - name: HF_TOKEN
                 valueFrom:
                   secretKeyRef:
                     name: hf-secret
                     key: hf_api_token
             ports:
             - containerPort: 80
             resources:
               limits:
                 google.com/tpu: 8
             livenessProbe:
               httpGet:
                 path: /health
                 port: 80
               initialDelaySeconds: 300
               periodSeconds: 120
     
     ---
     apiVersion: v1
     kind: Service
     metadata:
       name: service
     spec:
       selector:
         app: tgi-tpu
       ports:
         - name: http
           protocol: TCP
           port: 8080
           targetPort: 80
     

     Este manifiesto describe una implementación de TPU Optimum con un balanceador de carga interno en el puerto TCP 8080.

  2. Aplicar el archivo de manifiesto

     kubectl apply -f optimum-tpu-gemma-2b-2x4.yaml
     

  Llama3 8B

  1. Guarda el siguiente archivo de manifiesto como optimum-tpu-llama3-8b-2x4.yaml:

     apiVersion: apps/v1
     kind: Deployment
     metadata:
       name: tgi-tpu
     spec:
       replicas: 1
       selector:
         matchLabels:
           app: tgi-tpu
       template:
         metadata:
           labels:
             app: tgi-tpu
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-topology: 2x4
             cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
           containers:
           - name: tgi-tpu
             image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
             args:
             - --model-id=meta-llama/Meta-Llama-3-8B
             - --max-concurrent-requests=4
             - --max-input-length=8191
             - --max-total-tokens=8192
             - --max-batch-prefill-tokens=32768
             - --max-batch-size=16
             env:
               - name: HF_TOKEN
                 valueFrom:
                   secretKeyRef:
                     name: hf-secret
                     key: hf_api_token
             ports:
             - containerPort: 80
             resources:
               limits:
                 google.com/tpu: 8
             livenessProbe:
               httpGet:
                 path: /health
                 port: 80
               initialDelaySeconds: 300
               periodSeconds: 120
     ---
     apiVersion: v1
     kind: Service
     metadata:
       name: service
     spec:
       selector:
         app: tgi-tpu
       ports:
         - name: http
           protocol: TCP
           port: 8080
           targetPort: 80
     

     Este manifiesto describe una implementación de TPU Optimum con un balanceador de carga interno en el puerto TCP 8080.

  2. Aplicar el archivo de manifiesto

     kubectl apply -f optimum-tpu-llama3-8b-2x4.yaml
     

  Consulta los registros de la implementación en ejecución:

  kubectl logs -f -l app=tgi-tpu
  

  La salida debería ser similar a la siguiente:

  2024-07-09T22:39:34.365472Z  WARN text_generation_router: router/src/main.rs:295: no pipeline tag found for model google/gemma-2b
  2024-07-09T22:40:47.851405Z  INFO text_generation_router: router/src/main.rs:314: Warming up model
  2024-07-09T22:40:54.559269Z  INFO text_generation_router: router/src/main.rs:351: Setting max batch total tokens to 64
  2024-07-09T22:40:54.559291Z  INFO text_generation_router: router/src/main.rs:352: Connected
  2024-07-09T22:40:54.559295Z  WARN text_generation_router: router/src/main.rs:366: Invalid hostname, defaulting to 0.0.0.0
  

  Asegúrate de que el modelo se haya descargado por completo antes de pasar a la siguiente sección.

  Aplicar el modelo

  Configura la redirección de puertos al modelo:

  kubectl port-forward svc/service 8080:8080
  

  Interactuar con el servidor de modelos mediante curl

  Verifica tus modelos desplegados:

  En una nueva sesión de terminal, usa curl para chatear con el modelo:

  curl 127.0.0.1:8080/generate     -X POST     -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":40}}'     -H 'Content-Type: application/json'
  

  La salida debería ser similar a la siguiente:

  {"generated_text":"\n\nDeep learning is a subset of machine learning that uses artificial neural networks to learn from data.\n\nArtificial neural networks are inspired by the way the human brain works. They are made up of multiple layers"}
  

  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 los recursos desplegados

  Para evitar que se apliquen cargos en tu cuenta de Google Cloud por los recursos que has creado en esta guía, ejecuta el siguiente comando:

  gcloud container clusters delete CLUSTER_NAME \
    --location=ZONE
  

  Siguientes pasos

  • Consulta la documentación de Optimum TPU.
  • Descubre cómo puedes ejecutar modelos Gemma en GKE y cómo ejecutar cargas de trabajo de IA y aprendizaje automático optimizadas con las funciones de orquestación de la plataforma GKE.
  • Consulta más información sobre las TPUs en GKE.