Überwachte Feinabstimmung von Gemma 4 auf TPU v6e ausführen

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

Ziele

  • Cluster Toolkit und die zugehörigen Abhängigkeiten installieren
  • MaxText und die zugehörigen Abhängigkeiten installieren
  • Cluster Toolkit-Cluster bereitstellen
  • Hugging Face-Modell in das MaxText-Format konvertieren
  • SFT-Trainingsarbeitslast auf der TPU ausführen
  • Das feinabgestimmte Modell zur Bereitstellung wieder in das Hugging Face-Format konvertieren

Kosten

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

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

Für diese Anleitung benötigen Sie ein Hugging Face-Zugriffstoken. Sie können sich für ein kostenloses Konto bei Hugging Face registrieren. Nachdem Sie ein Konto haben, erstellen 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.
  • Akzeptieren Sie auf der Hugging Face-Website die Lizenz vereinbarung für das Modell, das Sie trainieren möchten. In dieser Anleitung wird das Modell gemma4-31b verwendet.

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:

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, indem Sie das folgende Skript ausführen:

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="gke-tpu-v6e"
export REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-31b"
export HF_TOKEN="YOUR_HF_TOKEN"
export GCS_BUCKET="YOUR_BUCKET_NAME"

Ersetzen Sie Folgendes:

  • YOUR_PROJECT_ID: die ID Ihres Google Cloud Projekts.
  • YOUR_REGION: die Region, in der Sie Ihren Cluster bereitstellen möchten.
  • YOUR_ZONE: die Zone, in der Sie Ihren Cluster bereitstellen möchten.
  • YOUR_REPOSITORY_NAME: der Name des Artifact Registry Repositorys für Ihre MaxText-Images.
  • YOUR_RESERVATION_NAME: der Name Ihrer Reservierung.
  • YOUR_HF_TOKEN: Ihr Hugging Face-Zugriffstoken.
  • YOUR_BUCKET_NAME: ein weltweit eindeutiger Name für einen Cloud Storage-Bucket.

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 der entsprechenden Anleitung unter Abhängigkeiten installieren in der Cluster Toolkit-Dokumentation.

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

Cluster Toolkit installieren

Installieren Sie das vorgefertigte Bundle für Cluster Toolkit gemäß der Anleitung unter Cluster Toolkit installieren.

MaxText-Container-Image vorbereiten

So bereiten Sie Ihr MaxText-Container-Image vor, einschließlich der Installation der erforderlichen Abhängigkeiten:

  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 ${REPOSITORY_NAME} \
    --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.3 --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}"

Cluster Toolkit-Cluster erstellen

So erstellen und stellen Sie einen Cluster Toolkit-Cluster mit 32 v6e-TPU-Chips bereit:

  1. Erstellen Sie eine benutzerdefinierte IAM-Rolle (Identity and Access Management) mit dem Namen 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. Erstellen Sie einen Cloud Storage-Bucket:

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

  3. Standardmäßig hat das Dienstkonto des Cluster-Knotenpools nicht die erforderlichen Berechtigungen, um in Ihren Cloud Storage-Bucket zu schreiben. Damit das Dienstkonto des Knotenpools in Ihren Cloud Storage-Bucket schreiben kann, müssen Sie ihm die Rolle Storage Admin zuweisen. Bearbeiten Sie dazu die Datei gke-tpu-v6e-advanced.yaml und aktualisieren Sie das Modul 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. Stellen Sie Ihren Cluster Toolkit-Cluster mit dem Blueprint gke-tpu-v6e-advanced.yaml bereit und übergeben Sie die erforderlichen Variablen mit dem 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=1,tpu_topology=4x8,authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    --download-dependencies \
    -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. Nachdem Sie den Cluster Toolkit-Cluster erstellt haben, konfigurieren Sie Docker:

    # Configure docker for pulling images
gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker ${REGION}-docker.pkg.dev --quiet

  2. Konfigurieren Sie Ihr Standardprojekt, Ihren Standardcluster und Ihren Standardstandort, um 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}

  3. Führen Sie das folgende Skript aus, 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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

Führen Sie den folgenden Befehl aus, um den Status des Konvertierungsjobs zu prüfen:

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

Trainingsarbeitslast starten

Nach Abschluss der Konvertierung können Sie die SFT-Arbeitslast mit dem folgenden Befehl starten:

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

Führen Sie den folgenden Befehl aus, um den Status des Trainingsjobs zu prüfen:

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

Trainiertes Modell wieder in das Hugging Face-Format konvertieren

Nach Abschluss der Trainingsarbeitslast konvertieren Sie das Modell wieder in das Hugging Face-Format:

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

Führen Sie den folgenden Befehl aus, um den Status des Konvertierungsjobs zu prüfen:

# 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 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/ - though again, it's ~2x the size of the original...

Bereinigen

Löschen Sie die in dieser Anleitung erstellten Ressourcen, um zusätzliche Kosten zu vermeiden.

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}

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.