Ejecuta el ajuste supervisado multihost en el modelo Qwen3-14b con MaxText

En este instructivo, se proporciona una guía paso a paso para ejecutar el ajuste supervisado (SFT) en el modelo Qwen3-14b con MaxText en Cloud TPU. Aprenderás a compilar una imagen de contenedor especializada, aprovisionar un clúster de Google Kubernetes Engine (GKE) con Pathways mediante el kit de tratamiento de datos acelerado (XPK) y ejecutar una carga de trabajo de entrenamiento de varios hosts.

Objetivos

• Aprende a compilar una imagen de contenedor de MaxText personalizada optimizada para el entrenamiento posterior.
• Aprovisiona un clúster de GKE con XPK con Pathways habilitado.
• Convierte el modelo Qwen3 14b del formato de Hugging Face al formato de MaxText.
• Ejecuta una carga de trabajo de entrenamiento de SFT de varios hosts en Cloud 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:

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Artifact Registry
• Compute Engine, for the temporary build virtual machine (VM)

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

• Verifica que tu cuenta de usuario o cuenta de servicio tenga las siguientes funciones:
  • roles/compute.admin, para crear la VM de compilación
  • roles/artifactregistry.admin, para administrar el repositorio de Docker
  • roles/storage.admin, para administrar el bucket de datos
  • roles/container.admin, para crear y administrar el clúster de Google Kubernetes Engine
  • roles/iam.serviceAccountAdmin, para crear la cuenta de servicio de la carga de trabajo
  • roles/resourcemanager.projectIamAdmin, para configurar las políticas de Identity and Access Management (IAM)
  • roles/iam.serviceAccountUser, para actuar como la cuenta de servicio
• Instala y, luego, inicializa Google Cloud CLI.
• Verifica que hayas instalado Python 3.12 o una versión posterior en tu estación de trabajo.

• 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 Access tokens.
  2. En la página Access tokens, haz clic en Create new token.
  3. Selecciona el tipo de token Read y, luego, ingresa un nombre para tu token.
  4. Se muestra tu token de acceso. Guarda el token en un lugar seguro.

Configura el 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="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"
export CLOUD_IMAGE_NAME="$REGION-docker.pkg.dev/$PROJECT/maxtext-images/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-14b"
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 que deseas usar.
• YOUR_ZONE: Es la zona que deseas usar.
• YOUR_CLUSTER_NAME: Es un nombre para tu clúster de Google Kubernetes Engine.
• YOUR_GCS_BUCKET: Es un nombre único para tu bucket de Cloud Storage.
• YOUR_RESERVATION_NAME: Es tu reserva de capacidad.
• YOUR_HF_TOKEN: Es tu token de acceso de Hugging Face.

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 maxtext-images \
       --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.1 --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 Google Kubernetes Engine

Para ejecutar el entrenamiento de SFT en el modelo Qwen3 14b, necesitas un clúster de Google Kubernetes Engine equipado con chips de TPU. Instala el kit de tratamiento de datos acelerado (XPK) y crea un clúster de GKE con compatibilidad con Pathways.

# Start with creating a new virtual environment to install XPK in.
VENV_DIR=venvp3
python3 -m venv $VENV_DIR
source $VENV_DIR/bin/activate
pip install xpk==1.14.0

xpk cluster create-pathways \
  --num-slices=${CLUSTER_NODEPOOL_COUNT} \
  --tpu-type=${TPU_TYPE} \
  --pathways-gce-machine-type=${PW_CPU_MACHINE_TYPE} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --cluster=${CLUSTER_NAME} \
  --custom-cluster-arguments="--enable-ip-alias" \
  --reservation=$RESERVATION \
  --default-pool-cpu-machine-type=n4-standard-16

gcloud container clusters get-credentials $CLUSTER_NAME \
  --location=$REGION \
  --project $PROJECT

Prepara el modelo para el entrenamiento

Convierte el modelo base al formato de MaxText con una carga de trabajo basada en CPU. No ejecutes esta tarea en varias máquinas en paralelo. El siguiente comando incluye una verificación para garantizar que la conversión se ejecute en un solo nodo TPU.

xpk workload create \
  --workload "qwen-hf-to-mt" \
  --docker-image $CLOUD_IMAGE_NAME \
  --cluster ${CLUSTER_NAME} \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --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}/qwen-3-14b/max-text-format/ \
  scan_layers=True \
  use_multimodal=False \
  skip_jax_distributed_system=true \
  hardware=cpu \
  --lazy_load_tensors=True"

Realiza un seguimiento del progreso de la conversión del modelo

Para realizar un seguimiento del progreso de la conversión, haz lo siguiente:

1. Para enumerar los pods programados en tu clúster de GKE, ejecuta el comando kubectl get pod.
2. Busca el pod llamado qwen-hf-to-mt-slice-job-0-0-HASH.
3. Para inspeccionar el resultado del pod en tiempo real, ejecuta el comando kubectl logs -f POD_NAME.

Inicia la carga de trabajo de entrenamiento

Una vez que se complete el proceso de conversión, puedes iniciar la carga de trabajo de ajuste de SFT con XPK.

xpk workload create-pathways \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-training" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 \
  python3 -m maxtext.trainers.post_train.sft.train_sft \
  run_name=sft \
  base_output_directory=gs://${GCS_BUCKET}/qwen-3-14b/trained/ \
  model_name=${MODEL_NAME} \
  load_parameters_path=gs://${GCS_BUCKET}/qwen-3-14b/max-text-format/0/items/ \
  hf_access_token=${HF_TOKEN} \
  per_device_batch_size=1 \
  steps=1000 \
  profiler=xplane \
  checkpoint_storage_use_zarr3=0 \
  checkpoint_storage_use_ocdbt=0 \
  enable_single_controller=True"

Supervisa la carga de trabajo de entrenamiento

Supervisa el estado de tu carga de trabajo con la interfaz de línea de comandos (CLI) de XPK.

xpk workload list --cluster ${CLUSTER_NAME} --project ${PROJECT} --zone ${ZONE}

Para ver los registros y el uso de la TPU, usa la Google Cloud consola. También puedes ver los registros ejecutando el siguiente comando:

kubectl logs -f qwen-training-pathways-head-0-0-HASH

Reemplaza HASH por el hash numérico en el nombre de tu pod. Para verificar el valor de este hash, ejecuta el comando kubectl get pod y examina la lista de pods que se muestra.

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 los puntos de control al formato de Hugging Face.

xpk workload create \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-mt-to-hf" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --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}/qwen-3-14b/trained/sft/checkpoints/1000/model_params/ \
  base_output_directory=gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ \
  skip_jax_distributed_system=true \
  hardware=cpu \
  scan_layers=True \
  use_multimodal=False \
  weight_dtype=bfloat16"

Para realizar un seguimiento del progreso de la conversión, ejecuta el comando kubectl logs -f qwen-mt-to-hf-slice-job-0-0-HASH, y reemplaza HASH por el hash numérico en el nombre de tu pod.

Una vez que se complete la conversión, tu modelo ajustado almacenado en gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ estará listo para usarse.

Limpia

Para evitar que se apliquen cargos adicionales, borra los recursos creados durante este instructivo, incluidos tu clúster de Google Kubernetes Engine, tu bucket de Cloud Storage y tu repositorio de Artifact Registry.

**Precaución:** Si planeas seguir usando el modelo que entrenaste en este instructivo, guarda una versión a la que puedas acceder antes de borrar el bucket de Cloud Storage que lo contiene.

Para borrar los recursos que creaste en este instructivo, ejecuta el siguiente comando:

xpk cluster delete --cluster $CLUSTER_NAME --project $PROJECT --zone $ZONE --force

gcloud storage rm --recursive gs://$GCS_BUCKET

gcloud artifacts repositories delete maxtext-images --location=$REGION --project=$PROJECT --quiet

¿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.