Überwachte Feinabstimmung mit mehreren Hosts für das Modell Qwen3-14b mit MaxText ausführen

In dieser Anleitung erfahren Sie Schritt für Schritt, wie Sie mit MaxText auf Cloud TPU ein überwachtes Fine-Tuning (Supervised Fine-Tuning, SFT) für das Modell Qwen3-14b ausführen. Sie lernen, wie Sie ein spezialisiertes Container-Image erstellen, einen Google Kubernetes Engine-Cluster (GKE) mit Pathways mithilfe des Accelerated Processing Kit (XPK) bereitstellen und eine Trainingsarbeitslast mit mehreren Hosts ausführen.

Ziele

• Sie erfahren, wie Sie ein benutzerdefiniertes MaxText-Container-Image erstellen, das für die Nachbearbeitung optimiert ist.
• Sie stellen einen GKE-Cluster mit XPK und aktiviertem Pathways bereit.
• Sie konvertieren das Modell Qwen3 14b vom Hugging Face-Format in das MaxText-Format.
• Sie führen eine SFT-Trainingsarbeitslast mit mehreren Hosts auf Cloud TPU aus.
• Sie konvertieren das feinabgestimmte Modell zur Bereitstellung wieder 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
• Artifact Registry
• Compute Engine, for the temporary build virtual machine (VM)

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

Neuen Google Cloud Nutzern vonsteht 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.

Hinweis

• Prüfen Sie, ob Ihr Nutzerkonto oder Dienstkonto die folgenden Rollen hat:
  • roles/compute.admin zum Erstellen der Build-VM
  • roles/artifactregistry.admin zum Verwalten des Docker-Repositorys
  • roles/storage.admin zum Verwalten des Daten-Buckets
  • roles/container.admin zum Erstellen und Verwalten des Google Kubernetes Engine-Clusters
  • roles/iam.serviceAccountAdmin zum Erstellen des Dienstkontos für die Arbeitslast
  • roles/resourcemanager.projectIamAdminzum Festlegen von IAM-Richtlinien (Identity and Access Management)
  • roles/iam.serviceAccountUser zum Handeln als Dienstkonto
• Installieren und initialisieren Sie das Google Cloud CLI.
• Prüfen Sie, ob Python 3.12 oder höher auf Ihrer Workstation installiert ist.

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

  1. Klicken Sie auf der Seite Welcome to Hugging Face (Willkommen bei Hugging Face) auf Ihren Kontoavatar und wählen Sie Access tokens (Zugriffstokens) aus.
  2. Klicken Sie auf der Seite Access tokens (Zugriffstokens) auf Create new token (Neues Token erstellen).
  3. Wählen Sie den Tokentyp Read (Lesen) aus und geben Sie einen Namen für Ihr Token ein.
  4. Ihr Zugriffstoken wird angezeigt. Speichern Sie das Token an einem sicheren Ort.

Umgebung einrichten

Richten Sie Ihre Umgebungsvariablen ein, indem Sie das folgende Skript ausführen:

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"
export CLOUD_IMAGE_NAME="$REGION-docker.pkg.dev/$PROJECT/maxtext-images/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-14b"
export HF_TOKEN="YOUR_HF_TOKEN"

Ersetzen Sie Folgendes:

• YOUR_PROJECT_ID: Ihre Google Cloud Projekt-ID
• YOUR_REGION: die Region, die Sie verwenden möchten
• YOUR_ZONE: die Zone, die Sie verwenden möchten
• YOUR_CLUSTER_NAME: ein Name für Ihren Google Kubernetes Engine-Cluster
• YOUR_GCS_BUCKET: ein eindeutiger Name für Ihren Cloud Storage-Bucket
• YOUR_RESERVATION_NAME: Ihre Kapazitätsreservierung
• YOUR_HF_TOKEN: Ihr Hugging Face-Zugriffstoken

MaxText-Container-Image vorbereiten

Führen Sie die folgenden Schritte aus, um Ihr MaxText-Container-Image vorzubereiten und die erforderlichen Abhängigkeiten zu installieren:

1. Erstellen Sie einen Cloud Storage-Bucket:

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

2. Erstellen Sie ein Artifact Registry-Repository:

   gcloud artifacts repositories create maxtext-images \
       --repository-format=docker \
       --location=$REGION \
       --project=$PROJECT \
       --description="Docker repository for MaxText images in $REGION" || true

3. Erstellen Sie im Stammverzeichnis Ihres Repositorys eine Datei mit dem Namen cloudbuild.yaml und dem folgenden Inhalt:

   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.1 --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:
     # We use a high-CPU machine to match the n4-standard-16 from the VM tutorial
     machineType: 'E2_HIGHCPU_32'

4. Erstellen Sie mit Cloud Build Ihr MaxText-Docker-Image:

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

Google Kubernetes Engine-Cluster erstellen

Um das SFT-Training für das Modell Qwen3 14b auszuführen, benötigen Sie einen Google Kubernetes Engine-Cluster mit TPU-Chips. Installieren Sie das Accelerated Processing Kit (XPK) und erstellen Sie einen GKE-Cluster mit Pathways-Unterstützung.

# Start with creating a new virtual environment to install XPK in.
VENV_DIR=venvp3
python3 -m venv $VENV_DIR
source $VENV_DIR/bin/activate
pip install xpk==1.14.0

xpk cluster create-pathways \
  --num-slices=${CLUSTER_NODEPOOL_COUNT} \
  --tpu-type=${TPU_TYPE} \
  --pathways-gce-machine-type=${PW_CPU_MACHINE_TYPE} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --cluster=${CLUSTER_NAME} \
  --custom-cluster-arguments="--enable-ip-alias" \
  --reservation=$RESERVATION \
  --default-pool-cpu-machine-type=n4-standard-16

gcloud container clusters get-credentials $CLUSTER_NAME \
  --location=$REGION \
  --project $PROJECT

Modell für das Training vorbereiten

Konvertieren Sie das Basismodell mit einer CPU-basierten Arbeitslast in das MaxText-Format. Führen Sie diese Aufgabe nicht parallel auf mehreren Maschinen aus. Der folgende Befehl enthält eine Prüfung, um sicherzustellen, dass die Konvertierung nur auf einem TPU-Knoten ausgeführt wird.

xpk workload create \
  --workload "qwen-hf-to-mt" \
  --docker-image $CLOUD_IMAGE_NAME \
  --cluster ${CLUSTER_NAME} \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --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}/qwen-3-14b/max-text-format/ \
  scan_layers=True \
  use_multimodal=False \
  skip_jax_distributed_system=true \
  hardware=cpu \
  --lazy_load_tensors=True"

Fortschritt der Modellkonvertierung verfolgen

So verfolgen Sie den Fortschritt der Konvertierung:

1. Führen Sie den Befehl kubectl get pod aus, um die Pods aufzulisten, die in Ihrem GKE-Cluster geplant sind.
2. Suchen Sie den Pod mit dem Namen qwen-hf-to-mt-slice-job-0-0-HASH.
3. Führen Sie den Befehl aus, um die Ausgabe des Pods in Echtzeit zu prüfen kubectl logs -f POD_NAME.

Trainingsarbeitslast starten

Nach Abschluss der Konvertierung können Sie die SFT-Feinabstimmungsarbeitslast mit XPK starten.

xpk workload create-pathways \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-training" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 \
  python3 -m maxtext.trainers.post_train.sft.train_sft \
  run_name=sft \
  base_output_directory=gs://${GCS_BUCKET}/qwen-3-14b/trained/ \
  model_name=${MODEL_NAME} \
  load_parameters_path=gs://${GCS_BUCKET}/qwen-3-14b/max-text-format/0/items/ \
  hf_access_token=${HF_TOKEN} \
  per_device_batch_size=1 \
  steps=1000 \
  profiler=xplane \
  checkpoint_storage_use_zarr3=0 \
  checkpoint_storage_use_ocdbt=0 \
  enable_single_controller=True"

Trainingsarbeitslast überwachen

Überwachen Sie den Status Ihrer Arbeitslast mit dem XPK-Befehlszeilentool.

xpk workload list --cluster ${CLUSTER_NAME} --project ${PROJECT} --zone ${ZONE}

Verwenden Sie die Google Cloud Console, um Logs und die TPU-Auslastung aufzurufen. Sie können Logs auch mit dem folgenden Befehl aufrufen:

kubectl logs -f qwen-training-pathways-head-0-0-HASH

Ersetzen Sie HASH durch den numerischen Hash im Namen Ihres Pods. Führen Sie den Befehl kubectl get pod aus und prüfen Sie die zurückgegebene Liste der Pods, um den Wert dieses Hash zu prüfen.

Trainiertes Modell wieder in das Hugging Face-Format konvertieren

Nach Abschluss der Trainingsarbeitslast konvertieren Sie die Prüfpunkte wieder in das Hugging Face-Format.

xpk workload create \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-mt-to-hf" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --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}/qwen-3-14b/trained/sft/checkpoints/1000/model_params/ \
  base_output_directory=gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ \
  skip_jax_distributed_system=true \
  hardware=cpu \
  scan_layers=True \
  use_multimodal=False \
  weight_dtype=bfloat16"

Führen Sie den Befehl kubectl logs -f qwen-mt-to-hf-slice-job-0-0-HASH, aus, um den Fortschritt der Konvertierung zu verfolgen. Ersetzen Sie HASH durch den numerischen Hash im Namen Ihres Pods.

Nach Abschluss der Konvertierung ist Ihr feinabgestimmtes Modell, das in gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ gespeichert ist, einsatzbereit.

Bereinigen

Löschen Sie die in dieser Anleitung erstellten Ressourcen, einschließlich Ihres Google Kubernetes Engine-Clusters, Ihres Cloud Storage-Bucket und Ihres Artifact Registry-Repositorys, um zusätzliche Kosten zu vermeiden.

Achtung: Wenn Sie das in dieser Anleitung trainierte Modell weiterhin verwenden möchten, speichern Sie eine Version, auf die Sie zugreifen können, bevor Sie den Cloud Storage-Bucket löschen, in dem es enthalten ist.

Führen Sie den folgenden Befehl aus, um die Ressourcen zu löschen, die Sie für diese Anleitung erstellt haben:

xpk cluster delete --cluster $CLUSTER_NAME --project $PROJECT --zone $ZONE --force

gcloud storage rm --recursive gs://$GCS_BUCKET

gcloud artifacts repositories delete maxtext-images --location=$REGION --project=$PROJECT --quiet

Nächste Schritte

• Weitere Informationen zu Cloud TPU finden Sie unter Einführung in Cloud TPU.
• Architektur- und Konfigurationsdetails für die v6e-32 TPU finden Sie unter TPU v6e.