Ejecuta SFT para Gemma 4 31B en una TPU v6e de varios hosts

En este instructivo, se muestra cómo ejecutar el ajuste supervisado (SFT) en un clúster de unidad de procesamiento tensorial (TPU) v6e de varios hosts con MaxText y el 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 la entrega.

Resumen de la carga de trabajo

La ejecución de esta carga de trabajo de SFT en Cloud TPU y GKE implica los siguientes componentes:

• Modelo: Gemma 4 31B (google/gemma-4-31b-it), un modelo avanzado de pesos abiertos. En este instructivo, los pesos 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 (v6e-32).
• Método de entrenamiento: SFT con MaxText. La carga de trabajo ejecuta una canalización de ajuste fino escalable en Cloud TPU v6e con JAX. Optimiza los parámetros del modelo en 1,000 pasos con el entrenador de SFT integrado de MaxText (maxtext.trainers.post_train.sft.train_sft) y los conjuntos de datos de ajuste de instrucciones.
• Conjunto de datos y plantilla: UltraChat 200k (HuggingFaceH4/ultrachat_200k), un conjunto de datos de diálogo de varios turnos de alta calidad y gran escala. La carga de trabajo usa este conjunto de datos para ajustar el modelo y mejorar sus capacidades de conversación y seguimiento de instrucciones.
• Infraestructura y framework: GKE, administrado con Cluster Toolkit, en Cloud TPU v6e (v6e-32, con una topología 4x8). El framework es MaxText (tpu_post_training:0.2.4) con 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.
• Ejecuta una carga de trabajo de entrenamiento de SFT en la TPU.
• 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:

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• 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 Realiza una limpieza.

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

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.

• En la tarjeta del modelo google/gemma-4-31b-it en Hugging Face, acepta el contrato de licencia del modelo.

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, haz lo siguiente:
  • Administrador de TPU (roles/tpu.admin)
  • Usuario de la cuenta de servicio (roles/iam.serviceAccountUser)
  • Administrador de Compute (roles/compute.admin)
  • Administrador de almacenamiento (roles/storage.admin)
  • Administrador de Kubernetes Engine (roles/container.admin)
  • Administrador de cuenta de servicio (roles/iam.serviceAccountAdmin)
  • Administrador de IAM de proyecto (roles/resourcemanager.projectIamAdmin)
  • 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 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 TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME="gemma4-31b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"

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 tu 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: 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, de modo que los IDs de cuenta de servicio basados en este se mantengan dentro del límite de 30 caracteres.
• 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.

De forma predeterminada, la secuencia de comandos establece MODEL_NAME en gemma4-31b (el nombre del modelo de MaxText para google/gemma-4-31b-it).

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

Nota: En la documentación de Instala Cluster Toolkit, se describe cómo descargar la versión más reciente en un directorio cluster-toolkit. Sin embargo, debido a que las variables de plano pueden diferir entre las versiones de Cluster Toolkit, para este instructivo, debes instalar la versión exacta especificada en CLUSTER_TOOLKIT_VERSION (v1.103.0) en lugar de la versión más reciente. Cuando extraes el paquete, se proporcionan el objeto binario gcluster y los planos de examples/ que usas en los pasos que siguen. Si extrajiste el paquete en un directorio cluster-toolkit, asegúrate de cambiar a ese directorio (cd cluster-toolkit) antes de ejecutar los comandos posteriores.

Autentica tu entorno

Para asegurarte de que Cluster Toolkit pueda aprovisionar recursos y enviar imágenes a Artifact Registry en tu nombre, autentica Google Cloud CLI y configura tus credenciales predeterminadas de la aplicación:

gcloud auth login
gcloud auth application-default login --no-launch-browser
gcloud config set project "${PROJECT:?}"
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. Prepara la configuración del plano y copia el plano de implementación de ejemplo en un directorio de trabajo aislado:

   mkdir -p tmp
   cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml tmp/

3. Actualiza tmp/gke-tpu-v6e-advanced.yaml para otorgar a la cuenta de servicio del grupo de nodos del clúster (gke-np-sa) el rol de administrador de almacenamiento (roles/storage.admin) para escribir puntos de control en tu bucket de Cloud Storage, cambia el tipo de máquina del grupo de nodos del sistema a e2-standard-8 y fija el grupo de nodos del sistema en tu zona de destino para evitar que se agote el stock de tipos de máquinas regionales:

   # Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
   sed -i "s/- storage.objectViewer/- storage.admin/" tmp/gke-tpu-v6e-advanced.yaml
   # Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
   sed -i "s/n2-standard-8/e2-standard-8/" tmp/gke-tpu-v6e-advanced.yaml
   # 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)]' tmp/gke-tpu-v6e-advanced.yaml

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

   ./gcluster deploy tmp/gke-tpu-v6e-advanced.yaml \
       --vars "project_id=${PROJECT:?},deployment_name=${CLUSTER_NAME:?},region=${REGION:?},zone=${ZONE:?},num_slices=1,tpu_topology=4x8,authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
       --download-dependencies \
       -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 hf-to-mt \
       --compute-type "${TPU_TYPE:?}" \
       --num-slices 1 \
       --image "${CLOUD_IMAGE_NAME:?}" \
       --await-job-completion \
       --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
         python3 -m maxtext.checkpoint_conversion.to_maxtext \
           model_name=${MODEL_NAME:?} \
           hf_access_token=${HF_TOKEN:?} \
           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"

   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 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 SFT:

   ./gcluster job submit --name sft \
       --compute-type "${TPU_TYPE:?}" \
       --num-slices 1 \
       --image "${CLOUD_IMAGE_NAME:?}" \
       --await-job-completion \
       --command "JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true JAX_TRACEBACK_FILTERING=off LIBTPU_INIT_ARGS=' --xla_tpu_scoped_vmem_limit_kib=61440 --xla_tpu_bf16_emission_mode=NATIVE_EMISSION --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true ' \
         python3 -m maxtext.trainers.post_train.sft.train_sft \
         run_name=sft \
         base_output_directory=gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/ \
         model_name=${MODEL_NAME} \
         load_parameters_path=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/0/items/ \
         hf_access_token=${HF_TOKEN:?} \
         dataset_type=hf \
         hf_path=HuggingFaceH4/ultrachat_200k \
         per_device_batch_size=1 steps=1000 \
         profiler=xplane \
         checkpoint_storage_use_zarr3=0 \
         checkpoint_storage_use_ocdbt=0 \
         skip_jax_distributed_system=False"

   El comando incluye la marca --await-job-completion, lo que significa que espera en tu terminal actual hasta que finalice el trabajo de entrenamiento.

2. Opcional: Para transmitir los registros de entrenamiento en vivo 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 sft --main-only -f

   Una vez que se completa la compilación inicial de XLA, el coordinador registra el progreso de salida para cada paso de la ejecución de 1,000 pasos en el formato completed step: <step>, seconds: ..., TFLOP/s/device: ..., loss: ..., con la pérdida de entrenamiento que disminuye de aproximadamente 1.57 en el paso 5 a aproximadamente 0.82 en el paso 1,000. Antes de continuar con la siguiente sección, espera hasta que finalice el paso 1000 y el resultado del comando ./gcluster job list muestre el estado de sft como Succeeded.

   Nota: Los JobSets completados en GKE se limpian automáticamente después de que vence el tiempo de actividad (TTL) configurado, que es de una hora de forma predeterminada (--gke-ttl-after-finished=1h). Las cargas de trabajo completadas no se enumeran en el resultado del comando ./gcluster job list después de que vence su TTL. Sin embargo, todos los puntos de control y los pesos del modelo se almacenan en tu bucket de Cloud Storage.

3. En la terminal original, verifica que los puntos de control de entrenamiento de SFT se generen en tu bucket de Cloud Storage:

   gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/sft/checkpoints/1000/model_params/"

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 mt-to-hf \
       --compute-type "${TPU_TYPE:?}" \
       --num-slices 1 \
       --image "${CLOUD_IMAGE_NAME:?}" \
       --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/sft/checkpoints/1000/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"

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

2. Opcional: Para transmitir los registros de conversión en vivo 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 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 finales de Hugging Face Safetensors se generen 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:?}" --robust --auto-approve
gcloud storage rm -r "gs://${GCS_BUCKET:?}"

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

¿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-32, consulta TPU v6e.