Esegui l'addestramento RL multi-host per Qwen3-30b-a3b su TPU v6e

Questo tutorial mostra come eseguire l'addestramento multi-host di apprendimento per rinforzo (RL) su un cluster di Tensor Processing Unit (TPU) v6e-32 utilizzando MaxText e Cluster Toolkit. Utilizzi Cluster Toolkit per eseguire un carico di lavoro di addestramento multihost ed esportare i risultati nel formato Hugging Face per la pubblicazione.

Obiettivi

  • Installa Cluster Toolkit e le relative dipendenze.
  • Esegui il deployment di un cluster Cluster Toolkit.
  • Converti un modello Hugging Face nel formato MaxText.
  • Esegui un carico di lavoro di addestramento RL sul cluster TPU v6e.
  • Converti di nuovo il modello ottimizzato nel formato Hugging Face per la pubblicazione.

Costi

In questo documento vengono utilizzati i seguenti componenti fatturabili di Google Cloud:

Per generare una stima dei costi in base all'utilizzo previsto, utilizza il calcolatore prezzi.

I nuovi utenti di Google Cloud potrebbero avere diritto a una prova senza costi.

Al termine delle attività descritte in questo documento, puoi evitare l'addebito di ulteriori costi eliminando le risorse che hai creato. Per saperne di più, consulta Esegui la pulizia.

Prima di iniziare

  • Per utilizzare questo tutorial, è necessario un token di accesso a Hugging Face. Puoi registrarti per un account senza costi su Hugging Face. Dopo aver creato un account, genera un token di accesso:

    1. Nella pagina Benvenuto in Hugging Face, fai clic sull'avatar del tuo account e seleziona Token di accesso.
    2. Nella pagina Token di accesso, fai clic su Crea nuovo token.
    3. Seleziona il tipo di token Lettura e inserisci un nome per il token.
    4. Viene visualizzato il token di accesso. Salva il token in un luogo sicuro.

  • Sul sito web di Hugging Face, accetta il contratto di licenza per il modello che prevedi di addestrare. Questo tutorial utilizza il modello qwen3-30b-a3b.

Per ottenere le autorizzazioni necessarie per completare questo tutorial, chiedi all'amministratore di concederti i seguenti ruoli IAM nel progetto:

Per saperne di più sulla concessione dei ruoli, consulta Gestisci l'accesso a progetti, cartelle e organizzazioni.

Potresti anche riuscire a ottenere le autorizzazioni richieste tramite i ruoli personalizzati o altri ruoli predefiniti.

Configura le variabili di ambiente

Imposta le variabili di ambiente:

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"

Sostituisci quanto segue:

  • YOUR_PROJECT_ID: l'ID del tuo Google Cloud progetto.
  • YOUR_REGION: la regione in cui vuoi eseguire il deployment del cluster.
  • YOUR_ZONE: la zona in cui vuoi eseguire il deployment del cluster.
  • YOUR_CLUSTER_NAME: il nome del tuo cluster Google Kubernetes Engine (fino a 20 caratteri).
  • YOUR_BUCKET_NAME: un nome univoco a livello globale per un bucket Cloud Storage.
  • YOUR_RESERVATION_NAME: il nome della prenotazione.
  • YOUR_HF_TOKEN: il token di accesso a Hugging Face.

Installa le dipendenze di Cluster Toolkit

Per completare questo tutorial da un client o una workstation Linux o macOS, segui i passaggi pertinenti in Installare le dipendenze nella documentazione di Cluster Toolkit.

Se utilizzi Cloud Shell, puoi saltare questa sezione.

Installa Cluster Toolkit

Installa il pacchetto predefinito per Cluster Toolkit nella directory di lavoro corrente seguendo le istruzioni riportate in Installare Cluster Toolkit.

Ad esempio, puoi scaricare ed estrarre il bundle nella directory di lavoro attuale come segue:

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

L'estrazione del bundle nella directory di lavoro attuale fornisce il file binario gcluster e i blueprint examples/ che utilizzerai nei passaggi successivi.

Crea il cluster Cluster Toolkit

Per creare ed eseguire il deployment di un cluster Cluster Toolkit con 32 chip TPU v6e, completa i seguenti passaggi:

  1. Crea un bucket Cloud Storage:

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

  2. Copia il progetto base di Cluster Toolkit nella directory di lavoro corrente:

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

  3. Per impostazione predefinita, il account di servizio del pool di nodi del cluster non dispone delle autorizzazioni necessarie per scrivere nel bucket Cloud Storage. Per consentire al account di servizio del pool di nodi di scrivere nel tuo bucket Cloud Storage, devi concedergli il ruolo Storage Admin. Per concedere questo ruolo, modifica il file gke-tpu-v6e-advanced.yaml aggiornando il modulo service-account denominato 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. Utilizza il comando gcluster deploy per eseguire il deployment del cluster Cluster Toolkit utilizzando il blueprint gke-tpu-v6e-advanced.yaml e passando le variabili richieste utilizzando il flag --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 l'autenticazione di Container Registry e concedi il ruolo Amministratore Storage (roles/storage.admin) ai tuoi service account 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

Converti il modello nel formato MaxText

Per addestrare il modello in formato MaxText, devi convertirlo dal formato Hugging Face al formato MaxText.

  1. Per semplificare i comandi successivi, utilizza il comandogcluster job config per configurare il progetto, il cluster e la località predefiniti:

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

  2. Utilizza il gcluster job submit comando per convertire il modello dal formato Hugging Face al formato MaxText e archiviarlo nel 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. Utilizza il comando gcluster job logs per controllare lo stato del job di conversione:

    # 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 che i file del modello convertiti siano disponibili nel bucket Cloud Storage:

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

Avvia il carico di lavoro di addestramento

Una volta completato il processo di conversione, avvia il carico di lavoro di addestramento 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'}}\""

  • Per controllare lo stato del job di addestramento:

    # 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

  • Per verificare che i checkpoint di addestramento siano generati nel bucket Cloud Storage:

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

Converti il modello addestrato di nuovo nel formato Hugging Face

Al termine del carico di lavoro di addestramento, converti nuovamente il modello nel formato 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"

  • Per controllare lo stato del job di conversione:

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

  • Per verificare che i pesi del modello Hugging Face addestrato e i file di configurazione siano presenti nel bucket Cloud Storage:

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

Esegui la pulizia

Per evitare addebiti aggiuntivi, utilizza il comandogcluster destroy per eliminare le risorse create durante questo tutorial:

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

Passaggi successivi