Ejecutar entrenamiento de aprendizaje por refuerzo multi-host para Qwen3-30b-a3b en 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. 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.
  • 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 RL en el clúster de TPU v6e.
  • Convierte 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 Realiza una limpieza.

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 que tengas 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 Tokens de acceso, haz clic en Crear token nuevo.
    3. Selecciona el tipo de token Read 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 del modelo que planeas entrenar. En este instructivo, se usa el modelo qwen3-30b-a3b.

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 TOPOLOGY="4x8"
export CLUSTER_NODEPOOL_COUNT=1
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-30b-a3b"
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 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 (hasta 20 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.
  • YOUR_HF_TOKEN: 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 del Cluster Toolkit.

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

Instala Cluster Toolkit

Instala el paquete compilado previamente para Cluster Toolkit en tu directorio de trabajo actual siguiendo las instrucciones en Instala Cluster Toolkit.

Por ejemplo, puedes descargar y extraer el paquete en tu directorio de trabajo actual de la siguiente manera:

wget -qO- "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:-v1.103.0}/gcluster_bundle_linux_amd64.tgz" | tar -xz

Si extraes el paquete en tu directorio de trabajo actual, se proporcionarán el objeto binario gcluster y los planos examples/ que usarás en los pasos siguientes.

Crea tu clúster de Cluster Toolkit

Para crear e implementar un clúster de Cluster Toolkit con 32 chips de 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. Copia el plano de Cluster Toolkit en tu directorio de trabajo actual:

    cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml .

  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, actualiza el módulo service-account llamado node_pool_service_account editando el archivo gke-tpu-v6e-advanced.yaml:

    - 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
    - artifactregistry.reader

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

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

  5. Configura la autenticación de Container Registry y otorga el rol de administrador de almacenamiento (roles/storage.admin) a tus cuentas de servicio de GKE:

    gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker "${REGION}-docker.pkg.dev" --quiet
gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-wl-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet
gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-np-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet

Convierte el modelo al formato de MaxText

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

  1. Para simplificar los comandos posteriores, usa el comandogcluster 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 qwen-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='Qwen/Qwen3-30B-A3B-Instruct-2507' \
  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"

  3. Usa el comando gcluster job logs para verificar el estado del trabajo de conversión:

    # 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 qwen-hf-to-mt --main-only -f

  4. Verifica que los archivos del modelo convertido 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

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

./gcluster job submit \
  --name="qwen-rl" \
  --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="(echo 190G > /sys/fs/cgroup/memory.max || echo 190G > /sys/fs/cgroup/memory/memory.limit_in_bytes) 2>/dev/null || true && \
    export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}') && \
    export VLLM_ENABLE_V1_MULTIPROCESSING=0 && \
    python3 -c \"import pathlib, tunix.generate.vllm_sampler as vs; p = pathlib.Path(vs.__file__); p.write_text(p.read_text().replace('reshard_chunk_size: Optional[int] = None', 'reshard_chunk_size: Optional[int] = 4').replace('reshard_chunk_size=self.config.reshard_chunk_size', 'reshard_chunk_size=4'))\" && \
    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} \
    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 \
    num_test_batches=0 \
    batch_size=8 \
    train_micro_batch_size=8 \
    max_target_length=512 \
    max_prefill_predict_length=256 \
    mu_dtype=bfloat16 \
    grad_dtype=bfloat16 \
    rollout_tensor_parallelism=1 \
    rollout_expert_parallelism=4 \
    trainer_devices_fraction=0.5 \
    sampler_devices_fraction=0.5 \
    tokenizer_path='Qwen/Qwen3-30B-A3B-Instruct-2507' \
    ici_tensor_parallelism=2 \
    ici_expert_parallelism=4 \
    ici_fsdp_parallelism=-1 \
    hbm_utilization_vllm=0.55 \
    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'}}\""

  • Para verificar el estado del trabajo de entrenamiento, sigue estos pasos:

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

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

# Check progress of the job
kubectl logs -f \
    -l jobset.sigs.k8s.io/replicatedjob-name=pathways-head \
    -c workload-container

  • Para verificar que los puntos de control del entrenamiento se generen en tu bucket de Cloud Storage, haz lo siguiente:

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

Cómo volver 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="qwen-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=True \
  use_multimodal=False \
  weight_dtype=bfloat16 \
  --override_model_architecture"

  • Para verificar el estado del trabajo de conversión, haz lo siguiente:

    # 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 qwen-mt-to-hf --main-only -f
# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

  • Para verificar que los pesos y los archivos de configuración del modelo de Hugging Face entrenado estén presentes en tu bucket de Cloud Storage, haz lo siguiente:

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

Realiza una limpieza

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

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

# To delete the local deployment folder and copied blueprint
rm -rf .ghpc "${CLUSTER_NAME}" gke-tpu-v6e-advanced.yaml

¿Qué sigue?