Esegui SFT per Gemma 4 31B su TPU v6e multi-host

Questo tutorial mostra come eseguire il fine-tuning supervisionato (SFT) su un cluster di Tensor Processing Unit (TPU) v6e multi-host utilizzando MaxText e Cluster Toolkit. Utilizza Cluster Toolkit per eseguire un workload di addestramento multi-host ed esportare i risultati nel formato Hugging Face per la pubblicazione.

Riepilogo del workload

L'esecuzione di questo workload SFT su Cloud TPU e GKE prevede i seguenti componenti:

  • Modello: Gemma 4 31B (google/gemma-4-31b-it), un modello open-weights avanzato. In questo tutorial, i pesi del modello vengono scaricati da Hugging Face, convertiti in un formato compatibile con MaxText ed è stato eseguito il deployment su un cluster GKE con infrastruttura Cloud TPU v6e (v6e-32).
  • Metodo di addestramento: SFT utilizzando MaxText. Il workload esegue una pipeline di ottimizzazione scalabile su Cloud TPU v6e utilizzando JAX. Ottimizza i parametri del modello in 1000 passaggi utilizzando l'addestratore SFT integrato di MaxText (maxtext.trainers.post_train.sft.train_sft) e i set di dati di ottimizzazione delle istruzioni.
  • Set di dati e modello: UltraChat 200k (HuggingFaceH4/ultrachat_200k), un set di dati di dialogo multi-turn di alta qualità e su larga scala. Il workload utilizza questo set di dati per ottimizzare il modello e migliorare le sue capacità di conversazione e di esecuzione delle istruzioni.
  • Infrastruttura e framework: GKE, gestito utilizzando Cluster Toolkit, su Cloud TPU v6e (v6e-32, con una topologia 4x8). Il framework è MaxText (tpu_post_training:0.2.4) con un backend JAX.

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 workload di addestramento SFT sulla TPU.
  • Converti 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, devi disporre di un token di accesso a Hugging Face. Puoi creare un account senza costi su Hugging Face. Dopo aver creato un account, genera un token di accesso:

    1. Nella pagina Welcome to 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.

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 TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME="gemma4-31b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"

Sostituisci quanto segue:

  • YOUR_PROJECT_ID: l'ID del progetto Google Cloud.
  • YOUR_REGION: la regione in cui si trova la prenotazione della TPU e in cui esegui il deployment del cluster.
  • YOUR_ZONE: la zona in cui si trova la prenotazione della TPU e in cui esegui il deployment del cluster.
  • YOUR_BUCKET_NAME: un nome univoco a livello globale per un bucket Cloud Storage.
  • YOUR_RESERVATION_NAME: il nome della prenotazione TPU v6e.
  • YOUR_HF_TOKEN: il tuo token di accesso a Hugging Face.

Per impostazione predefinita, lo script imposta MODEL_NAME su gemma4-31b (il nome del modello MaxText per google/gemma-4-31b-it) e CLUSTER_NAME su gke-tpu-v6e. Se personalizzi CLUSTER_NAME, il nome deve iniziare con una lettera minuscola, contenere solo lettere minuscole, numeri e trattini e avere una lunghezza massima di 20 caratteri, in modo che gli ID dei account di servizio generati non superino il limite di 30 caratteri.

Installa le dipendenze di Cluster Toolkit

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

Se utilizzi Cloud Shell, puoi saltare questa sezione.

Installa Cluster Toolkit

Scarica ed estrai il bundle predefinito per la versione specifica di Cluster Toolkit definita in CLUSTER_TOOLKIT_VERSION (v1.103.0) nella directory di lavoro attuale:

case "$(uname -s)-$(uname -m)" in
  Linux-x86_64)              BUNDLE="gcluster_bundle_linux_amd64.tgz" ;;
  Linux-aarch64|Linux-arm64) BUNDLE="gcluster_bundle_linux_arm64.tgz" ;;
  Darwin-x86_64)             BUNDLE="gcluster_bundle_mac_amd64.tgz" ;;
  Darwin-arm64)              BUNDLE="gcluster_bundle_mac_arm64.tgz" ;;
  *) BUNDLE=""; echo "Unsupported platform: $(uname -s)-$(uname -m)" >&2 ;;
esac
mkdir -p cluster-toolkit
curl -fsSL "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:?}/${BUNDLE:?}" \
  | tar -xz -C cluster-toolkit
cd cluster-toolkit

Autentica il tuo ambiente

Per assicurarti che Cluster Toolkit possa eseguire il provisioning delle risorse e il push delle immagini in Artifact Registry per tuo conto, autentica Google Cloud CLI e imposta le Credenziali predefinite dell'applicazione:

gcloud auth login
gcloud auth application-default login --no-launch-browser
gcloud config set project "${PROJECT:?}"
gcloud auth application-default set-quota-project "${PROJECT:?}"

gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker "${REGION:?}-docker.pkg.dev" --quiet

Crea il tuo 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. Esegui lo stage della configurazione del progetto iniziale e copia il progetto iniziale di deployment di esempio in una directory di lavoro isolata:

    mkdir -p tmp
cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml tmp/

  3. Aggiorna tmp/gke-tpu-v6e-advanced.yaml per concedere al account di servizio del pool di nodi del cluster (gke-np-sa) il ruolo Storage Admin (roles/storage.admin) per scrivere i checkpoint nel bucket Cloud Storage, cambia il tipo di macchina del pool di nodi di sistema in e2-standard-8 e assegna il node pool di sistema alla zona di destinazione per evitare esaurimenti di stock dei tipi di macchina regionali:

    # Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
sed -i "s/- storage.objectViewer/- storage.admin/" tmp/gke-tpu-v6e-advanced.yaml
# Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
sed -i "s/n2-standard-8/e2-standard-8/" tmp/gke-tpu-v6e-advanced.yaml
# Pin system node pool to a single zone to avoid cross-zone stockouts
sed -i '/system_node_pool_machine_type/a \      system_node_pool_zones: [$(vars.zone)]' tmp/gke-tpu-v6e-advanced.yaml

  4. Utilizza il comando gcluster deploy per eseguire il deployment del tuo cluster Cluster Toolkit utilizzando il progetto tmp/gke-tpu-v6e-advanced.yaml e passando le variabili richieste utilizzando il flag --vars:

    ./gcluster deploy tmp/gke-tpu-v6e-advanced.yaml \
    --vars "project_id=${PROJECT:?},deployment_name=${CLUSTER_NAME:?},region=${REGION:?},zone=${ZONE:?},num_slices=1,tpu_topology=4x8,authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    --download-dependencies \
    -l IGNORE \
    --auto-approve -w

Converti il modello nel formato MaxText

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

  1. Per semplificare i comandi successivi, utilizza il comando gcluster 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 comando gcluster job submit per convertire il modello dal formato Hugging Face al formato MaxText e archiviarlo nel tuo bucket Cloud Storage:

    ./gcluster job submit --name hf-to-mt \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --await-job-completion \
    --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:?}/${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"

    Il comando include il flag --await-job-completion, il che significa che attende nel terminale corrente fino al completamento della conversione.

  3. (Facoltativo) Per trasmettere in streaming i log di conversione mentre il job è in esecuzione, apri un terminale separato nella directory di lavoro in cui hai estratto gcluster ed esegui il comando gcluster job logs:

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

  4. Nel terminale originale, verifica che i file del modello convertito siano disponibili nel bucket Cloud Storage:

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

Avvia il workload di addestramento

  1. Al termine del processo di conversione, avvia il workload SFT:

    ./gcluster job submit --name sft \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --await-job-completion \
    --command "JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true JAX_TRACEBACK_FILTERING=off LIBTPU_INIT_ARGS=' --xla_tpu_scoped_vmem_limit_kib=61440 --xla_tpu_bf16_emission_mode=NATIVE_EMISSION --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true ' \
      python3 -m maxtext.trainers.post_train.sft.train_sft \
      run_name=sft \
      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:?} \
      dataset_type=hf \
      hf_path=HuggingFaceH4/ultrachat_200k \
      per_device_batch_size=1 steps=1000 \
      profiler=xplane \
      checkpoint_storage_use_zarr3=0 \
      checkpoint_storage_use_ocdbt=0 \
      skip_jax_distributed_system=False"

    Il comando include il flag --await-job-completion, il che significa che attende nel terminale corrente fino al termine del job di addestramento.

  2. (Facoltativo) Per trasmettere in streaming i log di addestramento in tempo reale mentre il job è in esecuzione, apri un terminale separato nella directory di lavoro in cui hai estratto gcluster ed esegui il comando gcluster job logs:

    # 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 sft --main-only -f

    Al termine della compilazione XLA iniziale, il coordinatore registra l'avanzamento dell'output per ogni passaggio dell'esecuzione di 1000 passaggi nel formato completed step: <step>, seconds: ..., TFLOP/s/device: ..., loss: ..., con la perdita di addestramento che diminuisce da circa 1.57 al passaggio 5 a circa 0.82 al passaggio 1000. Prima di procedere alla sezione successiva, attendi il completamento del passaggio 1000 e che l'output del comando ./gcluster job list mostri lo stato di sft come Succeeded.

  3. Nel terminale originale, verifica che i checkpoint di addestramento SFT siano generati nel bucket Cloud Storage:

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/sft/checkpoints/1000/model_params/"

Converti il modello addestrato nel formato Hugging Face

  1. Al termine del workload di addestramento, converti il modello nel formato Hugging Face:

    ./gcluster job submit --name mt-to-hf \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --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/sft/checkpoints/1000/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"

    Il comando include il flag --await-job-completion, il che significa che attende nel terminale corrente fino al completamento della conversione.

  2. (Facoltativo) Per trasmettere in streaming i log di conversione in tempo reale mentre il job è in esecuzione, apri un terminale separato nella directory di lavoro in cui hai estratto gcluster ed esegui il comando gcluster job logs:

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

# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

  3. Nel terminale originale, verifica che i pesi finali di Hugging Face Safetensors siano generati 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 comando gcluster destroy per eliminare le risorse create durante questo tutorial:

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

# To delete the local deployment folder and staged blueprint
rm -rf .ghpc "${CLUSTER_NAME:?}" tmp

Passaggi successivi

  • Per saperne di più su Cloud TPU, consulta Introduzione a Cloud TPU.
  • Per i dettagli sull'architettura e sulla configurazione della TPU v6e-32, consulta TPU v6e.