Eseguire l'addestramento di apprendimento per rinforzo multi-host per Gemma 4 26B su TPU v6e

Questo tutorial mostra come eseguire l'addestramento di reinforcement learning (RL) multi-host su un cluster di Tensor Processing Unit (TPU) v6e-64 utilizzando MaxText e Cluster Toolkit. Utilizzerai Cluster Toolkit per eseguire un carico di lavoro di addestramento multi-host ed esportare i risultati nel formato Hugging Face per l'erogazione.

Obiettivi

  • Installare Cluster Toolkit e le relative dipendenze.
  • Installare MaxText e le relative dipendenze.
  • Eseguire il deployment di un cluster Cluster Toolkit.
  • Convertire un modello Hugging Face nel formato MaxText.
  • Eseguire un carico di lavoro di addestramento RL sul cluster TPU v6e.
  • Convertire il modello ottimizzato nel formato Hugging Face per l'erogazione.

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 Google Cloud utenti 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 Liberare spazio.

Prima di iniziare

Per utilizzare questo tutorial, devi disporre di 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 intendi addestrare. Questo tutorial utilizza il modello gemma4-26b.

Per ottenere le autorizzazioni necessarie per completare questo tutorial, chiedi all'amministratore di concederti i seguenti ruoli IAM nel tuo 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.

Impostare 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 REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
export COMPUTE_TYPE="ct6e-standard-4t"
export TPU_TYPE="v6e-64"
export TOPOLOGY="8x8"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-26b"
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 tuo cluster.
  • YOUR_ZONE: la zona in cui vuoi eseguire il deployment del cluster.
  • YOUR_CLUSTER_NAME: il nome del cluster Google Kubernetes Engine.
  • YOUR_REPOSITORY_NAME: il nome del repository Artifact Registry per le immagini MaxText.
  • 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.

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

Installare Cluster Toolkit

Installa il pacchetto predefinito per Cluster Toolkit seguendo le istruzioni riportate in Installare Cluster Toolkit.

Preparare l'immagine container MaxText

Per preparare l'immagine container MaxText, inclusa l'installazione delle dipendenze richieste, completa i seguenti passaggi:

  1. Crea un bucket Cloud Storage:

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

  2. Crea un repository 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 file nella directory principale del repository con il nome file cloudbuild.yaml e i seguenti contenuti:

    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.4 --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:
  machineType: 'E2_HIGHCPU_32'

  4. Utilizza Cloud Build per creare l'immagine Docker MaxText:

    gcloud builds submit . \
    --project=$PROJECT \
    --region=$REGION \
    --substitutions=_CLOUD_IMAGE_NAME="${CLOUD_IMAGE_NAME}"

Creare il cluster Cluster Toolkit

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

  1. Crea un ruolo IAM (Identity and Access Management) personalizzato, denominato 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 Cloud Storage:

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

  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 bucket Cloud Storage, devi concedergli il ruolo Storage Admin. Per concedere questo ruolo, modifica il file gke-tpu-v6e-advanced.yaml aggiornando il service-account modulo 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            # Change from storage.objectViewer
    - artifactregistry.reader

  4. Applica le impostazioni personalizzate al blocco gke-tpu-v6e-cluster che sostituiscono le impostazioni predefinite di IPv6 e del tipo di macchina:

    - id: gke-tpu-v6e-cluster
  source: modules/scheduler/gke-cluster
  use: [gke-tpu-v6e-net-0, workload_service_account]
  settings:
    enable_private_ipv6_google_access: false
    system_node_pool_disk_size_gb: $(vars.system_node_pool_disk_size_gb)
    system_node_pool_taints: []
    enable_private_endpoint: false # Allows access from authorized public IPs
    enable_pathways_for_tpus: $(vars.enable_pathways_for_tpus)
    enable_dataplane_v2: true
    configure_workload_identity_sa: true

  5. Esegui il deployment del cluster Cluster Toolkit utilizzando il progetto base gke-tpu-v6e-advanced.yaml e passando le variabili richieste utilizzando il flag --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=${CLUSTER_NODEPOOL_COUNT},tpu_topology=${TOPOLOGY},authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    -l IGNORE --auto-approve -w

Convertire il modello nel formato MaxText

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

  1. Dopo aver creato il cluster 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. Per semplificare i comandi successivi, configura 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}

  3. Converti il modello dal formato Hugging Face al formato MaxText e archivialo nel bucket Cloud Storage:

    ./gcluster job submit \
  --name="gemma4-hf-to-mt" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --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='google/gemma-4-26b-a4b-it' \
    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"

  4. Controlla lo stato del job di conversione:

    # 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 gemma4-hf-to-mt --main-only -f \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

Avviare il carico di lavoro di addestramento

Al termine del processo di conversione, avvia il carico di lavoro di addestramento RL:

./gcluster job submit \
  --name="gemma4-training" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME}" \
  --compute-type="${COMPUTE_TYPE}" \
  --topology="${TOPOLOGY}" \
  --pathways \
  --pathways-gcs-location="gs://${GCS_BUCKET}/pathways/" \
  --env="GRPC_DNS_RESOLVER=native" \
  --pathways-proxy-env="GRPC_DNS_RESOLVER=native" \
  --pathways-server-env="GRPC_DNS_RESOLVER=native" \
  --pathways-worker-env="GRPC_DNS_RESOLVER=native" \
  --command="export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}'); \
      python3 -c \"import pathlib, tpu_inference.layers.common.fused_moe_gmm as f; p = pathlib.Path(f.__file__); p.write_text(p.read_text().replace('onehot_moe_permute_threshold: int = 0,', 'onehot_moe_permute_threshold: int = 100000,'))\"; \
      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} \
      scan_layers=False \
      load_parameters_path=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/0/items/ \
      hf_access_token=${HF_TOKEN} \
      num_batches=50 \
      batch_size=8 \
      rollout_tensor_parallelism=2 \
      rollout_expert_parallelism=4 \
      trainer_devices_fraction=0.5 \
      sampler_devices_fraction=0.5 \
      tokenizer_path='google/gemma-4-26b-a4b-it' \
      ici_tensor_parallelism=2 \
      ici_expert_parallelism=4 \
      hbm_utilization_vllm=0.55 \
      remat_policy=full \
      async_scheduling=False \
      allow_split_physical_axes=true \
      ragged_gather_reduce_fallback=True \
      vllm_hf_overrides='{architectures: [\"MaxTextForCausalLM\"]}' \
      vllm_additional_config=\"{'maxtext_config': {'model_name': '${MODEL_NAME}', 'allow_split_physical_axes': 'true', 'use_ragged_sort': 'false', 'ragged_gather_reduce_fallback': 'true', 'prefuse_moe_weights': 'true', 'weight_dtype': 'bfloat16'}}\""

Controlla lo stato del job di addestramento:

# 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 gemma4-training --main-only -f \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

Convertire di nuovo il modello addestrato nel formato Hugging Face

Al termine del carico di lavoro di addestramento, converti di nuovo il modello nel formato Hugging Face:

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

Controlla lo stato del job di conversione:

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

Libera spazio

Per evitare addebiti aggiuntivi, elimina le risorse create durante questo tutorial:

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

Passaggi successivi