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-32 con MaxText y Cluster Toolkit. Usa Cluster Toolkit para ejecutar una carga de trabajo de entrenamiento de varios hosts y exportar los resultados al formato de Hugging Face para entregarlos.

Resumen de la carga de trabajo

La ejecución de esta carga de trabajo de entrenamiento posterior de RL en Cloud TPU y GKE incluye los siguientes componentes:

  • Modelo: Gemma 4 26B A4B Instruct (google/gemma-4-26b-a4b-it), un modelo avanzado de pesos abiertos ajustado por instrucciones. En este instructivo, las ponderaciones del modelo se descargan de Hugging Face, se convierten en un formato compatible con MaxText y se implementan en un clúster de GKE con la infraestructura de Cloud TPU v6e.
  • Algoritmo de aprendizaje por refuerzo (RL): Optimización de políticas relativas de grupo (GRPO) con MaxText y Pathways La carga de trabajo ejecuta una canalización escalable posterior al entrenamiento con el entrenador de RL integrado de MaxText (maxtext.trainers.post_train.rl.train_rl) vinculado con un backend de muestreador de vLLM. GRPO optimiza la política generando un grupo de resultados por instrucción y evaluando las recompensas relativas sin necesidad de un modelo crítico independiente.
  • Conjunto de datos y plantilla: OpenMathInstruct-2 (openmathinstruct2_rl.json), un conjunto de datos de ajuste de instrucciones matemáticas a gran escala. La carga de trabajo usa plantillas de chat estructuradas durante la fase de lanzamiento del RL para mejorar las capacidades de razonamiento matemático en varios pasos.
  • Infraestructura y framework: GKE, administrado con Cluster Toolkit, en Cloud TPU v6e (porción ct6e-standard-4t, v6e-32 con una topología 4x8). El framework es MaxText (tpu_post_training:0.2.4) con Pathways y un backend de JAX.

Objetivos

  • Instala Cluster Toolkit y sus dependencias.
  • Implementa un clúster de Cluster Toolkit.
  • Convierte un modelo de Hugging Face al formato de MaxText.
  • Ejecutar 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 entregarlo.

Costos

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

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 Realiza una limpieza.

Antes de comenzar

  • Necesitas un token de acceso de Hugging Face para usar este instructivo. 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 Welcome to 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 Lectura y, luego, ingresa un nombre para el token.
    4. Se mostrará tu token de acceso. Guarda el token en un lugar seguro.

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 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 a través de 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 GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/maxtext-images/tpu_post_training:0.2.4"
export COMPUTE_TYPE="ct6e-standard-4t"
export TPU_TYPE="v6e-32"
export TOPOLOGY="4x8"
export CLUSTER_NODEPOOL_COUNT=1
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-26b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"
export HF_TOKEN="YOUR_HF_TOKEN"

Reemplaza lo siguiente:

  • YOUR_PROJECT_ID: Es el ID de tu proyecto de Google Cloud .
  • YOUR_REGION: Es la región en la que se encuentra tu reserva de TPU y en la que implementas el clúster.
  • YOUR_ZONE: Es la zona en la que se encuentra tu reserva de TPU y en la que implementas tu clúster.
  • YOUR_CLUSTER_NAME: Es el nombre del clúster de GKE. El nombre debe comenzar con una letra minúscula, contener solo letras minúsculas, números y guiones, y tener 20 caracteres o menos.
  • 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 de TPU v6e.
  • 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

Descarga y extrae el paquete compilado previamente para la versión específica de Cluster Toolkit definida en CLUSTER_TOOLKIT_VERSION (v1.103.0) en tu directorio de trabajo actual:

case "$(uname -s)-$(uname -m)" in
  Linux-x86_64)              BUNDLE="gcluster_bundle_linux_amd64.tgz" ;;
  Linux-aarch64|Linux-arm64) BUNDLE="gcluster_bundle_linux_arm64.tgz" ;;
  Darwin-x86_64)             BUNDLE="gcluster_bundle_mac_amd64.tgz" ;;
  Darwin-arm64)              BUNDLE="gcluster_bundle_mac_arm64.tgz" ;;
  *) BUNDLE=""; echo "Unsupported platform: $(uname -s)-$(uname -m)" >&2 ;;
esac
mkdir -p cluster-toolkit
curl -fsSL "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:?}/${BUNDLE:?}" \
  | tar -xz -C cluster-toolkit
cd cluster-toolkit

Autentica tu entorno

Antes de ejecutar las secuencias de comandos de implementación, debes autenticar Google Cloud CLI y configurar tus credenciales predeterminadas de la aplicación (ADC). Esto garantiza que Cluster Toolkit pueda aprovisionar recursos y enviar imágenes a Artifact Registry en tu nombre.

gcloud auth login
gcloud auth application-default login --no-launch-browser
gcloud auth application-default set-quota-project "${PROJECT:?}"

gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker "${REGION:?}-docker.pkg.dev" --quiet

Crea tu clúster de Cluster Toolkit

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

  1. Crea un bucket de Cloud Storage:

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

  2. Usa el comando gcluster deploy para implementar tu clúster de Cluster Toolkit con el plano gke-tpu-v6e-advanced.yaml y pasar las variables requeridas con la marca --vars:

    TMP_DIR=$(mktemp -d)
BLUEPRINT_FILE="${TMP_DIR}/gke-tpu-v6e-advanced.yaml"

cp ./examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml "${BLUEPRINT_FILE}"

# Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
sed -i "s/n2-standard-8/e2-standard-8/" "${BLUEPRINT_FILE}"
# Pin system node pool to a single zone to avoid cross-zone stockouts
sed -i '/system_node_pool_machine_type/a \      system_node_pool_zones: [$(vars.zone)]' "${BLUEPRINT_FILE}"
# Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
sed -i "s/- storage.objectViewer/- storage.admin/" "${BLUEPRINT_FILE}"

./gcluster deploy "${BLUEPRINT_FILE}" \
    --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 de MaxText

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

  1. Para simplificar los comandos posteriores, usa el comando gcluster job config para configurar 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:?}"

  2. Usa el comando gcluster job submit para convertir el modelo del formato de Hugging Face al formato de MaxText y almacenarlo en tu bucket de Cloud Storage:

    ./gcluster job submit \
  --name="gemma4-hf-to-mt" \
  --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=False \
    use_multimodal=False \
    skip_jax_distributed_system=true \
    checkpoint_storage_use_zarr3=0 \
    checkpoint_storage_use_ocdbt=0 \
    hardware=cpu \
    --lazy_load_tensors=True"

    El comando incluye la marca --await-job-completion, lo que significa que espera en tu terminal actual hasta que finalice la conversión.

  3. Opcional: Para transmitir los registros de conversión mientras se ejecuta el trabajo, abre una terminal independiente en el directorio de trabajo en el que extrajiste gcluster y ejecuta el comando gcluster job logs:

    # Use the list command to check status
./gcluster job list

# 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

  4. En la terminal original, verifica que los archivos del modelo convertidos estén disponibles en tu bucket de Cloud Storage:

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/max-text-format/"

Inicia la carga de trabajo de entrenamiento

  1. 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" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME:?}" \
  --compute-type="${COMPUTE_TYPE:?}" \
  --topology="${TOPOLOGY:?}" \
  --pathways \
  --pathways-gcs-location="gs://${GCS_BUCKET:?}/pathways/" \
  --gke-ttl-after-finished="24h" \
  --restarts=0 \
  --env="GRPC_DNS_RESOLVER=native" \
  --env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-proxy-env="GRPC_DNS_RESOLVER=native" \
  --pathways-proxy-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-server-env="GRPC_DNS_RESOLVER=native" \
  --pathways-server-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-worker-env="GRPC_DNS_RESOLVER=native" \
  --pathways-worker-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --command="export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}') && \
    export VLLM_ENABLE_V1_MULTIPROCESSING=0 && \
    python3 -c \"import pathlib, re; p = pathlib.Path('/deps/src/maxtext/trainers/post_train/rl/utils_rl.py'); p.write_text(re.sub('optax[.]adamw[(][^)]+[)]', 'optax.adafactor(learning_rate=learning_rate)', p.read_text()))\" && \
    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:?} \
    data_template_path=maxtext/examples/chat_templates/openmathinstruct2_rl.json \
    num_batches=50 \
    batch_size=4 \
    train_micro_batch_size=2 \
    max_target_length=640 \
    max_prefill_predict_length=256 \
    mu_dtype=bfloat16 \
    grad_dtype=bfloat16 \
    rollout_tensor_parallelism=4 \
    rollout_expert_parallelism=1 \
    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 \
    ici_fsdp_parallelism=-1 \
    hbm_utilization_vllm=0.7 \
    remat_policy=full \
    async_scheduling=False \
    allow_split_physical_axes=true \
    ragged_gather_reduce_fallback=True \
    enable_dp_attention=False \
    decode_sampling_temperature=0.8 \
    decode_sampling_top_k=50 \
    decode_sampling_nucleus_p=0.95 \
    learning_rate=2e-5 \
    learning_rate_schedule_steps=100 \
    rl.num_generations=4 \
    rl.reshard_chunk_size=4 \
    debug=True \
    vllm_hf_overrides='{\"architectures\": [\"MaxTextForCausalLM\"]}' \
    vllm_additional_config=\"{'maxtext_config': {'model_name': '${MODEL_NAME}', 'model_call_mode': 'inference', 'enable_dp_attention': false, 'allow_split_physical_axes': true, 'use_ragged_sort': false, 'ragged_gather_reduce_fallback': true, 'prefuse_moe_weights': true, 'weight_dtype': 'bfloat16'}}\""

  2. Verifica el estado, espera a que el Pod coordinador esté listo y transmite los registros del trabajo de entrenamiento:

    # Use the list command to check status (initially shows "Suspended" briefly while the CPU node pool scales up)
./gcluster job list

# Ensure kubectl is configured
gcloud container clusters get-credentials "${CLUSTER_NAME:?}" \
  --location="${REGION:?}" \
  --project="${PROJECT:?}"

# Wait for the Pathways head pod to be ready before streaming logs
kubectl wait --for=condition=Ready pod \
    -l job-name=gemma4-training-pathways-head-0 \
    --timeout=600s

# Stream training logs (runs for 50 steps; wait for "Post RL Training: ... mean_reward=..." and job status Succeeded before converting to Hugging Face)
kubectl logs -f \
    -l job-name=gemma4-training-pathways-head-0 \
    -c workload-container

    Inmediatamente después del envío, el resultado del comando ./gcluster job list podría mostrar brevemente Suspended mientras GKE escala verticalmente el grupo de nodos de CPU del coordinador de Pathways (cpu-np). Cuando comienza el trabajo, genera una evaluación inicial previa al entrenamiento en el paso 0 (Pre RL Training: ... mean_reward=...), pausa el resultado del registro durante varios minutos durante la compilación inicial de XLA y, luego, entrena durante 50 pasos, imprimiendo Reshard finished y generaciones de lanzamiento en cada paso sin números de paso.

    Antes de continuar con el siguiente paso, espera hasta que los registros muestren Post RL Training: ... mean_reward=... y el resultado del comando ./gcluster job list muestre el estado de gemma4-training como Succeeded.

  3. Verifica que los puntos de control de entrenamiento se generen en tu bucket de Cloud Storage:

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/rl/checkpoints/actor/"

Vuelve a convertir el modelo entrenado al formato de Hugging Face

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

    ./gcluster job submit \
  --name="gemma4-mt-to-hf" \
  --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=False \
      use_multimodal=False \
      weight_dtype=bfloat16 \
      --override_model_architecture"

  2. Opcional: Para supervisar los registros de conversión en vivo mientras el comando --await-job-completion anterior espera en tu primera terminal, abre una terminal independiente en el directorio de trabajo en el que extrajiste gcluster y usa el comando gcluster job logs:

    # Use the list command to check status
./gcluster job list 

# 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

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

  3. En la terminal original, verifica que los pesos y los archivos de configuración del modelo de Hugging Face entrenado estén presentes en tu bucket de Cloud Storage:

    gcloud storage ls -l --readable-sizes "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/hf-trained/"

Realiza una limpieza

Para evitar que se generen cargos adicionales, usa el comando gcluster destroy para borrar los recursos que se crearon durante este instructivo:

./gcluster destroy "${CLUSTER_NAME:?}" --auto-approve
gcloud storage rm -r "gs://${GCS_BUCKET:?}"

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

¿Qué sigue?