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

Estándar

En este instructivo, se muestra cómo entregar modelos de código abierto de modelos de lenguaje grandes (LLM) con unidades de procesamiento tensorial (TPU) en Google Kubernetes Engine (GKE) con el framework de entrega Optimum TPU desde Hugging Face. En este instructivo, descargarás modelos de código abierto de Hugging Face y los implementarás en un clúster Standard de GKE con un contenedor que ejecute Optimum TPU.

Esta guía proporciona un punto de partida si necesitas el control detallado, la escalabilidad, la resiliencia, la portabilidad y la rentabilidad de Kubernetes administrado cuando implementas y entregas tus cargas de trabajo de IA/AA.

Este instructivo está dirigido a clientes de IA generativa en el ecosistema de Hugging Face, usuarios nuevos o existentes de GKE, 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.

Los productos deGoogle Cloud , como GKE, Vertex AI y Compute Engine, admiten varias bibliotecas de servicio, como JetStream, vLLM y otras ofertas de socios. 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.

La TPU óptima admite las siguientes funciones:

• Procesamiento por lotes continuo
• Transmisión de tokens
• Muestreo multinomial y búsqueda voraz con transformadores

Objetivos

1. Prepara un clúster estándar de GKE con la topología de TPU recomendada según las características del modelo.
2. Implementa Optimum TPU en GKE.
3. Usa Optimum TPU para entregar los modelos compatibles a través de curl.

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

If you're using an existing project for this guide, verify that you have the permissions required to complete this guide. If you created a new project, then you already have the required permissions.

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

If you're using an existing project for this guide, verify that you have the permissions required to complete this guide. If you created a new project, then you already have the required permissions.

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

• Crea una cuenta de Hugging Face, si todavía no la tienes.
• Asegúrate de que tu proyecto tenga la cuota suficiente para Cloud TPU en GKE.

Roles obligatorios

Para obtener los permisos que necesitas para configurar clústeres y cargas de trabajo, pídele a tu administrador que te otorgue los siguientes roles de IAM en el proyecto:

• Administrador de cuenta de servicio (roles/iam.serviceAccountAdmin)
• Administra clústeres de GKE: Administrador de Kubernetes Engine (roles/container.admin)
• Compila y envía imágenes a Artifact Registry: Administrador de Artifact Registry (roles/artifactregistry.admin)

Para obtener más información sobre cómo otorgar roles, consulta Administra el acceso a proyectos, carpetas y organizaciones.

También puedes obtener los permisos necesarios a través de roles personalizados o cualquier otro rol predefinido.

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 consola de Google Cloud , activa Cloud Shell.

   Activa Cloud Shell

   En la parte inferior de la consola de Google Cloud , se inicia una sesión de Cloud Shell que muestra una ventana emergente con una línea de comandos. Cloud Shell es un entorno de shell con Google Cloud CLI ya instalada y con valores ya establecidos para el proyecto actual. La sesión puede tardar unos segundos en inicializarse.

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
   export HF_TOKEN=HF_TOKEN
   

   Reemplaza los siguientes valores:

   • PROJECT_ID: Es el Google Cloud ID del 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.
   • HF_TOKEN: Tu token de Hugging Face.

3. Clona el repositorio de Optimum TPU:

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

Obtén acceso al modelo

Puedes usar los modelos Gemma 2B o Llama3 8B. En este instructivo, nos enfocamos en estos dos modelos, pero Optimum TPU admite más modelos.

Gemma 2B

Para obtener acceso a los modelos de Gemma para la implementación en GKE, primero debes firmar el contrato de consentimiento de licencia y, luego, generar un token de acceso de Hugging Face.

Firma el acuerdo de consentimiento de licencia

Debes firmar el acuerdo de consentimiento para usar Gemma. Sigue estas instrucciones:

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.

Genera 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 Token nuevo.
3. Especifica el nombre que desees y un rol de al menos Read.
4. Haz clic en Generar un token.
5. Copia el token generado al portapapeles.

Llama3 8B

Debes firmar el acuerdo de consentimiento para usar Llama3 8b en el repositorio de Hugging Face

Genera 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 Token nuevo.
3. Especifica el nombre que desees y un rol de al menos Read.
4. Selecciona Genera un token.
5. Copia el token generado al portapapeles.

Cree un clúster de GKE

Crea un clúster estándar de GKE con 1 nodo de CPU:

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

Crea un grupo de nodos TPU

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

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

Si los recursos de TPU están disponibles, GKE aprovisiona el grupo de nodos. Si los recursos de TPU no están disponibles temporalmente, el resultado mostrará un mensaje de error GCE_STOCKOUT. Para solucionar problemas de disponibilidad de recursos, consulta Recursos de TPU insuficientes para satisfacer la solicitud de TPU.

Crea el contenedor

Ejecuta el siguiente comando para compilar la imagen:

cd optimum-tpu && make tpu-tgi

Envía la imagen al registro de 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

Crea un secreto de Kubernetes para las credenciales de Hugging Face

Crea un Secret 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 -

Implementa Optimum TPU

Para implementar Optimum TPU, 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.

Gemma 2B

1. Guarda el siguiente manifiesto de Deployment 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
   

   En este manifiesto, se describe una implementación de Optimum TPU con un balanceador de cargas interno en el puerto TCP 8080.

2. Aplica el manifiesto

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

Llama3 8B

1. Guarda el siguiente 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
   

   En este manifiesto, se describe una implementación de Optimum TPU con un balanceador de cargas interno en el puerto TCP 8080.

2. Aplica el manifiesto

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

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

kubectl logs -f -l app=tgi-tpu

El resultado debería ser similar al siguiente ejemplo:

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 continuar con la siguiente sección.

Entrega el modelo

Configura la redirección de puertos al modelo:

kubectl port-forward svc/service 8080:8080

Interactúa con el modelo con curl

Verifica tus modelos implementados:

En una sesión de terminal nueva, usa curl para chatear con tu 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'

El resultado debería ser similar al siguiente ejemplo:

{"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"}

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

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

gcloud container clusters delete CLUSTER_NAME \
  --location=REGION_NAME

¿Qué sigue?

• Explora la documentación de Optimum TPU.
• Descubre cómo puedes ejecutar modelos de Gemma en GKE y cómo ejecutar cargas de trabajo de IA/AA optimizadas con las funciones de organización de la plataforma de GKE.
• Obtén más información sobre las TPU en GKE.