Servir LLMs mediante TPUs de varios hosts en GKE con JetStream y Pathways

Estándar

En esta guía se explica cómo servir modelos de lenguaje extensos (LLMs) de vanguardia, como Llama 3.1 405B, en Google Kubernetes Engine (GKE) mediante unidades de procesamiento tensorial (TPUs) en varios nodos.

En esta guía se muestra cómo usar tecnologías portátiles de código abierto (Kubernetes, JetStream, Pathways on Cloud y la API LeaderWorkerSet [LWS]) para desplegar y servir cargas de trabajo de IA y aprendizaje automático en GKE. Para ello, se aprovechan el control granular, la escalabilidad, la resiliencia, la portabilidad y la rentabilidad de GKE.

Importante: Para acceder a Pasos, ponte en contacto con tu representante de cuentas de Google Cloud.

Fondo

Los modelos de lenguaje extenso han aumentado de tamaño y ya no caben en una sola porción de TPU de host. Para la inferencia de aprendizaje automático, puedes usar Pathways en Cloud para ejecutar inferencias de varios hosts a gran escala en GKE en varios nodos de TPU interconectados. En esta guía, se explica cómo aprovisionar un clúster de GKE con las slices de TPU de varios hosts, usar los archivos binarios de Pathways on Cloud, iniciar el servidor JetStream con el framework MaxText y hacer solicitudes de inferencia de varios hosts.

Si sirves un LLM mediante TPUs en GKE con JetStream, MaxText y Pathways, puedes crear una solución de servicio estable y lista para producción con todas las ventajas de Kubernetes gestionado, como la rentabilidad, la escalabilidad y la mayor disponibilidad. En esta sección se describen las tecnologías clave que se usan en este tutorial.

Acerca de las TPUs

Las TPUs son circuitos integrados para aplicaciones específicas (ASIC) desarrollados a medida por Google que se utilizan para acelerar modelos de aprendizaje automático e IA creados con frameworks como TensorFlow, PyTorch y JAX.

Antes de usar las TPUs en GKE, te recomendamos que completes el siguiente plan de formación:

1. Consulta la arquitectura del sistema de las TPU de Cloud para obtener información sobre la disponibilidad de las versiones actuales de las TPU.
2. Consulta información sobre las TPUs en GKE.

En este tutorial se explica cómo servir el modelo Llama 3.1-405B. GKE implementa el modelo en nodos de TPU v6e de varios hosts con topologías de TPU configuradas en función de los requisitos del modelo para servir peticiones con baja latencia.

Rutas de aprendizaje en Cloud

Pathways es una capa de orquestación a gran escala para aceleradores. Pathways se ha diseñado específicamente para permitir la exploración de nuevos sistemas e ideas de investigación de aprendizaje automático, al tiempo que mantiene el rendimiento de vanguardia de los modelos actuales. Pathways permite que un solo proceso de cliente de JAX coordine la computación en una o varias grandes porciones de TPU, lo que agiliza las computaciones de aprendizaje automático que abarcan cientos o miles de chips de TPU.

JetStream

JetStream es un framework de servicio de inferencia de código abierto desarrollado por Google. JetStream permite realizar inferencias de alto rendimiento, alto volumen de procesamiento y memoria optimizada en TPUs y GPUs. JetStream ofrece optimizaciones de rendimiento avanzadas, como la creación de minilotes continua, la optimización de la caché de valores de clave y las técnicas de cuantización, para facilitar la implementación de LLMs. JetStream permite que PyTorch/XLA y JAX TPU sirvan para optimizar el rendimiento.

MaxText

MaxText es una implementación de LLM de JAX de alto rendimiento, escalable y adaptable, creada a partir de bibliotecas de JAX de código abierto, como Flax, Orbax y Optax. La implementación de LLM solo con decodificador de MaxText está escrita en Python. Aprovecha el compilador XLA para conseguir un alto rendimiento sin necesidad de crear kernels personalizados.

Para obtener más información sobre los modelos y tamaños de parámetros más recientes que admite MaxText, consulta el repositorio del proyecto MaxText.

Llama 3.1 405B

Llama 3.1 405B es un modelo de lenguaje extenso de Meta diseñado para llevar a cabo diversas tareas de procesamiento del lenguaje natural, como la generación de texto, la traducción y la respuesta a preguntas. GKE ofrece la infraestructura necesaria para satisfacer las necesidades de entrenamiento y servicio distribuidos de modelos de esta escala.

Para obtener más información, consulta la documentación de Llama.

Arquitectura

En esta sección se describe la arquitectura de GKE que se usa en este tutorial. La arquitectura incluye un clúster de GKE Standard que aprovisiona TPUs y aloja componentes de JetStream y Pathways para desplegar y servir el modelo.

En el siguiente diagrama se muestran los componentes de esta arquitectura:

Esta arquitectura incluye los siguientes componentes:

• Un clúster regional de GKE Standard.
• Un grupo de nodos de segmento de TPU multihost que aloja el despliegue de JetStream y los componentes de Pathways.
• Pathways resource manager gestiona los recursos del acelerador y coordina la asignación de aceleradores a los trabajos de los usuarios.
• El Pathways client se coordina con el Pathways resource manager para determinar dónde se colocan los programas compilados para su ejecución.
• El Pathways worker se ejecuta y realiza cálculos en máquinas aceleradoras, y envía datos a tu carga de trabajo a través del servidor proxy IFRT.
• IFRT proxy client implementa la API Interim Framework Runtime (IFRT) de OSS y actúa como puente de comunicación entre tu carga de trabajo y los componentes de Pathways.
• El IFRT proxy server recibe solicitudes del IFRT proxy client y las reenvía al Pathways client, distribuyendo el trabajo.
• El contenedor JetStream-Pathways proporciona un servidor de inferencia basado en JAX que recibe solicitudes de inferencia y delega sus procesos de ejecución en Pathways workers.
• El componente Service distribuye el tráfico entrante a todas las réplicas de JetStream HTTP.
• JetStream HTTP es un servidor HTTP que acepta solicitudes como envoltorio del formato requerido de JetStream y las envía al cliente GRPC de JetStream.

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/resourcemanager.projectIamAdmin

  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.
  • Asegúrate de que tienes suficiente cuota para dieciséis chips TPU v6e PodSlice Lite. En este tutorial, usarás instancias bajo demanda.
  • Asegúrate de que tu Google Cloud proyecto esté incluido en la lista de permitidos de Pathways.

  Acceder al modelo

  Para acceder al punto de control de Meta Llama 3.1-405B para implementarlo en GKE, sigue estos pasos:

  1. Firma el contrato de consentimiento de licencia.
  2. Accede a la página de descargas de Meta Llama.
  3. Revisa y acepta los Términos y Condiciones del modelo para obtener la URL necesaria para descargarlo.
  4. Para descargar el punto de control del modelo, busca el ID del modelo correspondiente. Para ver una lista de los modelos admitidos y sus IDs, consulta la documentación de la CLI de Llama. Por ejemplo, usa Llama 3.1-405B-Instruct:bf16-mp16 para el modelo Llama 3.1-405B.

  Preparar el entorno

  En este tutorial, usarás Cloud Shell para gestionar los recursos alojados enGoogle Cloud. Cloud Shell tiene preinstalado el software que necesitas 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
     gcloud config set billing/quota_project PROJECT_ID
     export PROJECT_ID=$(gcloud config get project)
     export CLUSTER_NAME=CLUSTER_NAME
     export BUCKET_NAME=BUCKET_NAME
     export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
     export NODE_LOCATION=NODE_LOCATION
     export CLUSTER_VERSION=CLUSTER_VERSION
     export MACHINE_TYPE=ct6e-standard-4t
     export TPU_TYPE=v6e
     export TOPOLOGY=4x4
     export WORKERS_PER_SLICE=4
     

     Sustituye los siguientes valores:

     • PROJECT_ID: tu Google Cloud ID de proyecto.
     • CLUSTER_NAME: el nombre de tu clúster de GKE.
     • BUCKET_NAME: el nombre de tu segmento de Cloud Storage. No es necesario que especifiques el prefijo gs://.
     • CONTROL_PLANE_LOCATION: la región de Compute Engine en la que se encuentran tu clúster de GKE, tu bucket de Cloud Storage y tus nodos de TPU. La región contiene zonas en las que están disponibles los tipos de máquina de TPU v6e (por ejemplo, us-east1, us-east5, europe-west4, asia-northeast1 o us-south1).
     • NODE_LOCATION: la zona en la que están disponibles los recursos de TPU (por ejemplo, us-east1-d).
     • CLUSTER_VERSION: la versión de GKE, que debe admitir el tipo de máquina que quieras usar. Ten en cuenta que es posible que la versión predeterminada de GKE no esté disponible para tu TPU de destino. Para ver una lista de las versiones mínimas de GKE disponibles por tipo de máquina de TPU, consulta Disponibilidad de TPU en GKE.
     • MACHINE_TYPE: el tipo de máquina v6e.
     • TPU_TYPE: prefijo que se usa para asignar nombres a los grupos de nodos (v6e).
     • TOPOLOGY: la topología de la TPU v6e.
     • WORKERS_PER_SLICE: número de nodos por grupo de nodos o segmento de TPU.

  Crear y configurar Google Cloud recursos

  Para crear los recursos necesarios, siga estas instrucciones:

  Nota: Es posible que tengas que crear una reserva de capacidad para usar algunos aceleradores. Para obtener más información sobre cómo reservar y consumir recursos reservados, consulta el artículo Consumir recursos de zona reservados.

  Crear un clúster de GKE

  1. Crea un clúster estándar de GKE regional:

     gcloud container clusters create CLUSTER_NAME \
         --project=PROJECT_ID \
         --cluster-version=CLUSTER_VERSION \
         --location=CONTROL_PLANE_LOCATION \
         --scopes=cloud-platform \
         --machine-type=n2-standard-32
     

     La creación del clúster puede tardar varios minutos.

     Sustituye CLUSTER_VERSION por la versión del clúster adecuada.

  2. Crea un grupo de nodos de TPU v6e con una topología 4x4 y cuatro nodos cada uno:

     gcloud container node-pools create multihost-np \
     --project=PROJECT_ID \
     --location=CONTROL_PLANE_LOCATION \
     --node-locations=NODE_LOCATION \
     --cluster=CLUSTER_NAME \
     --machine-type=MACHINE_TYPE \
     --num-nodes=WORKERS_PER_SLICE \
     --tpu-topology=TOPOLOGY \
     --scopes cloud-platform \
     --placement-type=COMPACT \
     --workload-metadata=GCE_METADATA
     

     Nota: El símbolo --workload-metadata=GCE_METADATA indica que el grupo de nodos tiene inhabilitada la federación de identidades de carga de trabajo para GKE. Los grupos de nodos que ejecutan Pathways con JetStream y MaxText deben inhabilitar Workload Identity Federation para GKE.

  Configurar una cuenta de servicio para acceder a objetos de Storage

  Configura una cuenta de servicio de Kubernetes para que actúe como una cuenta de servicio de IAM.

  1. Crea una cuenta de servicio de IAM para tu aplicación:

     gcloud iam service-accounts create jetstream-pathways
     

  2. Añade una vinculación de política de gestión de identidades y accesos a tu cuenta de servicio de gestión de identidades y accesos para gestionar Cloud Storage. Esto se hace para que tu cuenta de servicio de gestión de identidades y accesos pueda acceder al segmento de almacenamiento donde se guardará tu punto de control:

     gcloud projects add-iam-policy-binding ${PROJECT} \
       --member "serviceAccount:jetstream-pathways@${PROJECT}.iam.gserviceaccount.com" \
       --role roles/storage.objectUser
     
     gcloud projects add-iam-policy-binding ${PROJECT} \
       --member "serviceAccount:jetstream-pathways@${PROJECT}.iam.gserviceaccount.com" \
       --role roles/storage.insightsCollectorService
     

  3. Anota la cuenta de servicio de Kubernetes con la dirección de correo de la cuenta de servicio de gestión de identidades y accesos.

     kubectl annotate serviceaccount default \
     iam.gke.io/gcp-service-account=jetstream-pathways@${PROJECT}.iam.gserviceaccount.com
     

  Configurar Docker para autenticarse en Artifact Registry

  Configura Docker para que se autentique en Artifact Registry y pueda extraer las imágenes de Pathways incluidas en la lista de permitidas:

  gcloud auth login
  gcloud auth configure-docker
  

  Conversión de punto de control

  Para convertir un punto de control de Meta Llama 3.1-405B en un punto de control de inferencia int8 compatible con MaxText, sigue los pasos que se indican en Conversión de puntos de control con Llama3.1-405B. Tu implementación usa el punto de control con la marca load_parameters_path.

  Crear un segmento de Cloud Storage para almacenar archivos temporales de Pathways

  Crea un segmento de Cloud Storage para almacenar tus archivos temporales de Pathways, como la caché de compilación:

  export PATHWAYS_BUCKET=PATHWAYS_BUCKET
  gcloud storage buckets create gs://$PATHWAYS_BUCKET
  

  Desplegar JetStream-MaxText y Pathways

  Implementa el servidor de modelos JetStream-MaxText y Pathways.

  Conectarse al clúster de GKE

  gcloud container clusters get-credentials "${CLUSTER}" --project "${PROJECT}" --location "${ZONE}"
  

  Desplegar la API LeaderWorkerSet (LWS)

  LWS es un recurso personalizado diseñado para implementar y gestionar aplicaciones distribuidas con estado, especialmente aquellas con una arquitectura de líder-trabajador. Es especialmente adecuada para cargas de trabajo de IA o aprendizaje automático en las que un modelo grande se fragmenta y se sirve en varios dispositivos de varios nodos.

  VERSION=v0.6.1
  kubectl apply --server-side -f https://github.com/kubernetes-sigs/lws/releases/download/$VERSION/manifests.yaml
  

  Espera a que el controlador LeaderWorkerSet esté totalmente disponible:

  kubectl wait deploy/lws-controller-manager -n lws-system --for=condition=available --timeout=5m
  

  La salida debería ser similar a la siguiente:

  deployment.apps/lws-controller-manager condition met
  

  Valida que el controlador LeaderWorkerSet se esté ejecutando en el espacio de nombres lws-system:

  kubectl get pod -n lws-system
  

  La salida debería ser similar a la siguiente:

  NAME                          READY   STATUS    RESTARTS    AGE
  lws-controller-manager-abcd   1/1     Running   0           40s
  lws-controller-manager-efgh   1/1     Running   0           40s
  

  Desplegar el manifiesto de la carga de trabajo

  1. Guarda el siguiente archivo de manifiesto como jetstream-pathways-llama-3-1-405b-4x4.yaml:

     apiVersion: leaderworkerset.x-k8s.io/v1
     kind: LeaderWorkerSet
     metadata:
       name: jetstream-pathways
       annotations:
         leaderworkerset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
     spec:
       replicas: 1
       leaderWorkerTemplate:
         leaderTemplate:
           metadata:
             labels:
               app: jetstream-pathways
           spec:
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 4x4
             tolerations:
             - key: "google.com/tpu"
               operator: "Exists"
               effect: "NoSchedule"
             containers:
             - name: pathways-proxy
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
               args:
               imagePullPolicy: Always
               ports:
               - containerPort: 38681
             - name: pathways-rm
               env:
               - name: HOST_ADDRESS
                 value: "$(LWS_LEADER_ADDRESS)"
               - name: TPU_SKIP_MDS_QUERY
                 value: "true"
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
               args:
               - --server_port=38677
               - --gcs_scratch_location=PATHWAYS_BUCKET
               - --node_type=resource_manager
               - --instance_count=1
               - --instance_type=tpuv6e:4x4
               imagePullPolicy: Always
               ports:
               - containerPort: 38677
             - name: jax-tpu
               image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
               env:
               - name: LOG_LEVEL
                 value: "INFO"
               args:
               - MaxText/configs/v5e/inference/llama3_405b_v5e-64.yml
               - model_name=llama3.1-405b
               - load_parameters_path=CHECKPOINT_PATH
               - max_prefill_predict_length=1024
               - max_target_length=2048
               - async_checkpointing=false
               - steps=1
               - ici_fsdp_parallelism=1
               - ici_autoregressive_parallelism=2
               - ici_tensor_parallelism=8
               - scan_layers=false
               - weight_dtype=bfloat16
               - per_device_batch_size=6
               - enable_single_controller=true
               - quantization=int8
               - quantize_kvcache=true
               - checkpoint_is_quantized=true
               - enable_model_warmup=true
               imagePullPolicy: Always
               ports:
               - containerPort: 9000
               startupProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 1
                 initialDelaySeconds: 600
                 failureThreshold: 10000
               livenessProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 60
                 failureThreshold: 10
               readinessProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 60
                 failureThreshold: 10
             - name: jetstream-http
               image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
               imagePullPolicy: Always
               ports:
               - containerPort: 8000
         size: 5
         workerTemplate:
           spec:
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 4x4
             tolerations:
             - key: "google.com/tpu"
               operator: "Exists"
               effect: "NoSchedule"
             containers:
             - name: worker
               args:
               - --server_port=38679
               - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
               - --gcs_scratch_location=PATHWAYS_BUCKET
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
               imagePullPolicy: Always
               ports:
               - containerPort: 38679
               resources:
                 limits:
                   google.com/tpu: "4"
     --- 
     apiVersion: v1
     kind: Service
     metadata:
       name: jetstream-svc
     spec:
       selector:
         app: jetstream-pathways
       ports:
       - protocol: TCP
         name: jetstream-http
         port: 8000
         targetPort: 8000

  2. Asigna al campo load_parameters_path la ruta del punto de control que se ha generado en el proceso de conversión del punto de control.

     • En el caso de un punto de control bf16, la ruta debe ser similar a gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
     • En el caso de un punto de control int8, debería ser similar a gs://OUTPUT_BUCKET_DIRECTORY/int8.

     Asigna al campo gcs_scratch_location el valor del segmento de Pathways que has creado anteriormente.

     perl -pi -e 's|CHECKPOINT_PATH|gs://OUTPUT_BUCKET_DIRECTORY/int8|g' jetstream-pathways-llama-3-1-405b-4x4.yaml
     perl -pi -e 's|PATHWAYS_BUCKET|gs://PATHWAYS_BUCKET|g' jetstream-pathways-llama-3-1-405b-4x4.yaml
     

     Descripción de las variables y los ajustes.

     En las siguientes descripciones se explican las variables y los ajustes relevantes de JetStream, Pathways y LeaderWorkerSet.

     • Jetstream
       • Este contenedor se ha creado a partir del proyecto JetStream de software libre. También encapsula la biblioteca de cliente basada en gRPC del proxy de IFRT, que proporciona un backend de JAX que puede delegar la ejecución de programas de TPU en el servidor proxy de Pathways.
       • Para obtener más información sobre las marcas de JetStream-MaxText, consulta las descripciones de las marcas del servidor de JetStream-MaxText.
       • Para usar un punto de control cuantificado, puedes definir las marcas quantization, quantize_kvcache o checkpoint_is_quantized. Si el punto de control es bf16, puedes omitir estas marcas.
     • Rutas: para ver una descripción de los componentes de las rutas relevantes, consulta Componentes de las rutas.
     • LeaderWorkerSet: para obtener más información sobre los ajustes de LeaderWorkerSet, consulta la definición de la API.

  Aplica el manifiesto de Deployment

  Aplica el manifiesto para implementar el servidor:

  kubectl apply -f jetstream-pathways-llama-3-1-405b-4x4.yaml
  

  El servidor del modelo debería iniciarse.

  Verificar el inicio del servidor del modelo

  Un modelo 405B puede tardar entre 10 y 20 minutos en restaurar el punto de control. También puedes esperar más tiempo durante el calentamiento del modelo si has habilitado la marca enable_model_warmup.

  kubectl logs -f jetstream-pathways-0 -c jax-tpu
  

  El resultado debería ser similar al siguiente:

  2025-03-02 02:15:07,682 - JetstreamLogger - INFO - Initializing the driver with 1 prefill engines and 1 generate engines in interleaved mode
  2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up prefill thread 0.
  2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up transfer thread 0.
  2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up generate thread 0.
  2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up detokenize thread 0.
  2025-03-02 02:15:07,685 - JetstreamLogger - INFO - Driver initialized.
  ...
  ...
  ...
  INFO:     Started server process [7]
  INFO:     Waiting for application startup.
  INFO:     Application startup complete.
  INFO:     Uvicorn running on http://0.0.0.0:9999 (Press CTRL+C to quit)
  

  Servir Llama 3.1-405b

  Para publicar el modelo Llama 3.1-405b, configura el reenvío de puertos:

  kubectl port-forward svc/jetstream-svc 8000:8000
  

  El reenvío de puertos te permite acceder al servicio desde fuera del clúster. Puedes acceder a la implementación de JetStream-Pathways a través del servicio ClusterIP de GKE. Solo se puede acceder a los servicios ClusterIP desde el clúster.

  Interactuar con el modelo

  En un terminal nuevo, ejecuta el siguiente comando:

  curl --request POST \
  --header "Content-type: application/json" \
  -s \
  localhost:8000/generate \
  --data \
  '{
      "prompt": "What are the top 5 programming languages",
      "max_tokens": 200
  }'
  

  La solicitud inicial puede tardar varios segundos en completarse debido al calentamiento del modelo. La salida debería ser similar a la siguiente:

  {
      "response": " for web development?\nThe top 5 programming languages for web development are:\n1. **JavaScript**: JavaScript is the most popular language for web development, used by over 90% of websites for client-side scripting. It's also popular for server-side programming with technologies like Node.js.\n2. **HTML/CSS**: HTML (Hypertext Markup Language) and CSS (Cascading Style Sheets) are not programming languages, but are essential for building websites. HTML is used for structuring content, while CSS is used for styling and layout.\n3. **Python**: Python is a popular language for web development, especially with frameworks like Django and Flask. It's known for its simplicity, flexibility, and large community of developers.\n4. **Java**: Java is a popular language for building enterprise-level web applications, especially with frameworks like Spring and Hibernate. It's known for its platform independence, strong security features, and large community of developers.\n5. **PHP**: PHP is a mature language for web"
  }
  

  Has hecho lo siguiente correctamente:

  1. Ha desplegado el servidor de modelos JetStream con MaxText y Pathways en GKE mediante TPUs.
  2. Se ha creado un punto de control de Llama 3.1-405B int8 a las gs://BUCKET_NAME.
  3. Se ha ofrecido el modelo y se ha interactuado con él.

  Publicación desagregada

  El servicio desagregado es una técnica para servir LLMs que divide las fases de prellenado y decodificación en hosts diferentes. Este enfoque optimiza el uso de los recursos, lo que puede mejorar el rendimiento y la latencia.

  • Rellenado automático: un paso hacia adelante en la petición de entrada para inicializar la caché de pares clave-valor.

  • Decodificación: procedimiento que genera tokens de salida de forma incremental, un token por paso y un valor de caché de clave-valor por iteración.

  1. Define las variables de entorno predeterminadas:

     export NODE_POOL_NAME=dis-v6e-8
     export NODE_POOL_SIZE=2
     export MACHINE_TYPE=ct6e-standard-4t
     export TOPOLOGY=2x4
     export WORKERS_PER_SLICE=2
     

  2. Crea dos grupos de nodos que usen nodos v6e-8:

     for i in $(seq 1 NODE_POOL_SIZE); do
       gcloud container node-pools create NODE_POOL_NAME-${i}-np \
       --project=PROJECT \
       --cluster=CLUSTER_NAME \
       --location=CONTROL_PLANE_LOCATION \
       --node-locations=NODE_LOCATION \
       --machine-type=MACHINE_TYPE \
       --num-nodes=WORKERS_PER_SLICE \
       --tpu-topology=TOPOLOGY \
       --scopes=cloud-platform \
       --workload-metadata=GCE_METADATA
     done
     

  Conversión de punto de control

  Para convertir un punto de control de Meta Llama 2-70B en un punto de control de inferencia int8 compatible con MaxText, sigue los pasos que se indican en Conversión de puntos de control con Llama2-70B. Selecciona Llama2-70B como modelo cuando aceptes los Términos y Condiciones de Meta. Tu implementación usa el punto de control con la marca load_parameters_path.

  Sustituye los siguientes parámetros en tu archivo checkpoint-job.yaml:

  - --meta_url=META_URL
  - --model_name=llama-2
  - --model_path=Llama-2-70b-chat
  - --output_directory=gs://BUCKET_NAME/maxtext/llama-2-70b
  

  El punto de control se usará en tu implementación con la marca load_parameters_path.

  Desplegar JetStream Pathways con el servicio desagregado

  1. Guarda el siguiente archivo de manifiesto como jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml:

     apiVersion: leaderworkerset.x-k8s.io/v1
     kind: LeaderWorkerSet
     metadata:
       name: jetstream-pathways
       annotations:
         leaderworkerset.sigs.k8s.io/subgroup-exclusive-topology: cloud.google.com/gke-nodepool
     spec:
       replicas: 1
       leaderWorkerTemplate:
         subGroupPolicy:
           subGroupSize: 2
         leaderTemplate:
           metadata:
             labels:
               app: jetstream-pathways
           spec:
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 2x4
             tolerations:
             - key: "google.com/tpu"
               operator: "Exists"
               effect: "NoSchedule"
             containers:
             - name: pathways-proxy
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
               args:
               - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
               - --server_port=38681
               - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
               - --xla_jf_auto_cross_replica_sharding=false
               - --xla_tpu_enable_windowed_einsum_for_reduce_scatter=false
               - --xla_tpu_enable_windowed_einsum_for_all_gather=false
               - --xla_tpu_prefer_latch_optimized_rhs_layouts=true
               - --xla_tpu_enable_experimental_fusion_cost_model=false
               - --xla_tpu_dot_dot_fusion_duplicated=false
               - --xla_tpu_dot_dot_fusion=true
               - --xla_jf_conv_input_fusion=true
               - --xla_jf_conv_output_fusion=true
               - --xla_tpu_rwb_fusion=false
               - --xla_tpu_copy_fusion_pad_unpad_ratio=0
               - --xla_tpu_licm_size_inflation_ratio=1
               - --xla_tpu_copy_elision_analysis_allowance=150000
               - --xla_tpu_copy_insertion_use_region_analysis_limit=10000
               - --xla_tpu_order_dot_after_layout=true
               - --xla_jf_rematerialization_percent_shared_memory_limit=100
               - --xla_tpu_use_repeated_instance_for_preferred_prefetch_time=true
               - --xla_tpu_enforce_prefetch_fifo_order=false
               - --xla_tpu_prefetch_interval_picker_size_override=6000000
               - --xla_tpu_async_copy_bandwidth_scaling_factor=1
               - --xla_tpu_nd_short_transfer_max_chunks=-1
               - --xla_tpu_enable_aggressive_broadcast_priority_update=true
               - --xla_tpu_alternate_memory_benefit_scaling_factor_for_large_buffers=SQRT
               - --xla_tpu_memory_bound_loop_optimizer_options=enabled:true
               - --xla_tpu_enable_copy_fusion=true
               - --xla_tpu_enable_cross_program_prefetch_freeing=false
               - --xla_tpu_enable_dot_strength_reduction=true
               - --xla_tpu_layout_use_dot_grouping=false
               - --xla_tpu_msa_inefficient_use_to_copy_ratio=0.5
               - --xla_tpu_reduce_loop_fusion_dup_with_unfusable_user=false
               - --xla_tpu_vector_load_fusion_window=1024
               - --xla_tpu_vector_store_fusion_window=256
               - --xla_jf_conv_reshape_fusion=false
               - --xla_tpu_input_conv_multi_users=false
               - --xla_tpu_enable_multi_level_input_dot_dot_fusion=false
               - --xla_tpu_enable_multi_level_output_dot_dot_fusion=false
               - --xla_tpu_dot_dot_fusion_separable_convs_only=false
               - --xla_tpu_enable_multi_level_nested_loop_fusion=true
               - --xla_tpu_nested_dot_fusion=true
               - --xla_tpu_enable_multi_level_nested_dot_fusion=false
               - --xla_jf_enable_multi_output_fusion=true
               - --xla_tpu_use_lp_llo_scheduler_for_dot_dot_fusions=false
               - --xla_tpu_enable_flash_attention=true
               imagePullPolicy: Always
               ports:
               - containerPort: 38681
             - name: pathways-rm
               env:       
               - name: HOST_ADDRESS
                 value: "$(LWS_LEADER_ADDRESS)"
               - name: TPU_SKIP_MDS_QUERY
                 value: "true"
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
               args:
               - --server_port=38677
               - --gcs_scratch_location=PATHWAYS_BUCKET
               - --node_type=resource_manager
               - --instance_count=2
               - --instance_type=tpuv6e:2x4
               imagePullPolicy: Always
               ports:
               - containerPort: 38677
             - name: jax-tpu
               image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
               args:
               - MaxText/configs/base.yml
               - tokenizer_path=assets/tokenizer.llama2
               - load_parameters_path=CHECKPOINT_PATH
               - max_prefill_predict_length=1024
               - max_target_length=2048
               - model_name=llama2-70b
               - ici_fsdp_parallelism=1
               - ici_autoregressive_parallelism=1
               - ici_tensor_parallelism=-1
               - scan_layers=false
               - weight_dtype=bfloat16
               - per_device_batch_size=27
               - checkpoint_is_quantized=true 
               - quantization=int8
               - quantize_kvcache=true
               - compute_axis_order=0,2,1,3
               - ar_cache_axis_order=0,2,1,3
               - stack_prefill_result_cache=True
               - inference_server=ExperimentalMaxtextDisaggregatedServer_8
               - inference_benchmark_test=True
               - enable_model_warmup=True
               env:
               - name: LOG_LEVEL
                 value: "INFO"
               imagePullPolicy: Always
               securityContext:
                 capabilities:
                   add: ["SYS_PTRACE", "NET_ADMIN", "SYS_TIME"]
               ports: 
               - containerPort: 9000
               startupProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 1
                 initialDelaySeconds: 240
                 failureThreshold: 10000
               livenessProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 60
                 failureThreshold: 100
               readinessProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 60
                 failureThreshold: 100
             - name: jetstream-http
               image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
               imagePullPolicy: Always
               ports:
               - containerPort: 8000
         size: 5
         workerTemplate:
           spec:
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 2x4
             containers:
             - name: worker
               args:
               - --server_port=38679
               - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
               - --gcs_scratch_location=PATHWAYS_BUCKET
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
               imagePullPolicy: Always
               ports:
               - containerPort: 38679
               resources:
                 limits:
                   google.com/tpu: "4"
     --- 
     apiVersion: v1
     kind: Service
     metadata:
       name: jetstream-svc
     spec:
       selector:
         app: jetstream-pathways
       ports:
       - protocol: TCP
         name: jetstream-http
         port: 8000
         targetPort: 8000

  2. Asigna al campo load_parameters_path la ruta del punto de control que se ha generado en el proceso de conversión del punto de control.

     • En el caso de un punto de control bf16, la ruta debe ser similar a gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
     • En el caso de un punto de control int8, debería ser similar a gs://OUTPUT_BUCKET_DIRECTORY/int8.

     Asigna al campo gcs_scratch_location el valor del segmento de Pathways que has creado anteriormente.

     perl -pi -e 's|CHECKPOINT_PATH|BUCKET_NAME/maxtext/llama-2-70b/int8|g' jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
     perl -pi -e 's|PATHWAYS_BUCKET|gs://PATHWAYS_BUCKET|g' jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
     

  3. Aplica el archivo de manifiesto:

     kubectl apply -f jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
     

     El servidor del modelo puede tardar un tiempo en restaurar el punto de control, en función de su tamaño. Un modelo de 70 mil millones puede tardar unos 8 minutos en restaurar el punto de control, incluido el calentamiento del modelo. Puedes consultar los registros para determinar el punto de preparación verificando el inicio del servidor del modelo y servir el modelo configurando el reenvío de puertos para poder interactuar con él.

  Has hecho lo siguiente correctamente:

  1. Desplegó el servidor de modelos JetStream con MaxText y Pathways en GKE mediante TPUs y el servicio desagregado.
  2. Se ha creado un punto de control de Llama 2-70B int8 a las gs://BUCKET_NAME.
  3. Se ha ofrecido el modelo y se ha interactuado con él.

  Solucionar problemas

  • Si recibes el mensaje Empty reply from server, es posible que el contenedor no haya terminado de descargar los datos del modelo. Vuelve a consultar los registros del pod para ver el mensaje Connected, que indica que el modelo está listo para usarse.
  • Si ves el mensaje Connection refused, comprueba que el reenvío de puertos esté activo.

  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 los siguientes comandos y sigue las indicaciones:

  gcloud container clusters delete CLUSTER_NAME --location=CONTROL_PLANE_LOCATION
  
  gcloud iam service-accounts delete jetstream-pathways@PROJECT_ID.iam.gserviceaccount.com
  
  gcloud storage rm --recursive gs://BUCKET_NAME
  

  Siguientes pasos

  • 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.
  • Consulta el repositorio de GitHub de JetStream.
  • Consulta Vertex AI Model Garden.