SFT für Gemma 4 31B auf einer TPU v6e mit mehreren Hosts ausführen

In dieser Anleitung erfahren Sie, wie Sie die überwachte Feinabstimmung (Supervised Fine-Tuning, SFT) auf einem Multi-Host-Tensor Processing Unit (TPU) v6e-Cluster mit MaxText und Cluster Toolkit ausführen. Mit dem Cluster Toolkit führen Sie eine Multi-Host-Trainingsarbeitslast aus und exportieren die Ergebnisse zur Bereitstellung wieder in das Hugging Face-Format.

Zusammenfassung der Arbeitslast

Für die Ausführung dieser SFT-Arbeitslast auf Cloud TPU und GKE sind die folgenden Komponenten erforderlich:

• Modell: Gemma 4 31B (google/gemma-4-31b-it), ein fortschrittliches Open-Weights-Modell. In dieser Anleitung werden die Modellgewichtungen von Hugging Face heruntergeladen, in ein MaxText-kompatibles Format konvertiert und in einem GKE-Cluster mit Cloud TPU v6e-Infrastruktur (v6e-32) bereitgestellt.
• Trainingsmethode: SFT mit MaxText. Die Arbeitslast führt eine skalierbare Feinabstimmungspipeline auf Cloud TPU v6e mit JAX aus. Die Modellparameter werden über 1.000 Schritte hinweg mit dem integrierten SFT-Trainer von MaxText (maxtext.trainers.post_train.sft.train_sft) und Datasets für die Feinabstimmung von Anweisungen optimiert.
• Dataset und Vorlage: UltraChat 200k (HuggingFaceH4/ultrachat_200k), ein umfangreiches, hochwertiges Dataset mit Multi-Turn-Dialogen. Die Arbeitslast verwendet dieses Dataset, um das Modell zu optimieren und seine Fähigkeiten zur Konversation und zum Befolgen von Anweisungen zu verbessern.
• Infrastruktur und Framework: GKE, verwaltet mit Cluster Toolkit, auf Cloud TPU v6e (v6e-32, mit einer 4x8-Topologie). Das Framework ist MaxText (tpu_post_training:0.2.4) mit einem JAX-Backend.

Ziele

• Cluster Toolkit und seine Abhängigkeiten installieren
• Cluster Toolkit-Cluster bereitstellen
• Hugging Face-Modell in das MaxText-Format konvertieren.
• Führen Sie eine SFT-Trainingsarbeitslast auf der TPU aus.
• Konvertieren Sie das optimierte Modell für die Bereitstellung zurück in das Hugging Face-Format.

Kosten

In diesem Dokument verwenden Sie die folgenden kostenpflichtigen Komponenten von Google Cloud:

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Compute Engine

Mit dem Preisrechner können Sie eine Kostenschätzung für Ihre voraussichtliche Nutzung vornehmen.

Neuen Nutzern von Google Cloud steht möglicherweise eine kostenlose Testversion zur Verfügung.

Nach Abschluss der in diesem Dokument beschriebenen Aufgaben können Sie weitere Kosten vermeiden, indem Sie die erstellten Ressourcen löschen. Weitere Informationen finden Sie unter Bereinigen.

Wichtig: Für diese Anleitung benötigen Sie Zugriff auf v6e-32-TPUs. In dieser Anleitung wird davon ausgegangen, dass Sie bereits eine Reservierung für diese TPUs in IhremGoogle Cloud -Projekt haben, damit die Kapazität bei Bedarf verfügbar ist. Weitere Informationen zu den Optionen für die Reservierung von Cloud TPU-Kapazität für ein zukünftiges Datum und eine zukünftige Uhrzeit finden Sie unter Cloud TPU-Reservierungen.

Hinweis

• Für diese Anleitung benötigen Sie ein Hugging Face-Zugriffstoken. Sie können sich bei Hugging Face für ein kostenloses Konto registrieren. Nachdem Sie ein Konto erstellt haben, generieren Sie ein Zugriffstoken:

  1. Klicken Sie auf der Seite Willkommen bei Hugging Face auf Ihren Konto-Avatar und wählen Sie Zugriffstokens aus.
  2. Klicken Sie auf der Seite Zugriffstokens auf Neues Token erstellen.
  3. Wählen Sie den Tokentyp Lesen aus und geben Sie einen Namen für das Token ein.
  4. Ihr Zugriffstoken wird angezeigt. Speichern Sie das Token an einem sicheren Ort.

• Akzeptieren Sie auf der google/gemma-4-31b-it-Modellkarte auf Hugging Face die Lizenzvereinbarung für das Modell.

Bitten Sie Ihren Administrator, Ihnen die folgenden IAM-Rollen für Ihr Projekt zuzuweisen, um die Berechtigungen zu erhalten, die Sie zum Ausführen dieser Anleitung benötigen:

• So schließen Sie diese Anleitung ab:
  • TPU-Administrator (roles/tpu.admin)
  • Service Account User (roles/iam.serviceAccountUser)
  • Compute-Administrator (roles/compute.admin)
  • Storage-Administrator (roles/storage.admin)
  • Kubernetes Engine-Administrator (roles/container.admin)
  • Service Account Admin (roles/iam.serviceAccountAdmin)
  • Projekt-IAM-Administrator (roles/resourcemanager.projectIamAdmin)
  • Service Usage-Administrator ()roles/serviceusage.serviceUsageAdmin

Weitere Informationen zum Zuweisen von Rollen finden Sie unter Zugriff auf Projekte, Ordner und Organisationen verwalten.

Sie können die erforderlichen Berechtigungen auch über benutzerdefinierte Rollen oder andere vordefinierte Rollen erhalten.

Umgebungsvariablen einrichten

Richten Sie die Umgebungsvariablen ein:

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"

Ersetzen Sie Folgendes:

• YOUR_PROJECT_ID: die ID Ihres Projekts in Google Cloud.
• YOUR_REGION: die Region, in der sich Ihre TPU-Reservierung befindet und in der Sie Ihren Cluster bereitstellen.
• YOUR_ZONE: die Zone, in der sich Ihre TPU-Reservierung befindet und in der Sie Ihren Cluster bereitstellen.
• YOUR_CLUSTER_NAME: Name Ihres GKE-Clusters. Der Name muss mit einem Kleinbuchstaben beginnen, darf nur Kleinbuchstaben, Ziffern und Bindestriche enthalten und darf maximal 20 Zeichen lang sein, damit die darauf basierenden Dienstkonto-IDs die Beschränkung von 30 Zeichen nicht überschreiten.
• YOUR_BUCKET_NAME: ein global eindeutiger Name für einen Cloud Storage-Bucket.
• YOUR_RESERVATION_NAME: der Name Ihrer TPU v6e-Reservierung.
• YOUR_HF_TOKEN: Ihr Hugging Face-Zugriffstoken.

Standardmäßig setzt das Script MODEL_NAME auf gemma4-31b (den MaxText-Modellnamen für google/gemma-4-31b-it).

Cluster Toolkit-Abhängigkeiten installieren

Wenn Sie diese Anleitung auf einem Linux- oder macOS-Client oder einer Linux- oder macOS-Workstation ausführen möchten, folgen Sie den entsprechenden Schritten unter Abhängigkeiten installieren in der Cluster Toolkit-Dokumentation.

Wenn Sie Cloud Shell verwenden, können Sie diesen Abschnitt überspringen.

Cluster Toolkit installieren

Laden Sie das vordefinierte Bundle für die in CLUSTER_TOOLKIT_VERSION (v1.103.0) definierte Version des Cluster Toolkits herunter und extrahieren Sie es in Ihr aktuelles Arbeitsverzeichnis:

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

Hinweis: In der Dokumentation Cluster Toolkit installieren wird beschrieben, wie Sie die neueste Version in ein cluster-toolkit-Verzeichnis herunterladen. Da sich Blueprint-Variablen jedoch je nach Cluster Toolkit-Version unterscheiden können, müssen Sie für diese Anleitung die genaue Version installieren, die in CLUSTER_TOOLKIT_VERSION (v1.103.0) angegeben ist, und nicht die neueste Version. Durch das Extrahieren des Bundles erhalten Sie die Binärdatei gcluster und die Blueprints examples/, die Sie in den folgenden Schritten verwenden. Wenn Sie das Bundle in ein cluster-toolkit-Verzeichnis extrahiert haben, müssen Sie in dieses Verzeichnis (cd cluster-toolkit) wechseln, bevor Sie die nachfolgenden Befehle ausführen.

Umgebung authentifizieren

Damit das Cluster Toolkit Ressourcen bereitstellen und Images in Ihrem Namen in Artifact Registry hochladen kann, müssen Sie die Google Cloud CLI authentifizieren und Ihre Standardanmeldedaten für Anwendungen festlegen:

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

Cluster Toolkit-Cluster erstellen

Führen Sie die folgenden Schritte aus, um einen Cluster Toolkit-Cluster mit 32 v6e-TPU-Chips zu erstellen und bereitzustellen:

1. Erstellen Sie einen Cloud Storage-Bucket:

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

2. Stellen Sie die Blueprint-Konfiguration bereit und kopieren Sie den Beispiel-Blueprint für die Bereitstellung in ein isoliertes Arbeitsverzeichnis:

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

3. Aktualisieren Sie tmp/gke-tpu-v6e-advanced.yaml, um dem Dienstkonto des Clusterknotenpools (gke-np-sa) die Rolle „Storage-Administrator“ (roles/storage.admin) zuzuweisen, damit Prüfpunkte in Ihren Cloud Storage-Bucket geschrieben werden können. Ändern Sie den Maschinentyp des Systemknotenpools in e2-standard-8 und pinnen Sie den Systemknotenpool an Ihre Zielzone, um regionale Engpässe bei Maschinentypen zu vermeiden:

   # 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. Verwenden Sie den Befehl gcluster deploy, um Ihren Cluster Toolkit-Cluster mit dem Blueprint tmp/gke-tpu-v6e-advanced.yaml bereitzustellen und die erforderlichen Variablen mit dem Flag --vars zu übergeben:

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

Modell in das MaxText-Format konvertieren

Wenn Sie das Modell im MaxText-Format trainieren möchten, müssen Sie es vom Hugging Face-Format in das MaxText-Format konvertieren.

1. Verwenden Sie den gcluster job config-Befehl, um Ihr Standardprojekt, Ihren Standardcluster und Ihren Standardstandort zu konfigurieren und nachfolgende Befehle zu vereinfachen:

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

2. Verwenden Sie den Befehl gcluster job submit, um das Modell vom Hugging Face-Format in das MaxText-Format zu konvertieren und in Ihrem Cloud Storage-Bucket zu speichern:

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

   Der Befehl enthält das Flag --await-job-completion. Das bedeutet, dass er in Ihrem aktuellen Terminal wartet, bis die Konvertierung abgeschlossen ist.

3. Optional: Wenn Sie die Konvertierungslogs streamen möchten, während der Job ausgeführt wird, öffnen Sie ein separates Terminal im Arbeitsverzeichnis, in dem Sie gcluster extrahiert haben, und führen Sie den Befehl gcluster job logs aus:

   # 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. Prüfen Sie im ursprünglichen Terminal, ob die konvertierten Modelldateien in Ihrem Cloud Storage-Bucket verfügbar sind:

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

Trainingsarbeitslast starten

1. Starten Sie nach Abschluss der Konvertierung die SFT-Arbeitslast:

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

   Der Befehl enthält das Flag --await-job-completion. Das bedeutet, dass er im aktuellen Terminal wartet, bis der Trainingsjob abgeschlossen ist.

2. Optional: Wenn Sie die Livetrainingslogs während der Ausführung des Jobs streamen möchten, öffnen Sie ein separates Terminal im Arbeitsverzeichnis, in dem Sie gcluster extrahiert haben, und führen Sie den Befehl gcluster job logs aus:

   # 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

   Nach Abschluss der ersten XLA-Kompilierung protokolliert der Koordinator den Ausgabefortschritt für jeden Schritt des 1.000-Schritte-Laufs im Format completed step: <step>, seconds: ..., TFLOP/s/device: ..., loss: .... Der Trainingsverlust sinkt dabei von etwa 1.57 bei Schritt 5 auf etwa 0.82 bei Schritt 1.000. Bevor Sie mit dem nächsten Abschnitt fortfahren, warten Sie, bis Schritt 1000 abgeschlossen ist und die Ausgabe des Befehls ./gcluster job list den Status von sft als Succeeded anzeigt.

   Hinweis: Abgeschlossene JobSets in GKE werden automatisch bereinigt, nachdem ihre konfigurierte Gültigkeitsdauer (TTL) abgelaufen ist. Diese beträgt standardmäßig eine Stunde (--gke-ttl-after-finished=1h). Abgeschlossene Arbeitslasten werden nach Ablauf ihrer TTL nicht mehr in der Ausgabe des Befehls ./gcluster job list aufgeführt. Alle Prüfpunkte und Modellgewichtungen werden jedoch in Ihrem Cloud Storage-Bucket gespeichert.

3. Prüfen Sie im ursprünglichen Terminal, ob die SFT-Trainingsprüfpunkte in Ihrem Cloud Storage-Bucket generiert wurden:

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

Trainiertes Modell wieder in das Hugging Face-Format konvertieren

1. Nachdem die Trainingsarbeitslast abgeschlossen ist, konvertieren Sie das Modell wieder in das Hugging Face-Format:

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

   Der Befehl enthält das Flag --await-job-completion. Das bedeutet, dass er in Ihrem aktuellen Terminal wartet, bis die Konvertierung abgeschlossen ist.

2. Optional: Wenn Sie die Live-Konvertierungslogs während der Ausführung des Jobs streamen möchten, öffnen Sie ein separates Terminal im Arbeitsverzeichnis, in dem Sie gcluster extrahiert haben, und führen Sie den Befehl gcluster job logs aus:

   # 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. Prüfen Sie im ursprünglichen Terminal, ob die endgültigen Hugging Face Safetensors-Gewichtungen in Ihrem Cloud Storage-Bucket generiert wurden:

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

Bereinigen

Verwenden Sie den Befehl gcluster destroy, um die in dieser Anleitung erstellten Ressourcen zu löschen und zusätzliche Kosten zu vermeiden:

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

Nächste Schritte

• Weitere Informationen zu Cloud TPU finden Sie unter Einführung in Cloud TPU.
• Details zur Architektur und Konfiguration der v6e-32-TPU finden Sie unter TPU v6e.