Entrega LLM 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.

Importante: Para obtener acceso a Pathways, comunícate con tu representante de cuenta de Google Cloud.

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 en Cloud para ejecutar inferencias a gran escala con varios hosts en GKE en varios nodos de TPU interconectados. En esta guía, aprenderás a aprovisionar un clúster de GKE con las porciones de TPU de varios hosts y a 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 tiempo que conserva 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.
• 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.
• El 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

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.

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.

Enable the API

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

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.

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 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. En la lista Seleccionar un rol, elige uno.
  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.
  • Asegúrate de que tu proyecto Google Cloud esté en la lista de entidades permitidas para Pathways.

  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: 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} \
       --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 electrónico de la cuenta de servicio de IAM.

     kubectl annotate serviceaccount default \
     iam.gke.io/gcp-service-account=jetstream-pathways@${PROJECT}.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}" --project "${PROJECT}" --location "${ZONE}"
  

  Implementa la API de LeaderWorkerSet (LWS)

  LWS es un recurso personalizado diseñado para implementar y administrar aplicaciones distribuidas y con estado, en particular, aquellas con una arquitectura de líder-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://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
     

     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 más tiempo 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 este:

  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 terminal nueva, 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 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 \
       --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|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 int8 de Llama 2-70B 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.

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