Entrena un LLM con JAX, Ray Train y TPU Trillium en GKE

Autopilot Estándar

En este instructivo, se muestra cómo entrenar el modelo de lenguaje grande (LLM) Llama 3 8B en Google Kubernetes Engine (GKE) con MaxText, Ray Train y TPU.

En este instructivo, se proporciona un recorrido completo de extremo a extremo, desde la configuración de la infraestructura de nube necesaria hasta el envío y la ejecución correcta de la carga de trabajo de entrenamiento en TPU de varios hosts.

Este instructivo está dirigido a los administradores y operadores de plataformas, y a los especialistas en IA y datos que quieren aprender a entrenar modelos grandes en una porción de TPU distribuida de varios hosts.

Fondo

La combinación de GKE, KubeRay, MaxText y TPU proporciona una plataforma potente y escalable para el entrenamiento de modelos a gran escala. En esta sección, se describen las tecnologías clave que se usan en esta guía:

JAX

JAX es una biblioteca de Python para la transformación de programas y el procesamiento de arrays orientados a aceleradores. Está diseñada para la computación numérica de alto rendimiento y el aprendizaje automático a gran escala.

JAX proporciona un sistema extensible para transformar funciones numéricas como jax.grad, jax.jit y jax.vmap, utilizando el compilador XLA para crear código altamente optimizado que se escale de manera eficiente en aceleradores como GPU y TPU. La potencia principal de JAX radica en su capacidad de composición, que permite a los usuarios combinar estas transformaciones para compilar programas numéricos complejos y de alto rendimiento para la ejecución distribuida.

MaxText

MaxText es un modelo de lenguaje grande (LLM) de código abierto y alto rendimiento diseñado para la escalabilidad y la personalización. MaxText se basa en JAX y está optimizado para ejecutarse de manera eficiente en Cloud TPU y GPU.

TPU

Las unidades de procesamiento tensorial (TPU) son aceleradores diseñados de forma personalizada creados por Google para optimizar las cargas de trabajo de aprendizaje automático. A diferencia de las CPU de uso general o las GPU de procesamiento paralelo, las TPU están altamente especializadas para los cálculos masivos de matrices y tensores en la base del aprendizaje profundo, lo que las hace eficientes en esta tarea específica. La principal ventaja de las TPU es el rendimiento a escala.

En este instructivo, se usa TPU Trillium, que es la sexta generación de TPU. Para obtener más información, consulta Beneficios de usar TPU Trillium.

KubeRay

KubeRay es un operador de Kubernetes que proporciona una forma unificada de implementar, administrar y supervisar aplicaciones de Ray en Kubernetes. El operador de KubeRay se instala y administra a través del complemento Ray en GKE, que es la forma recomendada de implementar y administrar clústeres de Ray en GKE.

Objetivos

En este instructivo, se muestra cómo realizar lo siguiente:

1. Configura un clúster de GKE con un grupo de nodo TPU de varios hosts.
2. Configura KubeRay para administrar el entorno de entrenamiento distribuido.
3. Compila una imagen de Docker personalizada que contenga dependencias de MaxText, Ray y JAX.
4. Crea una secuencia de comandos de entrenamiento de Python que use JaxTrainer de Ray Train para organizar el bucle de entrenamiento de MaxText en la porción de TPU.
5. Define un RayCluster recurso personalizado para aprovisionar los nodos principales y de trabajador con los recursos de TPU necesarios.
6. Envía el trabajo de entrenamiento al RayCluster y supervisa su progreso.
7. Usa Cloud Storage para almacenar los puntos de control del modelo.

Antes de comenzar

Accede a tu Google Cloud cuenta. 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.

Instala Google Cloud CLI.

Nota: Si ya instalaste gcloud CLI, asegúrate de tener la versión más reciente ejecutando gcloud components update.

Si usas un proveedor de identidad externo (IdP), primero debes acceder a la gcloud CLI con tu identidad federada.

Para inicializar gcloud CLI, ejecuta el siguiente comando:

gcloud init

Crea o selecciona un Google Cloud proyecto.

Roles necesarios para seleccionar o crear un proyecto

• Selecciona un proyecto: Para seleccionar un proyecto, no se requiere un rol de IAM específico. Puedes seleccionar cualquier proyecto en el que se te haya otorgado un rol.
• Crea un proyecto: Para crear un proyecto, necesitas el rol de creador de proyectos (roles/resourcemanager.projectCreator), que contiene el resourcemanager.projects.create permiso. Obtén más información para otorgar roles.

Nota: Si no planeas conservar los recursos creados durante este procedimiento, crea un proyecto en lugar de seleccionar un proyecto existente. Cuando termines, puedes borrar el proyecto y quitar todos los recursos asociados con él.

• Crea un proyecto de: Google Cloud

  gcloud projects create PROJECT_ID

  Reemplaza PROJECT_ID por un nombre para el Google Cloud proyecto de que estás creando.

• Selecciona el Google Cloud proyecto de que creaste:

  gcloud config set project PROJECT_ID

  Reemplaza PROJECT_ID por el nombre de tu Google Cloud proyecto de.

Verifica que la facturación esté habilitada para tu Google Cloud proyecto.

Habilita las APIs necesarias:

Roles necesarios para habilitar las APIs

Para habilitar las APIs, necesitas el rol de IAM de administrador de Service Usage (roles/serviceusage.serviceUsageAdmin), que contiene el serviceusage.services.enable permiso. Obtén más información para otorgar roles.

gcloud services enable container.googleapis.com cloudbuild.googleapis.com

Instala Google Cloud CLI.

Nota: Si ya instalaste gcloud CLI, asegúrate de tener la versión más reciente ejecutando gcloud components update.

Si usas un proveedor de identidad externo (IdP), primero debes acceder a la gcloud CLI con tu identidad federada.

Para inicializar gcloud CLI, ejecuta el siguiente comando:

gcloud init

Crea o selecciona un Google Cloud proyecto.

Roles necesarios para seleccionar o crear un proyecto

• Selecciona un proyecto: Para seleccionar un proyecto, no se requiere un rol de IAM específico. Puedes seleccionar cualquier proyecto en el que se te haya otorgado un rol.
• Crea un proyecto: Para crear un proyecto, necesitas el rol de creador de proyectos (roles/resourcemanager.projectCreator), que contiene el resourcemanager.projects.create permiso. Obtén más información para otorgar roles.

Nota: Si no planeas conservar los recursos creados durante este procedimiento, crea un proyecto en lugar de seleccionar un proyecto existente. Cuando termines, puedes borrar el proyecto y quitar todos los recursos asociados con él.

• Crea un proyecto de: Google Cloud

  gcloud projects create PROJECT_ID

  Reemplaza PROJECT_ID por un nombre para el Google Cloud proyecto de que estás creando.

• Selecciona el Google Cloud proyecto de que creaste:

  gcloud config set project PROJECT_ID

  Reemplaza PROJECT_ID por el nombre de tu Google Cloud proyecto de.

Verifica que la facturación esté habilitada para tu Google Cloud proyecto.

Habilita las APIs necesarias:

Roles necesarios para habilitar las APIs

Para habilitar las APIs, necesitas el rol de IAM de administrador de Service Usage (roles/serviceusage.serviceUsageAdmin), que contiene el serviceusage.services.enable permiso. Obtén más información para otorgar roles.

gcloud services enable container.googleapis.com cloudbuild.googleapis.com

• Otorga roles a tu cuenta de usuario. Ejecuta el siguiente comando una vez para cada uno de los siguientes roles de IAM: roles/container.admin, roles/iam.serviceAccountAdmin, roles/cloudbuild.builds.editor

  gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

  Reemplaza lo siguiente:

  • PROJECT_ID: ID del proyecto
  • USER_IDENTIFIER: El identificador de tu cuenta de usuario Por ejemplo, [e-mail].
  • ROLE: El rol de IAM que otorgas a tu cuenta de usuario

• Debido a que en este instructivo se utiliza TPU Trillium (v6e), selecciona una región o zona con disponibilidad. Para obtener más información, consulta Cuotas de Cloud TPU.

Prepara el entorno

En este instructivo, usarás Cloud Shell. Cloud Shell viene preinstalado con las herramientas de línea de comandos de gcloud, helm y kubectl que se usan en este instructivo.

1. Ve a la Google Cloud consola.

2. En la parte superior de la Google Cloud ventana de la consola de, haz clic en el botón Activar Cloud Shell .

   Se abrirá una sesión de Cloud Shell en un marco nuevo en la Google Cloud consola, y se mostrará una ventana de línea de comandos.

3. Crea y activa un entorno virtual de Python:

   python3 -m venv ray-env
   source ray-env/bin/activate
   

4. Instala la CLI de Ray y otras dependencias:

   pip install "ray[default]==2.49.1"
   

5. Configura las siguientes variables de entorno:

   export PROJECT_ID=$(gcloud config get project)
   export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
   export GS_BUCKET=GS_BUCKET
   export KSA_NAME=KSA_NAME
   export NAMESPACE=default
   export CLUSTER_NAME=CLUSTER_NAME
   export REGION=REGION
   export ZONE=ZONE
   export ARTIFACT_REGISTRY=ARTIFACT_REGISTRY
   

   Reemplaza lo siguiente:

   • GS_BUCKET: el nombre del bucket de Cloud Storage.
   • KSA_NAME: el nombre de la cuenta de servicio de Kubernetes.
   • CLUSTER_NAME: es el nombre del clúster nuevo.
   • REGION: la región en la que está disponible tu capacidad de TPU Trillium.
   • ZONE: la zona en la que está disponible tu capacidad de TPU Trillium. Para obtener más información, consulta la disponibilidad de TPU en GKE.
   • ARTIFACT_REGISTRY: el nombre del repositorio de Artifact Registry.

Crea un clúster de GKE

Puedes configurar KubeRay en TPU en un clúster de GKE Autopilot o Estándar. Te recomendamos que uses un clúster en modo Autopilot para una experiencia de Kubernetes completamente administrada. Para elegir el modo de operación de GKE que se adapte mejor a tus cargas de trabajo, consulta Acerca de los modos de operación de GKE.

Autopilot

1. En Cloud Shell, ejecuta el siguiente comando:

   gcloud container clusters create-auto $CLUSTER_NAME \
       --enable-ray-operator \
       --machine-type=n1-standard-16 \
       --location=$REGION
   

2. Para comunicarte con tu clúster, configura kubectl :

   gcloud container clusters get-credentials CLUSTER_NAME \
       --location=$ZONE
   

Estándar

1. En Cloud Shell, crea un clúster estándar que habilite el complemento del operador de Ray ejecutando el siguiente comando:

   gcloud container clusters create $CLUSTER_NAME \
       --addons=RayOperator \
       --addons GcsFuseCsiDriver \
       --machine-type=n1-standard-16 \
       --workload-pool=$PROJECT_ID.svc.id.goog \
       --location=$ZONE
   

   Este comando también habilita GcsFuseCsiDriver, que permite que los pods activen buckets de Cloud Storage como sistemas de archivos locales. La creación del clúster puede tomar varios minutos.

2. Para comunicarte con tu clúster, configura kubectl:

   gcloud container clusters get-credentials CLUSTER_NAME \
       --location=LOCATION
   

3. Crea un grupo de nodos de porción de TPU de varios hosts:

   gcloud container node-pools create v6e-16 \
       --location=$ZONE \
       --cluster=$CLUSTER_NAME \
       --machine-type=ct6e-standard-4t \
       --threads-per-core=1 \
       --tpu-topology=4x4 \
       --num-nodes=4
   

GKE aprovisiona un grupo de nodos que consta de cuatro VMs de TPU Trillium (v6e), que se configuran juntas como una porción de TPU de varios hosts, con una topología 4x4, que está lista para cargas de trabajo de entrenamiento distribuido.

El clúster de GKE habilitado para el operador de Ray instala automáticamente KubeRay y el webhook de KubeRay TPU en tu clúster.

Configura un bucket de Cloud Storage y una cuenta de servicio

1. Crea un bucket de Cloud Storage para los puntos de control compartidos entre los nodos TPU de varios hosts.

   gsutil mb -p ${PROJECT_ID} -c STANDARD -l ${REGION} gs://${GS_BUCKET}
   

2. Para habilitar el acceso al bucket de Cloud Storage, crea una cuenta de servicio de Kubernetes:

   kubectl create serviceaccount ${KSA_NAME} --namespace ${NAMESPACE}
   

3. Para habilitar el acceso al bucket de Cloud Storage, agrega las vinculaciones de políticas de IAM necesarias a la cuenta de servicio:

   gcloud storage buckets add-iam-policy-binding gs://${GS_BUCKET} \
       --member "principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/${NAMESPACE}/sa/${KSA_NAME}" \
       --role "roles/storage.objectUser"
   

Crear la secuencia de comandos de entrenamiento

La siguiente secuencia de comandos usa JaxTrainer de Ray Train para ejecutar un trabajo de entrenamiento distribuido de MaxText. La secuencia de comandos configura el entorno de entrenamiento para un grupo de nodos de porción de TPU de varios hosts y ejecuta el trabajo de entrenamiento de MaxText en cada nodo trabajador. La función train_loop_per_worker encapsula el punto de entrada principal de MaxText y usa el programador distribuido de Ray para ejecutar el entrenador de MaxText en una porción de TPU de varios hosts.

1. Guarda la siguiente secuencia de comandos de Python como maxtext_ray_trainer.py:

   import os
   from absl import app
   import logging
   from typing import Sequence
   import ray
   from ray.train.v2.api.config import ScalingConfig, RunConfig
   from ray.train.v2.jax import JaxTrainer
   
   def train_loop_per_worker(config):
       from MaxText.train import main as maxtext_main
   
       argv = config["argv"]
       maxtext_main(argv)
   
   def main(argv: Sequence[str]):
       trainer = JaxTrainer(
           train_loop_per_worker=train_loop_per_worker,
           train_loop_config={"argv": argv},
           scaling_config=ScalingConfig(
               use_tpu=True,
               num_workers=4,
               topology="4x4",
               accelerator_type="TPU-V6E",
               resources_per_worker={"TPU": 4},
               placement_strategy="SPREAD",
           ),
           run_config=RunConfig(
               name="maxtext_jaxtrainer",
               worker_runtime_env={
                   "env_vars": {
                       "JAX_PLATFORMS": "tpu",
                       "ENABLE_PJRT_COMPATIBILITY": "true",
                       "TPU_SLICE_BUILDER_DUMP_CHIP_FORCE": "true",
                       "TPU_SLICE_BUILDER_DUMP_ICI": "true",
                       "XLA_FLAGS": "--xla_dump_to=/tmp/xla_dump_file --xla_dump_hlo_as_proto",
                   }
               },
           ),
       )
       result = trainer.fit()
       logging.info("Training complete!")
       ray.shutdown()
   
   if __name__ == "__main__":
       app.run(main)

2. Para alojar la imagen personalizada, crea un repositorio de Artifact Registry:

   gcloud artifacts repositories create ${ARTIFACT_REGISTRY} \
       --repository-format=docker --location=${REGION} && \
   gcloud auth configure-docker ${REGION}-docker.pkg.dev
   

3. Para compilar una imagen que incluya dependencias de Ray y MaxText para el entrenamiento, crea un Dockerfile:

   # Start from a Ray base image which includes JaxTrainer API.
   # Maxtext with TPU requires Python 3.12.
   FROM rayproject/ray:2.49.1-py312
   
   USER root
   RUN groupadd -r ray 2>/dev/null || true && usermod -g ray ray
   
   RUN sudo apt-get update -y \
     && sudo apt-get install --no-install-recommends -y git \
     && sudo rm -rf /var/lib/apt/lists/*
   
   WORKDIR /app
   
   # Clone the Maxtext repo and build from source, installing TPU dependencies.
   RUN git clone https://github.com/AI-Hypercomputer/maxtext.git
   
   RUN pip install --no-cache-dir uv
   
   RUN cd maxtext && \
       uv pip install --no-cache --system -e .[tpu] --resolution=lowest && \
       install_maxtext_github_deps
   
   # Copy the Ray Maxtext trainer to run on the remote container.
   COPY maxtext_ray_trainer.py .
   
   RUN chown -R ray:ray .
   ENV PYTHONPATH=/app/maxtext/src:/app/maxtext:/app
   USER ray

   Nota: En este instructivo, se usa una imagen personalizada porque las dependencias son demasiado grandes para instalarlas con el comando de Ray runtime_env.

4. Compila, etiqueta y envía la imagen de Docker a Artifact Registry:

   export DOCKER_IMAGE=${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_REGISTRY}/ray-maxtext:latest
   gcloud builds submit --tag ${DOCKER_IMAGE}
   

Entrena el modelo

1. Guarda el siguiente manifiesto de muestra como maxtext-tpu-cluster.yaml:

   apiVersion: ray.io/v1
   kind: RayCluster
   metadata:
     name: maxtext-tpu-cluster
   spec:
     headGroupSpec:
       rayStartParams: {}
       template:
         metadata:
           annotations:
             gke-gcsfuse/volumes: "true"
             gke-gcsfuse/cpu-limit: "0"
             gke-gcsfuse/memory-limit: "0"
             gke-gcsfuse/ephemeral-storage-limit: "0"
         spec:
           serviceAccountName: ${KSA_NAME}
           containers:
             - name: ray-head
               image: ${DOCKER_IMAGE}
               imagePullPolicy: IfNotPresent
               ports:
               - containerPort: 6379
                 name: gcs-server
               - containerPort: 8265
                 name: dashboard
               - containerPort: 10001
                 name: client
               resources:
                 limits:
                   memory: "16Gi"
                 requests:
                   cpu: "8"
                   memory: "16Gi"
               volumeMounts:
               - name: gcs-fuse-csi-ephemeral
                 mountPath: /data
               - name: dshm
                 mountPath: /dev/shm
           volumes:
           - name: gcs-fuse-cache
             emptyDir:
               medium: Memory
           - name: dshm
             emptyDir:
               medium: Memory
           - name: gcs-fuse-csi-ephemeral
             csi:
               driver: gcsfuse.csi.storage.gke.io
               volumeAttributes:
                 bucketName: ${GS_BUCKET}
                 mountOptions: "implicit-dirs"
     workerGroupSpecs:
       - replicas: 1
         numOfHosts: 4
         groupName: tpu-group
         rayStartParams: {}
         template:
           metadata:
             annotations:
               gke-gcsfuse/volumes: "true"
               gke-gcsfuse/cpu-limit: "0"
               gke-gcsfuse/memory-limit: "0"
               gke-gcsfuse/ephemeral-storage-limit: "0"
           spec:
             serviceAccountName: ${KSA_NAME}
             containers:
               - name: ray-worker
                 image: ${DOCKER_IMAGE}
                 imagePullPolicy: IfNotPresent
                 resources:
                   limits:
                     memory: 200G
                     google.com/tpu: "4"
                   requests:
                     cpu: "8"
                     memory: 200G
                     google.com/tpu: "4"
                 env:
                   - name: JAX_PLATFORMS
                     value: tpu
                   - name: ENABLE_PJRT_COMPATIBILITY
                     value: "true"
                 volumeMounts:
                 - name: gcs-fuse-csi-ephemeral
                   mountPath: /data
                 - name: dshm
                   mountPath: /dev/shm
             volumes:
             - name: gcs-fuse-cache
               emptyDir:
                 medium: Memory
             - name: dshm
               emptyDir:
                 medium: Memory
             - name: gcs-fuse-csi-ephemeral
               csi:
                 driver: gcsfuse.csi.storage.gke.io
                 volumeAttributes:
                   bucketName: ${GS_BUCKET}
                   mountOptions: "implicit-dirs"
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 4x4

   La especificación de RayCluster anterior crea un grupo de trabajador TPU con cuatro trabajadores (numOfHosts: 4) por réplica. Cada trabajador solicita cuatro chips TPU (google.com/tpu: "4"). Los trabajadores se programarán en un nodo que ejecute TPU Trillium (tpu-v6e-slice) y que forme parte de la misma porción de varios hosts colocada. KubeRay escala los cuatro trabajadores de forma atómica, y GKE inicia las variables de entorno JAX necesarias, así como las afinidades de pods para la programación, a través de un webhook mutante.

2. Para configurar los valores necesarios en el archivo YAML, crea el RayCluster con envsubst:

   envsubst < maxtext-tpu-cluster.yaml | kubectl apply -f -
   

3. Verifica que el clúster esté listo y en ejecución:

   kubectl get rayclusters maxtext-tpu-cluster
   

   El resultado debería ser similar al siguiente ejemplo:

   NAME                  DESIRED WORKERS   AVAILABLE WORKERS   CPUS   MEMORY        GPUS   STATUS   AGE
   maxtext-tpu-cluster   4                 4                   40     798027216Ki   0      ready    11m
   

4. Para acceder al panel de Ray a través del servicio principal de Ray, establece una sesión de redirección de puertos:

   kubectl port-forward svc/maxtext-tpu-cluster-head-svc 8265:8265 2>&1 >/dev/null &
   

5. Verifica que se pueda acceder al RayCluster desde tu entorno local:

   ray list nodes --address http://localhost:8265
   

   El resultado debería ser similar al siguiente ejemplo:

   ======== List: 2025-09-13 03:53:16.988269 ========
   Stats:
   ------------------------------
   Total: 5
   Table:
   ------------------------------
       NODE_ID                                                   NODE_IP    IS_HEAD_NODE    STATE    STATE_MESSAGE    NODE_NAME    RESOURCES_TOTAL                  LABELS
   0  92c79d04c34b659c1e3044f7642ad3fd47eb16f290785237149fab56  10.84.0.9
   (...)
   

6. Envía la secuencia de comandos de JaxTrainer al RayCluster y verifica que RayJob se complete correctamente:

   ray job submit \
     --address http://localhost:8265 \
     -- python /app/maxtext_ray_trainer.py \
         /app/maxtext/src/MaxText/configs/base.yml \
          base_output_directory=/data/ \
         dataset_type=synthetic \
         per_device_batch_size=1 \
         max_target_length=4096 \
         model_name=llama3-8b \
         steps=100 \
         ici_fsdp_parallelism=4 \
         ici_tensor_parallelism=4 \
         run_name=rayjob-8b-4096-tp4-4x4
   

   El comando anterior envía la secuencia de comandos de Python, que llama al código de Ray de JaxTrainer al RayCluster. El comando ray job submit incluye algunos argumentos específicos de MaxText para pasar a la configuración del modelo.

   Nota: La secuencia de comandos puede tardar un par de minutos en completarse.

   En tu terminal, deberías ver un resultado similar al siguiente:

   (RayTrainWorker pid=21663, ip=10.12.3.6) completed step: 99, seconds: 1.100, TFLOP/s/device: 179.739, Tokens/s/device: 3725.218, total_weights: 65536, loss: 0.000 [repeated 3x across cluster]
   
   ------------------------------------------
   Job 'raysubmit_zCrJcWnuymMQv4C3' succeeded
   ------------------------------------------
   

Completado correctamente: Completaste el entrenamiento con JaxTrainer de Ray Train y MaxText en grupos de nodo TPU de varios hosts distribuidos.

Limpia

Para evitar que se apliquen cargos a tu Google Cloud cuenta por los recursos usados en este instructivo, borra el proyecto que contiene los recursos o conserva el proyecto y borra los recursos individuales.

1. Borra el RayCluster:

   kubectl delete raycluster maxtext-tpu-cluster
   

2. Borra el clúster de GKE:

   gcloud container clusters delete $CLUSTER_NAME --zone=$ZONE
   

3. Borra el bucket de Cloud Storage:

   gsutil rm -r gs://${GS_BUCKET}
   

4. Borra el repositorio de Artifact Registry:

   gcloud artifacts repositories delete ${ARTIFACT_REGISTRY} --location=${REGION} --quiet
   

¿Qué sigue?

• Obtén más información sobre Ray en Kubernetes.
• Obtén información para entregar vLLM en GKE con TPU.
• Obtén información para entregar SDXL en GKE con TPU.
• Obtén más información sobre las TPU en GKE.