Entrega LLMs con TPU de varios hosts en GKE con JetStream y Pathways

Estándar

En esta guía, se muestra cómo entregar modelos de lenguaje grandes (LLMs) de vanguardia, como Llama 3.1 405B, en Google Kubernetes Engine (GKE) con 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 de LeaderWorkerSet [LWS]) para implementar y entregar cargas de trabajo de IA/AA en GKE, aprovechando el control detallado, la escalabilidad, la resiliencia, la portabilidad y la rentabilidad de GKE.

Fondo

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

Con la entrega de un LLM con TPU en GKE con JetStream, MaxText y Pathways, puedes compilar una solución de entrega sólida y lista para la producción con todos los beneficios de Kubernetes administrado, incluida la rentabilidad, escalabilidad y disponibilidad mayor. En esta sección, se describen las tecnologías clave que se usan en este instructivo.

Acerca de las 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 TensorFlow, PyTorch y JAX.

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 Llama 3.1-405B. GKE implementa el modelo en los nodos TPU v6e de varios hosts con topologías de TPU configuradas según los requisitos del modelo para entregar mensajes con baja latencia.

Rutas de aprendizaje en Cloud

Pathways es una capa de organización a gran escala para aceleradores. Pathways está diseñado explícitamente para permitir la exploración de nuevos sistemas e ideas de investigación de AA, al mismo tiempo que mantiene el rendimiento de vanguardia para los modelos actuales. Pathways permite que un solo proceso cliente de JAX coordine la computación en una o más porciones de TPU grandes, lo que optimiza los cálculos de AA que abarcan cientos o miles de chips de TPU.

JetStream

JetStream es un framework de entrega de inferencia de código abierto que desarrolla Google. JetStream permite la inferencia de alto rendimiento, alta capacidad de procesamiento y con optimización de memoria en TPU y GPU. JetStream proporciona optimizaciones de rendimiento avanzadas, incluidas técnicas de procesamiento por lotes, optimizaciones de la caché de KV y de cuantización continuas, para facilitar la implementación de LLM. JetStream permite que PyTorch/XLA y JAX TPU entreguen un rendimiento óptimo.

MaxText

MaxText es una implementación de LLM JAX escalable y adaptable, compilada en bibliotecas JAX de código abierto como Flax, Orbax y Optax. La implementación de LLM solo de decodificador de MaxText se escribe en Python. Aprovecha en gran medida el compilador XLA para lograr un alto rendimiento sin necesidad de compilar 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 de MaxText.

Llama 3.1 405b

Llama 3.1 405B es un modelo de lenguaje grande de Meta diseñado para una variedad de tareas de procesamiento de lenguaje natural, como la generación de texto, la traducción y la respuesta de preguntas. GKE ofrece la infraestructura necesaria para admitir las necesidades de entrenamiento y entrega 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 instructivo. La arquitectura incluye un clúster de GKE Standard que aprovisiona TPU y aloja componentes de JetStream y Pathways para implementar y entregar 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 porción de TPU de varios hosts que aloja los componentes de implementación y de Pathways de JetStream.
• El Pathways resource manager administra los recursos del acelerador y coordina la asignación de aceleradores para los trabajos del usuario.
• El objeto Pathways client se coordina con el objeto 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 de vuelta a tu carga de trabajo a través del servidor proxy de IFRT.
• IFRT proxy client implementa la API de 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 de Service distribuye el tráfico entrante a todas las réplicas de JetStream HTTP.
• JetStream HTTP es un servidor HTTP que acepta solicitudes como wrapper del formato requerido de JetStream y las envía al cliente GRPC de JetStream.

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

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

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

• Asegúrate de tener los siguientes roles en el proyecto: roles/container.admin, roles/iam.serviceAccountAdmin, roles/resourcemanager.projectIamAdmin

  Verifica los roles

  1. En la consola de Google Cloud , dirígete a la página IAM.

     Ir a IAM
  2. Selecciona el proyecto.

  3. En la columna Principal, busca todas las filas que te identifiquen a ti o a un grupo en el que se te incluya. Para saber en qué grupos estás incluido, comunícate con tu administrador.

  4. Para todas las filas en las que se te especifique o se te incluya, verifica la columna Rol para ver si la lista de roles incluye los roles necesarios.

  Otorga los roles

  1. En la consola de Google Cloud , dirígete a la página IAM.

     Ir a IAM
  2. Selecciona el proyecto.
  3. Haz clic en Otorgar acceso.

  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. Haz clic en Seleccionar un rol y, luego, busca el rol.
  6. Para otorgar roles adicionales, haz clic en Agregar otro rol y agrega uno más.
  7. Haz clic en Guardar.

• Asegúrate de tener suficiente cuota para dieciséis chips TPU v6e PodSlice Lite. En este instructivo, usarás instancias bajo demanda.

Obtén acceso al modelo

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

1. Firma el acuerdo 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 obtener una lista de los modelos compatibles 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.

Prepare el entorno

En este instructivo, usarás Cloud Shell para administrar recursos alojados enGoogle Cloud. Cloud Shell tiene preinstalado el software que necesitas para este instructivo, incluidos kubectl y la CLI de gcloud.

Para configurar tu entorno con Cloud Shell, sigue estos pasos:

1. En la Google Cloud consola, haz clic en Activar Cloud Shell en la Google Cloud consola 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
   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
   

   Reemplaza los siguientes valores:

   • PROJECT_ID: Es el Google Cloud ID del proyecto.
   • CLUSTER_NAME: Es el nombre del clúster de GKE.
   • BUCKET_NAME: Es el nombre de tu bucket de Cloud Storage. No es necesario especificar el prefijo gs://.
   • CONTROL_PLANE_LOCATION: Es la región de Compute Engine 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 v6e están disponibles (por ejemplo, us-east1, us-east5, europe-west4, asia-northeast1 o us-south1).
   • NODE_LOCATION: Es la zona en la que los recursos de TPU están disponibles (por ejemplo, us-east1-d).
   • CLUSTER_VERSION: la versión de GKE, que debe ser compatible con el tipo de máquina que deseas usar. Ten en cuenta que es posible que la versión predeterminada de GKE no tenga disponibilidad para la TPU de destino. Para obtener 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: Es el tipo de máquina v6e.
   • TPU_TYPE: Es un prefijo que se usa para nombrar grupos de nodos (v6e).
   • TOPOLOGY: Es la topología de la TPU v6e.
   • WORKERS_PER_SLICE: Es la cantidad de nodos por grupo de nodos o porción de TPU.

Crea y configura recursos de Google Cloud

Para crear los recursos necesarios, sigue estas instrucciones:

Nota: Es posible que debas crear una reserva de capacidad para usar algunos aceleradores. Para obtener más información sobre cómo reservar y consumir recursos reservados, consulta Consume recursos zonales reservados.

Crea un clúster de GKE

1. Crea un clúster de GKE Standard 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 tomar varios minutos.

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

2. Crea un grupo de nodos 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 para cargas de trabajo para GKE. Los grupos de nodos que ejecutan Pathways con JetStream y MaxText deben inhabilitar la federación de identidades para cargas de trabajo para GKE.

Configura 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. Agrega una vinculación de política de IAM para que tu cuenta de servicio de IAM administre Cloud Storage. Esto permite que tu cuenta de servicio de IAM acceda al bucket de almacenamiento en el que se guardará tu punto de control:

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

3. Anota la cuenta de servicio de Kubernetes con la dirección de correo electrónico de la cuenta de servicio de IAM.

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

Configura Docker para que se autentique 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 entidades 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, completa 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.

Crea un bucket de Cloud Storage para almacenar archivos temporales de Pathways

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

Implementa JetStream-MaxText y Pathways

Implementa el servidor de modelos de JetStream-MaxText y Pathways.

Conéctate al clúster de GKE

gcloud container clusters get-credentials "${CLUSTER_NAME}" --project "${PROJECT_ID}"  \
    --location "${CONTROL_PLANE_LOCATION}"

Implementa la API de LeaderWorkerSet (LWS)

LWS es un recurso personalizado diseñado para implementar y administrar aplicaciones distribuidas con estado, en particular, aquellas con una arquitectura de líder y trabajador. Es especialmente adecuado para cargas de trabajo de IA/AA en las que un modelo grande se fragmenta y se entrega en varios dispositivos en 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é completamente disponible:

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

El resultado debería ser similar al siguiente ejemplo:

deployment.apps/lws-controller-manager condition met

Valida que el controlador LeaderWorkerSet se ejecute en el espacio de nombres lws-system:

kubectl get pod -n lws-system

El resultado debería ser similar al siguiente ejemplo:

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

Implementa el manifiesto de la carga de trabajo

1. Guarda el siguiente 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. Establece el valor del campo load_parameters_path en la ruta de acceso del punto de control que se produjo en el proceso de conversión del punto de control.

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

   Establece el valor del campo gcs_scratch_location en el bucket de Pathways que creaste anteriormente.

   perl -pi -e "s|CHECKPOINT_PATH|gs://${BUCKET_NAME}/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
   

   Es la descripción de las variables y la configuración.

   En las siguientes descripciones, se explican las variables y la configuración pertinentes para JetStream, Pathways y LeaderWorkerSet.

   • Jetstream
     • Este contenedor se compila a partir del proyecto de OSS JetStream. También encapsula la biblioteca 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 de 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 establecer las marcas quantization, quantize_kvcache o checkpoint_is_quantized. Si el punto de control es bf16, puedes omitir estas marcas.
   • Rutas de aprendizaje: Para obtener una descripción de los componentes relevantes de las rutas de aprendizaje, consulta Componentes de las rutas de aprendizaje.
   • LeaderWorkerSet: Para obtener más información sobre la configuración 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

Se debería iniciar el servidor de modelos.

Verifica el inicio del servidor de modelos

Un modelo 405B puede tardar entre 10 y 20 minutos en restablecer el punto de control. También es posible que debas esperar un tiempo adicional durante el calentamiento del modelo si habilitaste la marca enable_model_warmup.

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

El resultado es similar a lo 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)

Cómo publicar Llama 3.1-405b

Para entregar el modelo Llama 3.1-405b, configura la redirección de puertos:

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

El reenvío de puertos te permite acceder al Service 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 de ClusterIP desde el clúster.

Interactúa con el modelo

En una nueva terminal, 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 a la preparación del modelo. El resultado debería ser similar al siguiente ejemplo:

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

Completaste correctamente lo siguiente:

1. Se implementó el servidor de modelos de JetStream con MaxText y Pathways en GKE con TPU.
2. Se creó un punto de control int8 de Llama 3.1-405B en gs://BUCKET_NAME.
3. Se entregó el modelo y se interactuó con él.

Publicación desagregada

La entrega desagregada es una técnica para entregar LLMs mientras se dividen las etapas de precompletado y decodificación en diferentes hosts. Este enfoque optimiza el uso de los recursos, lo que puede mejorar la capacidad de procesamiento y la latencia.

• Prefill: Es un pase hacia adelante en la instrucción de entrada para inicializar la caché de pares clave-valor.

• Decodificación: Es un procedimiento que genera tokens de salida de forma incremental, un token por paso y un valor de KV-cache por iteración.

1. Configura 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_ID} \
     --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, completa los pasos que se indican en Conversión de puntos de control con Llama2-70B. Selecciona Llama2-70B como tu modelo cuando confirmes los Términos y Condiciones de Meta. Tu implementación usa el punto de control con la marca load_parameters_path.

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

Implementa JetStream Pathways con la entrega desagregada

1. Guarda el siguiente 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. Establece el valor del campo load_parameters_path en la ruta de acceso del punto de control que se produjo en el proceso de conversión del punto de control.

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

   Establece el valor del campo gcs_scratch_location en el bucket de Pathways que creaste anteriormente.

   perl -pi -e "s|CHECKPOINT_PATH|gs://${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 manifiesto

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

   El servidor del modelo puede tardar un tiempo en restablecer el punto de control, según su tamaño. Un modelo de 70B puede tardar unos 8 minutos en restablecer el punto de control, incluido el calentamiento del modelo. Puedes observar los registros para determinar el punto de preparación verificando el inicio del servidor del modelo y publicando el modelo configurando el reenvío de puertos para que puedas interactuar con el modelo.

Completaste correctamente lo siguiente:

1. Se implementó el servidor de modelos de JetStream con MaxText y Pathways en GKE con TPU y entrega desagregada.
2. Se creó un punto de control de Llama 2-70B int8 en gs://BUCKET_NAME.
3. Se entregó el modelo y se interactuó con él.

Soluciona 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 verificar los registros del Pod en busca del mensaje Connected, que indica que el modelo está listo para entregar.
• Si ves el mensaje Connection refused, verifica que tu redirección de puertos esté activa.

Realiza una limpieza

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

¿Qué sigue?

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