Ejecuta el entrenamiento de aprendizaje por refuerzo de varios hosts para Gemma 4 26B en la TPU v6e

En este instructivo, se muestra cómo ejecutar el entrenamiento de aprendizaje por refuerzo (RL) de varios hosts en un clúster de unidad de procesamiento tensorial (TPU) v6e-64 con MaxText y Cluster Toolkit. Usas Cluster Toolkit para ejecutar una carga de trabajo de entrenamiento de varios hosts y exportar los resultados al formato de Hugging Face para la entrega.

Objetivos

• Instala Cluster Toolkit y sus dependencias.
• Instala MaxText y sus dependencias.
• Implementa un clúster de Cluster Toolkit.
• Convierte un modelo de Hugging Face al formato MaxText.
• Ejecuta una carga de trabajo de entrenamiento de RL en el clúster de TPU v6e.
• Vuelve a convertir el modelo ajustado al formato de Hugging Face para la entrega.

Costos

En este documento, usarás los siguientes componentes facturables de Google Cloud:

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Artifact Registry
• Compute Engine

Para generar una estimación de costos en función del uso previsto, usa la calculadora de precios.

Es posible que los usuarios de Google Cloud nuevos cumplan con los requisitos para acceder a una prueba gratuita.

Cuando completes las tareas que se describen en este documento, podrás borrar los recursos que creaste para evitar que se te siga facturando. Para obtener más información, consulta Libera espacio.

Importante: Para completar este instructivo, debes tener acceso a las TPUs v6e-64. En este instructivo, se supone que ya tienes una reserva para estas TPUs en tu Google Cloud proyecto para asegurarte de tener la capacidad disponible cuando la necesites. Para obtener más información sobre las opciones para reservar capacidad de Cloud TPU para una fecha y hora futuras, consulta Acerca de las reservas de Cloud TPU.

Antes de comenzar

Para usar este instructivo, necesitas un token de acceso de Hugging Face. Puedes registrarte para obtener una cuenta gratuita en Hugging Face. Después de tener una cuenta, genera un token de acceso:

1. En la página de bienvenida a Hugging Face, haz clic en el avatar de tu cuenta y selecciona Tokens de acceso.
2. En la página Tokens de acceso, haz clic en Crear token nuevo.
3. Selecciona el tipo de token Leer y, luego, ingresa un nombre para tu token.
4. Se mostrará tu token de acceso. Guarda el token en un lugar seguro.

• En el sitio web de Hugging Face, acepta el acuerdo de licencia para el modelo que planeas entrenar. En este instructivo, se usa el modelo gemma4-26b.

Si quieres obtener los permisos que necesitas para completar este instructivo, pídele a tu administrador que te otorgue los siguientes roles de IAM en tu proyecto:

• Para completar este instructivo:
  • Administrador de TPU (roles/tpu.admin)
  • Usuario de la cuenta de servicio (roles/iam.serviceAccountUser)
  • Editor de Compute (roles/compute.editor)
• Para preparar tu imagen de contenedor de MaxText:
  • Editor de Cloud Build (roles/cloudbuild.builds.editor)
  • Administrador de Artifact Registry (roles/artifactregistry.admin)
  • Administrador de almacenamiento (roles/storage.admin)
  • Administrador de Service Usage (roles/serviceusage.serviceUsageAdmin)

Para obtener más información sobre cómo otorgar roles, consulta Administra el acceso a proyectos, carpetas y organizaciones.

También puedes obtener los permisos necesarios mediante roles personalizados o cualquier otro rol predefinido.

Configura tus variables de entorno

Configura las variables de entorno:

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
export COMPUTE_TYPE="ct6e-standard-4t"
export TPU_TYPE="v6e-64"
export TOPOLOGY="8x8"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-26b"
export HF_TOKEN="YOUR_HF_TOKEN"

Reemplaza lo siguiente:

• YOUR_PROJECT_ID: Es el ID de tu Google Cloud proyecto.
• YOUR_REGION: Es la región en la que deseas implementar tu clúster.
• YOUR_ZONE: Es la zona en la que deseas implementar tu clúster.
• YOUR_CLUSTER_NAME: Es el nombre de tu clúster de Google Kubernetes Engine.
• YOUR_REPOSITORY_NAME: Es el nombre del repositorio de Artifact Registry para tus imágenes de MaxText.
• YOUR_BUCKET_NAME: Es un nombre único a nivel global para un bucket de Cloud Storage.
• YOUR_RESERVATION_NAME: Es el nombre de tu reserva.
• YOUR_HF_TOKEN: Es tu token de acceso de Hugging Face.

Instala las dependencias de Cluster Toolkit

Para completar este instructivo desde un cliente o una estación de trabajo de Linux o macOS, sigue los pasos pertinentes en Instala dependencias en la documentación de Cluster Toolkit.

Si usas Cloud Shell, puedes omitir esta sección.

Instala Cluster Toolkit

Para instalar el paquete compilado previamente de Cluster Toolkit, sigue las instrucciones que se indican en Instala Cluster Toolkit.

Prepara tu imagen de contenedor de MaxText

Para preparar tu imagen de contenedor de MaxText, incluida la instalación de las dependencias necesarias, completa los siguientes pasos:

1. Crea un bucket de Cloud Storage:

   gcloud storage buckets create gs://$GCS_BUCKET --project=$PROJECT --location=$REGION || true

2. Crea un repositorio de Artifact Registry.

   gcloud artifacts repositories create ${REPOSITORY_NAME} \
       --repository-format=docker \
       --location=$REGION \
       --project=$PROJECT \
       --description="Docker repository for MaxText images in $REGION" || true

3. Crea un archivo en el directorio raíz de tu repositorio con el nombre de archivo cloudbuild.yaml y el siguiente contenido:

   steps:
     - name: 'gcr.io/cloud-builders/docker'
       entrypoint: 'bash'
       args:
         - '-c'
         - |
           set -euo pipefail
   
           # 0. Install prerequisites (if needed)
           apt-get update && apt-get install -y curl || apk add curl || true
   
           # 1. Install uv
           curl -LsSf https://astral.sh/uv/install.sh | sh
           source $$HOME/.local/bin/env
   
           # 2. Setup Python environment and install MaxText runner
           uv venv --python 3.12 --seed maxtext_venv
           source maxtext_venv/bin/activate
           uv pip install maxtext[runner]==0.2.4 --resolution=lowest
   
           # 3. Build the Docker image (Cloud Build has Docker pre-configured)
           build_maxtext_docker_image WORKFLOW=post-training
   
           # 4. Tag the image properly
           docker tag maxtext_base_image ${_CLOUD_IMAGE_NAME}
   
   # Cloud Build automatically pushes images listed here
   images:
     - '${_CLOUD_IMAGE_NAME}'
   
   options:
     machineType: 'E2_HIGHCPU_32'

4. Usa Cloud Build para compilar tu imagen de Docker de MaxText:

   gcloud builds submit . \
       --project=$PROJECT \
       --region=$REGION \
       --substitutions=_CLOUD_IMAGE_NAME="${CLOUD_IMAGE_NAME}"

Crea tu clúster de Cluster Toolkit

Para crear y, luego, implementar un clúster de Cluster Toolkit con 64 chips de TPU v6e, completa los siguientes pasos:

1. Crea un rol personalizado de Identity and Access Management (IAM), llamado gke.gcsfuse.profileUser:

   # The GKE TPU v6e blueprint uses GCS Fuse CSI Storage Profiles which requires a custom IAM role.
   # If this role is not already created in your project, you must create it before deploying.
   gcloud iam roles create gke.gcsfuse.profileUser \
     --project=${PROJECT} \
     --title="GKE GCSFuse Profile User" \
     --description="Allows scanning GCS buckets for objects, retrieving bucket metadata, and creating Anywhere Caches." \
     --permissions="storage.objects.list,storage.buckets.get,storage.anywhereCaches.create,storage.anywhereCaches.get,storage.anywhereCaches.list,storage.anywhereCaches.update"

2. Crea un bucket de Cloud Storage:

   gcloud storage buckets create gs://$GCS_BUCKET --project=$PROJECT --location=$REGION || true

3. De forma predeterminada, la cuenta de servicio del grupo de nodos del clúster no tiene los permisos necesarios para escribir en tu bucket de Cloud Storage. Para permitir que la cuenta de servicio del grupo de nodos escriba en tu bucket de Cloud Storage, debes otorgarle el rol Storage Admin. Para otorgar este rol, edita el archivo gke-tpu-v6e-advanced.yaml actualizando el service-account módulo llamado node_pool_service_account:

   - id: node_pool_service_account
     source: modules/project/service-account
     settings:
       name: gke-np-sa
       project_roles:
       - logging.logWriter
       - monitoring.metricWriter
       - monitoring.viewer
       - stackdriver.resourceMetadata.writer
       - storage.admin            # Change from storage.objectViewer
       - artifactregistry.reader

4. Aplica la configuración personalizada al bloque gke-tpu-v6e-cluster que anula la configuración predeterminada de IPv6 y del tipo de máquina:

   - id: gke-tpu-v6e-cluster
     source: modules/scheduler/gke-cluster
     use: [gke-tpu-v6e-net-0, workload_service_account]
     settings:
       enable_private_ipv6_google_access: false
       system_node_pool_disk_size_gb: $(vars.system_node_pool_disk_size_gb)
       system_node_pool_taints: []
       enable_private_endpoint: false # Allows access from authorized public IPs
       enable_pathways_for_tpus: $(vars.enable_pathways_for_tpus)
       enable_dataplane_v2: true
       configure_workload_identity_sa: true

5. Para implementar tu clúster de Cluster Toolkit, usa el plano gke-tpu-v6e-advanced.yaml y pasa las variables necesarias con la marca --vars:

   ./gcluster deploy examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml \
       --vars "project_id=${PROJECT},deployment_name=${CLUSTER_NAME},region=${REGION},zone=${ZONE},num_slices=${CLUSTER_NODEPOOL_COUNT},tpu_topology=${TOPOLOGY},authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
       -l IGNORE --auto-approve -w

Convierte el modelo al formato MaxText

Para entrenar el modelo en formato MaxText, debes convertirlo del formato de Hugging Face al formato MaxText.

1. Después de terminar de crear tu clúster de Cluster Toolkit, configura Docker:

   # Configure docker for pulling images
   gcloud auth configure-docker gcr.io --quiet
   gcloud auth configure-docker ${REGION}-docker.pkg.dev --quiet

2. Para simplificar los comandos posteriores, configura tu proyecto, clúster y ubicación predeterminados:

   # Configure gcluster Defaults
   ./gcluster job config set project ${PROJECT}
   ./gcluster job config set cluster ${CLUSTER_NAME}
   ./gcluster job config set location ${REGION}

3. Convierte el modelo del formato de Hugging Face al formato MaxText y almacénalo en tu bucket de Cloud Storage:

   ./gcluster job submit \
     --name="gemma4-hf-to-mt" \
     --cluster="${CLUSTER_NAME}" \
     --project="${PROJECT}" \
     --location="${REGION}" \
     --num-slices=1 \
     --image="${CLOUD_IMAGE_NAME}" \
     --compute-type="${COMPUTE_TYPE}" \
     --topology="${TOPOLOGY}" \
     --await-job-completion \
     --command="[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
       python3 -m maxtext.checkpoint_conversion.to_maxtext \
       model_name=${MODEL_NAME} \
       hf_access_token=${HF_TOKEN} \
       --hf_model_path='google/gemma-4-26b-a4b-it' \
       base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/ \
       scan_layers=True \
       use_multimodal=False \
       skip_jax_distributed_system=true \
       checkpoint_storage_use_zarr3=0 \
       checkpoint_storage_use_ocdbt=0 \
       hardware=cpu \
       --lazy_load_tensors=True"

4. Verifica el estado del trabajo de conversión:

   # Use the list command to check status
   ./gcluster job list \
       --cluster ${CLUSTER_NAME} \
       --project ${PROJECT} \
       --location ${REGION}
   
   # Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
   ./gcluster job logs gemma4-hf-to-mt --main-only -f \
       --cluster ${CLUSTER_NAME} \
       --project ${PROJECT} \
       --location ${REGION}

Inicia la carga de trabajo de entrenamiento

Una vez que se complete el proceso de conversión, inicia la carga de trabajo de entrenamiento de RL:

./gcluster job submit \
  --name="gemma4-training" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME}" \
  --compute-type="${COMPUTE_TYPE}" \
  --topology="${TOPOLOGY}" \
  --pathways \
  --pathways-gcs-location="gs://${GCS_BUCKET}/pathways/" \
  --env="GRPC_DNS_RESOLVER=native" \
  --pathways-proxy-env="GRPC_DNS_RESOLVER=native" \
  --pathways-server-env="GRPC_DNS_RESOLVER=native" \
  --pathways-worker-env="GRPC_DNS_RESOLVER=native" \
  --command="export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}'); \
      python3 -c \"import pathlib, tpu_inference.layers.common.fused_moe_gmm as f; p = pathlib.Path(f.__file__); p.write_text(p.read_text().replace('onehot_moe_permute_threshold: int = 0,', 'onehot_moe_permute_threshold: int = 100000,'))\"; \
      JAX_PLATFORMS=proxy,cpu ENABLE_PATHWAYS_PERSISTENCE=1 \
      python3 -m maxtext.trainers.post_train.rl.train_rl \
      run_name=rl \
      base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/trained/ \
      model_name=${MODEL_NAME} \
      scan_layers=False \
      load_parameters_path=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/0/items/ \
      hf_access_token=${HF_TOKEN} \
      num_batches=50 \
      batch_size=8 \
      rollout_tensor_parallelism=2 \
      rollout_expert_parallelism=4 \
      trainer_devices_fraction=0.5 \
      sampler_devices_fraction=0.5 \
      tokenizer_path='google/gemma-4-26b-a4b-it' \
      ici_tensor_parallelism=2 \
      ici_expert_parallelism=4 \
      hbm_utilization_vllm=0.55 \
      remat_policy=full \
      async_scheduling=False \
      allow_split_physical_axes=true \
      ragged_gather_reduce_fallback=True \
      vllm_hf_overrides='{architectures: [\"MaxTextForCausalLM\"]}' \
      vllm_additional_config=\"{'maxtext_config': {'model_name': '${MODEL_NAME}', 'allow_split_physical_axes': 'true', 'use_ragged_sort': 'false', 'ragged_gather_reduce_fallback': 'true', 'prefuse_moe_weights': 'true', 'weight_dtype': 'bfloat16'}}\""

Verifica el estado del trabajo de entrenamiento:

# Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs gemma4-training --main-only -f \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

Vuelve a convertir el modelo entrenado al formato de Hugging Face

Una vez que se complete la carga de trabajo de entrenamiento, vuelve a convertir el modelo al formato de Hugging Face:

./gcluster job submit \
  --name="gemma4-mt-to-hf" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME}" \
  --compute-type="${COMPUTE_TYPE}" \
  --topology="${TOPOLOGY}" \
  --await-job-completion \
  --command="[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
    python3 -m maxtext.checkpoint_conversion.to_huggingface \
      model_name=${MODEL_NAME} \
      hf_access_token=${HF_TOKEN} \
      load_parameters_path=gs://${GCS_BUCKET}/${MODEL_NAME}/trained/rl/checkpoints/actor/50/model_params/ \
      base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ \
      skip_jax_distributed_system=true \
      hardware=cpu \
      scan_layers=True \
      use_multimodal=False \
      weight_dtype=bfloat16 \
      --override_model_architecture"

Verifica el estado del trabajo de conversión:

# Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs gemma4-mt-to-hf --main-only -f \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

Limpia

Para evitar que se generen cargos adicionales, borra los recursos creados durante este instructivo:

./gcluster destroy ${CLUSTER_NAME} --robust
gcloud storage rm -r gs://${GCS_BUCKET}
gcloud artifacts repositories delete ${REPOSITORY_NAME} --location=${REGION} --project=${PROJECT} --quiet

# To delete the local deployment folder
rm -rf .ghpc ${CLUSTER_NAME}

¿Qué sigue?

• Para obtener más información sobre Cloud TPU, consulta Introducción a Cloud TPU.
• Para obtener detalles sobre la arquitectura y la configuración de la TPU v6e-64, consulta TPU v6e.
• Para obtener más información sobre Cluster Toolkit, consulta Descripción general de Cluster Toolkit.