Exécuter un entraînement RL multi-hôte pour Qwen3-30b-a3b sur TPU v6e

Ce tutoriel explique comment exécuter un entraînement par apprentissage par renforcement (RL) multi-hôte sur un cluster v6e-32 de Tensor Processing Unit (TPU) à l'aide de MaxText et de Cluster Toolkit. Vous utilisez Cluster Toolkit pour exécuter une charge de travail d'entraînement multihôte et exporter les résultats au format Hugging Face pour le serving.

Objectifs

  • Installez Cluster Toolkit et ses dépendances.
  • Déployez un cluster Cluster Toolkit.
  • Convertissez un modèle Hugging Face au format MaxText.
  • Exécutez une charge de travail d'entraînement par renforcement sur le cluster TPU v6e.
  • Reconvertissez le modèle affiné au format Hugging Face pour le diffuser.

Coûts

Dans ce document, vous utilisez les composants facturables suivants de Google Cloud :

Pour obtenir une estimation des coûts en fonction de votre utilisation prévue, utilisez le simulateur de coût.

Les nouveaux utilisateurs de Google Cloud peuvent bénéficier d'un essai sans frais.

Une fois que vous avez terminé les tâches décrites dans ce document, supprimez les ressources que vous avez créées pour éviter que des frais vous soient facturés. Pour en savoir plus, consultez la section Effectuer un nettoyage.

Avant de commencer

  • Vous avez besoin d'un jeton d'accès Hugging Face pour suivre ce tutoriel. Vous pouvez créer un compte sans frais sur Hugging Face. Une fois que vous avez un compte, générez un jeton d'accès :

    1. Sur la page Bienvenue sur Hugging Face, cliquez sur l'avatar de votre compte, puis sélectionnez Jetons d'accès.
    2. Sur la page Jetons d'accès, cliquez sur Créer un jeton.
    3. Sélectionnez le type de jeton Lecture et saisissez un nom pour votre jeton.
    4. Votre jeton d'accès s'affiche. Enregistrez le jeton dans un endroit sûr.

  • Sur le site Web Hugging Face, acceptez le contrat de licence du modèle que vous prévoyez d'entraîner. Ce tutoriel utilise le modèle qwen3-30b-a3b.

Pour obtenir les autorisations nécessaires pour suivre ce tutoriel, demandez à votre administrateur de vous accorder les rôles IAM suivants sur votre projet :

Pour en savoir plus sur l'attribution de rôles, consultez la page Gérer l'accès aux projets, aux dossiers et aux organisations.

Vous pouvez également obtenir les autorisations requises avec des rôles personnalisés ou d'autres rôles prédéfinis.

Configurer vos variables d'environnement

Configurez vos variables d'environnement :

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"

Remplacez les éléments suivants :

  • YOUR_PROJECT_ID : ID de votre projet Google Cloud .
  • YOUR_REGION : région dans laquelle vous souhaitez déployer votre cluster.
  • YOUR_ZONE : zone dans laquelle vous souhaitez déployer votre cluster.
  • YOUR_CLUSTER_NAME : nom de votre cluster Google Kubernetes Engine (20 caractères maximum).
  • YOUR_BUCKET_NAME : nom unique au niveau mondial pour un bucket Cloud Storage.
  • YOUR_RESERVATION_NAME : nom de votre réservation.
  • YOUR_HF_TOKEN : votre jeton d'accès Hugging Face.

Installer les dépendances de Cluster Toolkit

Pour suivre ce tutoriel depuis un client ou une station de travail Linux ou macOS, suivez les étapes correspondantes dans Installer les dépendances de la documentation Cluster Toolkit.

Si vous utilisez Cloud Shell, vous pouvez ignorer cette section.

Installer Cluster Toolkit

Installez le bundle prédéfini pour Cluster Toolkit dans votre répertoire de travail actuel en suivant les instructions de la page Installer Cluster Toolkit.

Par exemple, vous pouvez télécharger et extraire le bundle dans votre répertoire de travail actuel comme suit :

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

L'extraction du bundle dans votre répertoire de travail actuel fournit le binaire gcluster et les plans examples/ que vous utiliserez dans les étapes suivantes.

Créer votre cluster Cluster Toolkit

Pour créer et déployer un cluster Cluster Toolkit avec 32 puces TPU v6e, procédez comme suit :

  1. Créez un bucket Cloud Storage :

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

  2. Copiez le plan Cluster Toolkit dans votre répertoire de travail actuel :

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

  3. Par défaut, le compte de service du pool de nœuds de votre cluster ne dispose pas des autorisations requises pour écrire dans votre bucket Cloud Storage. Pour autoriser le compte de service du pool de nœuds à écrire dans votre bucket Cloud Storage, vous devez lui attribuer le rôle Storage Admin. Pour attribuer ce rôle, modifiez le fichier gke-tpu-v6e-advanced.yaml en mettant à jour le module service-account nommé 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
    - artifactregistry.reader

  4. Utilisez la commande gcluster deploy pour déployer votre cluster Cluster Toolkit à l'aide du blueprint gke-tpu-v6e-advanced.yaml et transmettez les variables requises à l'aide de l'option --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. Configurez l'authentification Container Registry et accordez le rôle d'administrateur de l'espace de stockage (roles/storage.admin) à vos comptes de service 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

Convertir le modèle au format MaxText

Pour entraîner le modèle au format MaxText, vous devez le convertir du format Hugging Face au format MaxText.

  1. Pour simplifier les commandes suivantes, utilisez la commande gcluster job config pour configurer votre projet, votre cluster et votre emplacement par défaut :

    # Configure gcluster Defaults
./gcluster job config set project "${PROJECT}"
./gcluster job config set cluster "${CLUSTER_NAME}"
./gcluster job config set location "${REGION}"

  2. Exécutez la commande gcluster job submit pour convertir le modèle du format Hugging Face au format MaxText et le stocker dans votre bucket 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. Utilisez la commande gcluster job logs pour vérifier l'état du job de conversion :

    # 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. Vérifiez que les fichiers de modèle convertis sont disponibles dans votre bucket Cloud Storage :

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

Démarrer la charge de travail d'entraînement

Une fois le processus de conversion terminé, démarrez la charge de travail d'entraînement 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'}}\""

  • Pour vérifier l'état du job d'entraînement :

    # 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

  • Pour vérifier que les points de contrôle de l'entraînement sont générés dans votre bucket Cloud Storage :

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

Reconvertir le modèle entraîné au format Hugging Face

Une fois la charge de travail d'entraînement terminée, reconvertissez le modèle au format 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"

  • Pour vérifier l'état du job de conversion :

    # 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/

  • Pour vérifier que les fichiers de configuration et les pondérations du modèle Hugging Face entraîné sont présents dans votre bucket Cloud Storage :

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

Effectuer un nettoyage

Pour éviter des frais supplémentaires, utilisez la commande gcluster destroy pour supprimer les ressources créées au cours de ce tutoriel :

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

Étapes suivantes