Gemma 3 in einem A4-GKE-Cluster mit mehreren Hosts abstimmen

In dieser Anleitung wird beschrieben, wie Sie ein Gemma 3-LLM (Large Language Model) in einem GKE Autopilot-Cluster mit mehreren Hosts auf Google Cloudabstimmen. Dieser Cluster verwendet zwei virtuelle A4-Maschinen (VM) mit insgesamt 16 NVIDIA B200-GPUs.

Die beiden in dieser Anleitung beschriebenen Hauptprozesse sind:

  1. Stellen Sie einen leistungsstarken GKE-Cluster mit mehreren Hosts mit GKE Autopilot bereit. Im Rahmen dieser Bereitstellung erstellen Sie ein benutzerdefiniertes VM-Image mit der erforderlichen vorinstallierten Software.
  2. Nachdem der Cluster bereitgestellt wurde, führen Sie einen verteilten Feinabstimmungsjob mit den Skripts aus, die in dieser Anleitung enthalten sind. Für den Job wird die Hugging Face Accelerate-Bibliothek verwendet.

Diese Anleitung richtet sich an Entwickler von maschinellem Lernen (ML), Forscher, Plattformadministratoren und ‑operatoren sowie an Daten- und KI-Spezialisten, die daran interessiert sind, GKE-Cluster auf Google Cloud bereitzustellen, um LLMs zu trainieren.

Ziele

  • Über Hugging Face auf das Gemma 3-Modell zugreifen

  • Bereiten Sie Ihre Umgebung vor.

  • A4-GKE-Cluster erstellen und bereitstellen

  • Optimieren Sie das Gemma 3-Modell mit der Hugging Face Accelerate-Bibliothek mit vollständig fragmentierter Datenparallelität (Fully Sharded Data Parallelism, FSDP).

  • den Job überwachen

  • bereinigen.

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.

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

Hinweis

  1. Melden Sie sich in Ihrem Google Cloud -Konto an. Wenn Sie mit Google Cloudnoch nicht vertraut sind, erstellen Sie ein Konto, um die Leistungsfähigkeit unserer Produkte in der Praxis sehen und bewerten zu können. Neukunden erhalten außerdem ein Guthaben von 300 $, um Arbeitslasten auszuführen, zu testen und bereitzustellen.

  2. Installieren Sie die Google Cloud CLI.

  3. Wenn Sie einen externen Identitätsanbieter (IdP) verwenden, müssen Sie sich zuerst mit Ihrer föderierten Identität in der gcloud CLI anmelden.

  4. Führen Sie den folgenden Befehl aus, um die gcloud CLI zu initialisieren: 

    gcloud init

  5. Erstellen Sie ein Google Cloud Projekt oder wählen Sie eines aus.

    Rollen, die zum Auswählen oder Erstellen eines Projekts erforderlich sind

    • Projekt auswählen: Für die Auswahl eines Projekts ist keine bestimmte IAM-Rolle erforderlich. Sie können jedes Projekt auswählen, für das Ihnen eine Rolle zugewiesen wurde.
    • Projekt erstellen: Zum Erstellen eines Projekts benötigen Sie die Rolle „Projektersteller“ (roles/resourcemanager.projectCreator), die die Berechtigung resourcemanager.projects.create enthält. Weitere Informationen zum Zuweisen von Rollen

    • So erstellen Sie ein Google Cloud Projekt:

      gcloud projects create PROJECT_ID

      Ersetzen Sie PROJECT_ID durch einen Namen für das Google Cloud Projekt, das Sie erstellen.

    • Wählen Sie das von Ihnen erstellte Google Cloud Projekt aus:

      gcloud config set project PROJECT_ID

      Ersetzen Sie PROJECT_ID durch den Namen Ihres Projekts in Google Cloud .

  6. Prüfen Sie, ob die Abrechnung für Ihr Google Cloud Projekt aktiviert ist.

  7. Aktivieren Sie die erforderlichen APIs:

    Rollen, die zum Aktivieren von APIs erforderlich sind

    Zum Aktivieren von APIs benötigen Sie die Berechtigung serviceusage.services.enable. Wenn Sie das Projekt erstellt haben, haben Sie diese Berechtigung wahrscheinlich bereits über die Rolle „Inhaber“ (roles/owner). Andernfalls können Sie diese Berechtigung über die Rolle „Service Usage-Administrator“ (roles/serviceusage.serviceUsageAdmin) erhalten. Informationen zum Zuweisen von Rollen 

    gcloud services enable compute.googleapis.com container.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com

  8. Installieren Sie die Google Cloud CLI.

  9. Wenn Sie einen externen Identitätsanbieter (IdP) verwenden, müssen Sie sich zuerst mit Ihrer föderierten Identität in der gcloud CLI anmelden.

  10. Führen Sie den folgenden Befehl aus, um die gcloud CLI zu initialisieren: 

    gcloud init

  11. Erstellen Sie ein Google Cloud Projekt oder wählen Sie eines aus.

    Rollen, die zum Auswählen oder Erstellen eines Projekts erforderlich sind

    • Projekt auswählen: Für die Auswahl eines Projekts ist keine bestimmte IAM-Rolle erforderlich. Sie können jedes Projekt auswählen, für das Ihnen eine Rolle zugewiesen wurde.
    • Projekt erstellen: Zum Erstellen eines Projekts benötigen Sie die Rolle „Projektersteller“ (roles/resourcemanager.projectCreator), die die Berechtigung resourcemanager.projects.create enthält. Weitere Informationen zum Zuweisen von Rollen

    • So erstellen Sie ein Google Cloud Projekt:

      gcloud projects create PROJECT_ID

      Ersetzen Sie PROJECT_ID durch einen Namen für das Google Cloud Projekt, das Sie erstellen.

    • Wählen Sie das von Ihnen erstellte Google Cloud Projekt aus:

      gcloud config set project PROJECT_ID

      Ersetzen Sie PROJECT_ID durch den Namen Ihres Projekts in Google Cloud .

  12. Prüfen Sie, ob die Abrechnung für Ihr Google Cloud Projekt aktiviert ist.

  13. Aktivieren Sie die erforderlichen APIs:

    Rollen, die zum Aktivieren von APIs erforderlich sind

    Zum Aktivieren von APIs benötigen Sie die Berechtigung serviceusage.services.enable. Wenn Sie das Projekt erstellt haben, haben Sie diese Berechtigung wahrscheinlich bereits über die Rolle „Inhaber“ (roles/owner). Andernfalls können Sie diese Berechtigung über die Rolle „Service Usage-Administrator“ (roles/serviceusage.serviceUsageAdmin) erhalten. Informationen zum Zuweisen von Rollen 

    gcloud services enable compute.googleapis.com container.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com
    14.

  14. Weisen Sie Ihrem Nutzerkonto Rollen zu. Führen Sie den folgenden Befehl für jede der folgenden IAM-Rollen einmal aus: roles/compute.admin, roles/iam.serviceAccountUser, roles/cloudbuild.builds.editor, roles/artifactregistry.admin, roles/storage.admin, roles/serviceusage.serviceUsageAdmin 

    gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

    Ersetzen Sie Folgendes:

    • PROJECT_ID: Ihre Projekt-ID.
    • USER_IDENTIFIER: Die Kennung für Ihr Nutzerkonto . Beispiel: myemail@example.com
    • ROLE: Die IAM-Rolle, die Sie Ihrem Nutzerkonto zuweisen.
  15. Aktivieren Sie das Standarddienstkonto für Ihr Google Cloud Projekt: 
      export PROJECT_NUMBER="$(gcloud projects describe "PROJECT_ID" --format "value(project_number)")"
  gcloud iam service-accounts enable "${PROJECT_NUMBER}-compute@developer.gserviceaccount.com \
    --project=PROJECT_ID
  16. Weisen Sie dem Standarddienstkonto die Rolle „Bearbeiter“ (roles/editor) zu: 
    gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \
    --role=roles/editor
  17. Erstellen Sie lokale Anmeldedaten zur Authentifizierung für Ihr Nutzerkonto: 
    gcloud auth application-default login
  18. Aktivieren Sie OS Login für Ihr Projekt: 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  19. Melden Sie sich in einem Hugging Face-Konto an oder erstellen Sie ein Konto.

Über Hugging Face auf Gemma 3 zugreifen

So greifen Sie mit Hugging Face auf Gemma 3 zu:

  1. Bei Hugging Face anmelden
  2. Erstellen Sie ein Hugging Face-Zugriffstoken für write.
    Klicken Sie auf Mein Profil > Einstellungen > Zugriffstokens > +Neues Token erstellen.
  3. Kopieren und speichern Sie den write access-Tokenwert. Sie benötigen sie später in dieser Anleitung.

Umgebung vorbereiten

So bereiten Sie Ihre Umgebung vor:

export PROJECT_ID="YOUR_PROJECT_ID"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export CLUSTER_REGION="YOUR_CLUSTER_REGION"
export RESERVATION="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HUGGING_FACE_TOKEN"
export ARTIFACT_REPO_LOCATION="YOUR_ARTIFACT_REGISTRY_LOCATION"
export NUM_NODES="YOUR_NUMBER_OF_NODES"
export NETWORK="default"

gcloud config set project "${PROJECT_ID}"
gcloud config set billing/quota_project "${PROJECT_ID}"

Ersetzen Sie Folgendes:

  • PROJECT_ID: Der Name des Google Cloud Projekts, in dem Sie den GKE-Cluster erstellen möchten.

  • CLUSTER_NAME: der Name des GKE-Cluster, der erstellt werden soll.

  • CLUSTER_REGION: die Region, in der Sie Ihren GKE-Cluster erstellen möchten. Sie können den Cluster nur in der Region erstellen, in der Ihre Reservierung vorhanden ist.

  • RESERVATION: die Kennung für Ihre reservierte Kapazität.

  • HF_TOKEN: Das Hugging Face-Zugriffstoken, das Sie im vorherigen Abschnitt erstellt haben.

  • ARTIFACT_REPO_LOCATION: der Speicherort, an dem Sie Ihr Artifact Registry-Repository erstellen möchten.

  • NUM_NODES: Die Anzahl der Knoten, die der GKE-Cluster haben soll.

  • NETWORK: Das zu verwendende Netzwerk.

GKE-Cluster mit mehreren Hosts im Autopilot-Modus erstellen

Erstellen Sie einen GKE-Cluster mit mehreren Hosts im Autopilot-Modus:

gcloud container clusters create-auto "${CLUSTER_NAME}" \
    --project="${PROJECT_ID}" \
    --location="${CLUSTER_REGION}" \
    --release-channel=rapid

Das Erstellen des GKE-Cluster kann einige Zeit dauern. Rufen Sie in der Google Cloud Console die Seite Kubernetes-Cluster auf, um zu prüfen, ob Google Cloud den Cluster erstellt hat.

Anmeldedaten für Ihren GKE-Cluster abrufen

Konfigurieren Sie kubectl für die Kommunikation mit dem neu erstellten GKE-Cluster, indem Sie den folgenden Befehl ausführen:

gcloud container clusters get-credentials "${CLUSTER_NAME}" \
    --location="${CLUSTER_REGION}"

Kubernetes-Secret für Hugging Face-Anmeldedaten erstellen

So erstellen Sie ein Kubernetes-Secret für Hugging Face-Anmeldedaten:

  1. Konfigurieren Sie kubectl für die Kommunikation mit Ihrem GKE-Cluster:

    gcloud container clusters get-credentials "${CLUSTER_NAME}" \
    --location="${CLUSTER_REGION}"

  2. Erstellen Sie ein Kubernetes-Secret zum Speichern Ihres Hugging Face-Tokens:

    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token="${HF_TOKEN}" \
    --dry-run=client -o yaml | kubectl apply -f -

Arbeitslast vorbereiten

So bereiten Sie Ihre Arbeitslast vor:

  1. Arbeitslastskripts erstellen:

  2. Container für die Feinabstimmung mit Docker und Cloud Build erstellen

Arbeitslastskripts erstellen

So erstellen Sie die Skripts, die von Ihrer Arbeitslast zum Feinabstimmen verwendet werden:

  1. Erstellen Sie ein Verzeichnis für die Arbeitslastskripts. Verwenden Sie dieses Verzeichnis als Arbeitsverzeichnis.

    mkdir llm-finetuning-gemma
cd llm-finetuning-gemma

  2. Erstellen Sie die Datei cloudbuild.yaml, um Google Cloud Build zu verwenden. Mit dieser Datei wird der Container für Ihre Arbeitslast erstellt und in Artifact Registry gespeichert:

    steps:
- name: 'gcr.io/cloud-builders/docker'
  args: [ 'build', '-t', '$_ARTIFACT_REPO_LOCATION-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-multihost-gpu:1.0.0', '.' ]
images:
- '$_ARTIFACT_REPO_LOCATION-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-multihost-gpu:1.0.0'

  3. Erstellen Sie eine Dockerfile-Datei, um die Umgebung zu definieren und die Abhängigkeiten zu installieren, die zum Ausführen des Feinabstimmungsjobs erforderlich sind:

    FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
RUN apt-get update && \
    apt-get -y install python3 python3-dev gcc python3-pip python3-venv git curl vim
RUN python3 -m venv /opt/venv
ENV PATH="/opt/venv/bin:/usr/local/nvidia/bin:$PATH"
ENV LD_LIBRARY_PATH="/usr/local/nvidia/lib64:$LD_LIBRARY_PATH"
RUN pip3 install setuptools wheel packaging ninja
RUN pip3 install torch torchvision torchaudio  --index-url https://download.pytorch.org/whl/cu128

RUN pip3 install \
    transformers==4.53.3 \
    datasets==4.0.0 \
    accelerate==1.9.0 \
    evaluate==0.4.5 \
    bitsandbytes==0.46.1 \
    trl==0.19.1 \
    peft==0.16.0 \
    tensorboard==2.20.0 \
    protobuf==6.31.1 \
    sentencepiece==0.2.0
COPY finetune.py /finetune.py
COPY accel_fsdp_gemma3_config.yaml /accel_fsdp_gemma3_config.yaml
CMD accelerate launch --config_file accel_fsdp_gemma3_config.yaml finetune.py

  4. Erstellen Sie die Datei accel_fsdp_gemma3_config.yaml. Diese Konfigurationsdatei weist Hugging Face Accelerate an, den Optimierungsjob auf mehrere GPUs aufzuteilen.

    compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
enable_cpu_affinity: false
fsdp_config:
  fsdp_activation_checkpointing: false
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false
  fsdp_reshard_after_forward: true
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: Gemma3DecoderLayer
  fsdp_version: 2
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 2
num_processes: 16
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

  5. Erstellen Sie die Datei finetune.yaml.

    apiVersion: resource.k8s.io/v1
kind: ResourceClaimTemplate
metadata:
  name: mrdma
spec:
  spec:
    devices:
      requests:
      - name: mrdma
        exactly:
          deviceClassName: mrdma.google.com
---
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: finetune-jobset
  namespace: default
spec:
  failurePolicy:
    maxRestarts: 2
  replicatedJobs:
  - name: workers
    replicas: 1
    template:
      spec:
        parallelism: ${NUM_NODES}
        completions: ${NUM_NODES}
        backoffLimit: 0
        template:
          metadata:
            annotations:
              kubectl.kubernetes.io/default-container: finetuner
          spec:
            terminationGracePeriodSeconds: 600
            restartPolicy: OnFailure
            nodeSelector:
              cloud.google.com/compute-class: "Accelerator"
              cloud.google.com/gke-accelerator: "nvidia-b200"
              cloud.google.com/reservation-name: ${RESERVATION}
              cloud.google.com/reservation-affinity: "specific"
              cloud.google.com/gke-gpu-driver-version: latest
            containers:
            - name: finetuner
              image: $IMAGE_URL
              command: ["bash", "-c"]
              args:
              - |
                NUM_PROCESSES=$(( ${NUM_NODES} * 8 ))
                accelerate launch \
                  --config_file accel_fsdp_gemma3_config.yaml \
                  --num_machines ${NUM_NODES} \
                  --num_processes ${NUM_PROCESSES} \
                  --machine_rank ${JOB_COMPLETION_INDEX} \
                  --main_process_ip finetune-jobset-workers-0-0.finetune-jobset.default.svc.cluster.local \
                  --main_process_port 29500 \
                  finetune.py \
                  --model_id google/gemma-3-12b-pt \
                  --output_dir gemma-12b-text-to-sql \
                  --per_device_train_batch_size 8 \
                  --gradient_accumulation_steps 8 \
                  --num_train_epochs 3 \
                  --learning_rate 1e-5 \
                  --save_strategy steps \
                  --save_steps 100
              resources:
                limits:
                  nvidia.com/gpu: "8"
              env:
              - name: HF_TOKEN
                valueFrom:
                  secretKeyRef:
                    name: hf-secret
                    key: hf_api_token
              - name: NUM_NODES
                value: "${NUM_NODES}"
              volumeMounts:
              - mountPath: /dev/shm
                name: dshm
            volumes:
            - name: dshm
              emptyDir:
                medium: Memory

  6. Erstellen Sie die Datei finetune.py.

    import torch
import argparse
import subprocess
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig
from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model
from trl import SFTTrainer, SFTConfig
from huggingface_hub import login

def get_args():
    def str2bool(v):
        if isinstance(v, bool):
            return v
        return v.lower() in ("yes", "true", "t", "1")

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_id", type=str, default="google/gemma-3-12b-pt", help="Hugging Face model ID")
    parser.add_argument("--hf_token", type=str, default=None, help="Hugging Face token for private models")
    parser.add_argument("--trust_remote", type=str2bool, default=False, help="Trust remote code when loading tokenizer")
    parser.add_argument("--use_fast", type=str2bool, default=True, help="Determines if a fast Rust-based tokenizer should be used")
    parser.add_argument("--dataset_name", type=str, default="philschmid/gretel-synthetic-text-to-sql", help="Hugging Face dataset name")
    parser.add_argument("--output_dir", type=str, default="gemma-12b-text-to-sql", help="Directory to save model checkpoints")

    # LoRA arguments
    parser.add_argument("--lora_r", type=int, default=16, help="LoRA attention dimension")
    parser.add_argument("--lora_alpha", type=int, default=16, help="LoRA alpha scaling factor")
    parser.add_argument("--lora_dropout", type=float, default=0.05, help="LoRA dropout probability")
    # SFTConfig arguments
    parser.add_argument("--max_seq_length", type=int, default=512, help="Maximum sequence length")
    parser.add_argument("--num_train_epochs", type=int, default=3, help="Number of training epochs")
    parser.add_argument("--per_device_train_batch_size", type=int, default=8, help="Batch size per device during training")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient accumulation steps")
    parser.add_argument("--learning_rate", type=float, default=1e-5, help="Learning rate")
    parser.add_argument("--logging_steps", type=int, default=10, help="Log every X steps")
    parser.add_argument("--save_strategy", type=str, default="steps", help="Checkpoint save strategy")
    parser.add_argument("--save_steps", type=int, default=100, help="Save checkpoint every X steps")
    parser.add_argument("--push_to_hub", action='store_true', help="Push model back up to HF")
    parser.add_argument("--hub_private_repo", type=str2bool, default=True, help="Push to a private repo")
    return parser.parse_args()

def main():
    args = get_args()
    # --- 1. Setup and Login ---
    if args.hf_token:
        login(args.hf_token)
    # --- 2. Create and prepare the fine-tuning dataset ---
    dataset = load_dataset(args.dataset_name, split="train")
    dataset = dataset.shuffle().select(range(12500))
    dataset = dataset.train_test_split(test_size=2500/12500)
    # --- 3. Configure Model and Tokenizer ---
    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8:
        torch_dtype_obj = torch.bfloat16
        torch_dtype_str = "bfloat16"
    else:
        torch_dtype_obj = torch.float16
        torch_dtype_str = "float16"
    tokenizer = AutoTokenizer.from_pretrained(args.model_id, trust_remote_code=args.trust_remote, use_fast=args.use_fast)
    tokenizer.pad_token = tokenizer.eos_token
    gemma_chat_template = (
        "{{ bos_token }}"
        "{% if messages[0]['role'] == 'system' %}{{ messages[0]['content'] }}{% endif %}"
        "{% for message in messages %}"
        "{% if message['role'] == 'user' %}<start_of_turn>user\n{{ message['content'] }}<end_of_turn>\n{% elif message['role'] == 'assistant' %}<start_of_turn>model\n{{ message['content'] }}<end_of_turn>\n{% endif %}"
        "{% endfor %}"
    )
    tokenizer.chat_template = gemma_chat_template
    # --- 4. Define the Formatting Function ---
    def formatting_func(example):
        system_message = "You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA."
        user_prompt = "Given the <USER_QUERY> and the <SCHEMA>, generate the corresponding SQL command to retrieve the desired data, considering the query's syntax, semantics, and schema constraints.\n\n<SCHEMA>\n{context}\n</SCHEMA>\n\n<USER_QUERY>\n{question}\n</USER_QUERY>\n"

        messages = [
            {"role": "user", "content": user_prompt.format(question=example["sql_prompt"], context=example["sql_context"])},
            {"role": "assistant", "content": example["sql"]}
        ]
        return tokenizer.apply_chat_template(messages, tokenize=False)
    # --- 5. Load Model and Apply PEFT ---
    config = AutoConfig.from_pretrained(args.model_id)
    config.use_cache = False
    print("Loading base model...")
    model = AutoModelForCausalLM.from_pretrained(
        args.model_id,
        config=config,
        attn_implementation="eager",
        torch_dtype=torch_dtype_obj,
    )

    model = prepare_model_for_kbit_training(model)
    peft_config = LoraConfig(
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        r=args.lora_r,
        bias="none",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        task_type="CAUSAL_LM",
    )
    print("Applying PEFT configuration...")
    model = get_peft_model(model, peft_config)
    model.print_trainable_parameters()
    # --- 6. Configure Training Arguments ---
    training_args = SFTConfig(
        output_dir=args.output_dir,
        max_seq_length=args.max_seq_length,
        num_train_epochs=args.num_train_epochs,
        per_device_train_batch_size=args.per_device_train_batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate,
        logging_steps=args.logging_steps,
        save_strategy=args.save_strategy,
        save_steps=args.save_steps,
        packing=False,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        optim="adamw_torch",
        fp16=True if torch_dtype_obj == torch.float16 else False,
        bf16=True if torch_dtype_obj == torch.bfloat16 else False,
        max_grad_norm=0.3,
        warmup_ratio=0.03,
        lr_scheduler_type="constant",
        push_to_hub=args.push_to_hub,
        hub_private_repo=args.hub_private_repo,
        report_to="tensorboard",
        dataset_kwargs={
            "add_special_tokens": False,
            "append_concat_token": True,
        }
    )
    # --- 7. Create Trainer and Start Training ---
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        formatting_func=formatting_func,
    )
    print("Starting training...")
    trainer.train()
    print("Training finished.")
    # --- 8. Save the final model ---
    print(f"Saving final model to {args.output_dir}")
    trainer.save_model(args.output_dir)
    if torch.distributed.is_initialized():
        torch.distributed.destroy_process_group()

if __name__ == "__main__":
    main()

Container für die Feinabstimmung mit Docker und Cloud Build erstellen

  1. Artifact Registry-Docker-Repository erstellen:

    gcloud artifacts repositories create gemma \
    --repository-format=docker \
    --location="${ARTIFACT_REPO_LOCATION}" \
    --description="Repository for Gemma fine tuning workload containers" || true

  2. Installieren Sie die benutzerdefinierten Ressourcendefinitionen (CRDs) für JobSet, die für die Orchestrierung von Arbeitslasten mit mehreren Hosts erforderlich sind.

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.12.0/manifests.yaml

    JobSet ist eine Kubernetes-Erweiterung, mit der Gruppen von zusammengehörigen Jobs verwaltet werden. Weitere Informationen zu JobSet finden Sie in der externen JobSet-Dokumentation.

  3. Führen Sie im Verzeichnis llm-finetuning-gemma, das Sie in einem vorherigen Schritt erstellt haben, den folgenden Befehl aus, um das Feinabstimmungs-Image zu erstellen und per Push an Artifact Registry zu übertragen.

    gcloud builds submit . --substitutions=_ARTIFACT_REPO_LOCATION="${ARTIFACT_REPO_LOCATION}"

  4. Exportieren Sie die Bild-URL. Sie benötigen sie in einem späteren Schritt dieser Anleitung:

    export IMAGE_URL="${ARTIFACT_REPO_LOCATION}-docker.pkg.dev/${PROJECT_ID}/gemma/finetune-gemma-gpu:1.0.0"

Arbeitslast für das Feinabstimmen starten

So starten Sie die Arbeitslast für das Fine-Tuning:

  1. Wenden Sie das Manifest für die Feinabstimmung an, um den Job für die Feinabstimmung zu erstellen:

    envsubst '${RESERVATION} ${IMAGE_URL} ${NUM_NODES}' < finetune.yaml | kubectl apply -f -

    Da Sie Cluster im GKE Autopilot-Modus verwenden, kann es einige Minuten dauern, bis Ihr GPU-fähiger Knoten gestartet wird.

  2. Überwachen Sie den Job mit folgendem Befehl:

    watch kubectl get pods

  3. Prüfen Sie die Logs des primären Worker-Pods mit dem folgenden Befehl:

    kubectl logs -l "job-name=finetune-jobset-workers-0" -f

    Die Jobressource lädt die Modelldaten herunter und optimiert das Modell dann auf allen acht GPUs. Der Download dauert etwa fünf Minuten. Nach dem Download dauert das Feinabstimmungsverfahren etwa zweieinhalb Stunden.

Arbeitslast überwachen

Sie können die Nutzung der GPUs in Ihrem GKE-Cluster überwachen, um zu prüfen, ob Ihr Fine-Tuning-Job effizient ausgeführt wird. Öffnen Sie dazu den folgenden Link in Ihrem Browser:

echo "https://console.cloud.google.com/kubernetes/clusters/details/${CLUSTER_REGION}/${CLUSTER_NAME}/observability?mods=monitoring_api_prod&project=${PROJECT_ID}&pageState=("timeRange":("duration":"PT1H"),"nav":("section":"gpu"),"groupBy":("groupByType":"namespacesTop5"))"

Wenn Sie Ihre Arbeitslast überwachen, sehen Sie Folgendes:

  • GPU-Nutzung: Bei einem fehlerfreien Feinabstimmungsjob sollte die Nutzung aller 8 GPUs während des gesamten Trainings ansteigen und sich auf einem hohen Niveau stabilisieren.
  • Job-Dauer: Die Ausführung des Jobs sollte auf dem angegebenen A4-Cluster etwa 10 Minuten dauern.

Bereinigen

Damit Ihrem Google Cloud-Konto die in dieser Anleitung verwendeten Ressourcen nicht in Rechnung gestellt werden, löschen Sie entweder das Projekt, das die Ressourcen enthält, oder Sie behalten das Projekt und löschen die einzelnen Ressourcen.

Ressourcen löschen

  1. Führen Sie den folgenden Befehl aus, um Ihr JobSet zu löschen:

    kubectl delete jobset finetune-jobset

  2. Führen Sie den folgenden Befehl aus, um Ihren GKE-Cluster zu löschen:

    gcloud container clusters delete "${CLUSTER_NAME}" \
    --region="${CLUSTER_REGION}"

  3. Führen Sie den folgenden Befehl aus, um Ihr Artifact Registry-Repository zu löschen:

    gcloud artifacts repositories delete gemma \
    --location="${ARTIFACT_REPO_LOCATION}" \
    --quiet

Projekt löschen

Google Cloud -Projekt löschen:

gcloud projects delete PROJECT_ID

Nächste Schritte