Ejecuta el ajuste supervisado en Gemma 4 en la TPU v6e

En este instructivo, se muestra cómo ejecutar el ajuste supervisado (SFT) en un clúster de unidades de procesamiento tensorial (TPU) v6e con MaxText y Cluster Toolkit. Usarás 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 SFT en la TPU.
  • 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:

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

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

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 mediante roles personalizados o cualquier otro rol predefinido.

Configura tus variables de entorno

Para configurar tus variables de entorno, ejecuta la siguiente secuencia de comandos:

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="gke-tpu-v6e"
export REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-31b"
export HF_TOKEN="YOUR_HF_TOKEN"
export GCS_BUCKET="YOUR_BUCKET_NAME"

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_REPOSITORY_NAME: Es el nombre del repositorio de Artifact Registry para tus imágenes de MaxText.
  • YOUR_RESERVATION_NAME: Es el nombre de tu reserva.
  • YOUR_HF_TOKEN: Es tu token de acceso de Hugging Face.
  • YOUR_BUCKET_NAME: Es un nombre único a nivel global para un bucket de Cloud Storage.

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.3 --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:
  # We use a high-CPU machine to match the n4-standard-16 from the VM tutorial
  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 32 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 módulo 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. Para implementar tu clúster de Cluster Toolkit, usa el blueprint 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=1,tpu_topology=4x8,authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    --download-dependencies \
    -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. Para convertir el modelo del formato de Hugging Face al formato MaxText y almacenarlo en tu bucket de Cloud Storage, ejecuta la siguiente secuencia de comandos:

    ./gcluster job submit --name hf-to-mt \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

Para verificar el estado del trabajo de conversión, ejecuta el siguiente comando:

# 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 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, puedes iniciar la carga de trabajo de SFT ejecutando el siguiente comando:

./gcluster job submit --name sft \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

Para verificar el estado del trabajo de entrenamiento, ejecuta el siguiente comando:

# 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 sft --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 mt-to-hf \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

Para verificar el estado del trabajo de conversión, ejecuta el siguiente comando:

# 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 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/ - though again, it's ~2x the size of the original...

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