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

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 últimas optimizaciones 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.

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.

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