Gemma 4 in einem A4-GKE-Cluster abstimmen

Über Hugging Face auf Gemma 4 zugreifen

So greifen Sie mit Hugging Face auf Gemma 4 zu:

1. Bei Hugging Face anmelden
2. Erstellen Sie ein Hugging Face-Zugriffstoken für write.
   Klicken Sie auf 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_REGION"
export RESERVATION="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export ARTIFACT_REPO_LOCATION="YOUR_ARTIFACT_REGISTRY_LOCATION"
export NETWORK="default"

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

Ersetzen Sie Folgendes:

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

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

• YOUR_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.

• YOUR_RESERVATION_NAME: Die Kennung für Ihre reservierte Kapazität.

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

• YOUR_ARTIFACT_REGISTRY_LOCATION: die Google Cloud Region, in der Sie Ihr Artifact Registry-Repository erstellen möchten. Um die Latenz beim Abrufen von Images zu minimieren, empfehlen wir, dieselbe Region wie für YOUR_REGION zu verwenden.

GKE-Cluster im Autopilot-Modus erstellen

Führen Sie den folgenden Befehl aus, um einen GKE-Cluster im Autopilot-Modus zu erstellen:

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

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.

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-gpu:2.0.0'
     - '.'
   images:
   - '$_ARTIFACT_REPO_LOCATION-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-gpu:2.0.0'
   options:
     logging: CLOUD_LOGGING_ONLY

3. Erstellen Sie eine Dockerfile-Datei, um den Job für die Feinabstimmung auszuführen:

   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 && \
       rm -rf /var/lib/apt/lists/*
   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_gemma4_config.yaml /accel_fsdp_gemma4_config.yaml
   CMD accelerate launch --config_file accel_fsdp_gemma4_config.yaml finetune.py

4. Erstellen Sie die Datei accel_fsdp_gemma4_config.yaml. Diese Konfigurationsdatei weist Hugging Face Accelerate an, den Abstimmungsjob 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_backward_prefetch: BACKWARD_PRE
     fsdp_cpu_ram_efficient_loading: true
     fsdp_forward_prefetch: true
     fsdp_offload_params: false
     fsdp_reshard_after_forward: true
     fsdp_sharding_strategy: FULL_SHARD
     fsdp_state_dict_type: FULL_STATE_DICT
     fsdp_sync_module_states: true
     fsdp_transformer_layer_cls_to_wrap: Gemma4TextDecoderLayer
     fsdp_version: 1
   machine_rank: 0
   main_training_function: main
   mixed_precision: bf16
   num_machines: 1
   num_processes: 8
   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: batch/v1
   kind: Job
   metadata:
     name: finetune-job
     namespace: default
   spec:
     backoffLimit: 2
     template:
       metadata:
         annotations:
           kubectl.kubernetes.io/default-container: finetuner
       spec:
         terminationGracePeriodSeconds: 600
         containers:
         - name: finetuner
           image: $IMAGE_URL
           command: ["accelerate","launch"]
           args:
           - "--config_file"
           - "accel_fsdp_gemma4_config.yaml"
           - "finetune.py"
           - "--model_id"
           - "google/gemma-4-31B-it"
           - "--output_dir"
           - "gemma-31b-text-to-sql"
           - "--per_device_train_batch_size"
           - "2"
           - "--gradient_accumulation_steps"
           - "32"
           - "--num_train_epochs"
           - "3"
           - "--learning_rate"
           - "1e-5"
           - "--save_strategy"
           - "steps"
           - "--save_steps"
           - "100"
           resources:
             limits:
               nvidia.com/gpu: "8"
               memory: "1000Gi"
               ephemeral-storage: "350Gi"
             requests:
               nvidia.com/gpu: "8"
               memory: "1000Gi"
               ephemeral-storage: "350Gi"
           env:
           - name: HF_TOKEN
             valueFrom:
               secretKeyRef:
                 name: hf-secret
                 key: hf_api_token
           volumeMounts:
           - mountPath: /dev/shm
             name: dshm
         volumes:
         - name: dshm
           emptyDir:
             medium: Memory
             sizeLimit: 64Gi
         nodeSelector:
           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
         restartPolicy: OnFailure

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():
       parser = argparse.ArgumentParser()
       parser.add_argument(
           "--model_id", type=str, default="google/gemma-4-31b-it",
           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=bool, default="False",
           help="Trust remote code when loading tokenizer")
       parser.add_argument(
           "--use_fast", type=bool, 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-31b-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=32,
           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_length", type=int, default=1024,
           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=2,
           help="Batch size per device during training")
       parser.add_argument(
           "--gradient_accumulation_steps", type=int, default=32,
           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=bool, 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,
       )
       if tokenizer.pad_token is None:
           tokenizer.pad_token = tokenizer.eos_token
       # --- 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": "system", "content": system_message},
               {
                   "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="sdpa",
           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="all-linear",
           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_length=args.max_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,
           label_names=["domain"],
           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_steps=0.03,
           lr_scheduler_type="constant",
           push_to_hub=args.push_to_hub,
           report_to="tensorboard",
       )
       # --- 7. Create Trainer and Start Training ---
       trainer = SFTTrainer(
           model=model,
           args=training_args,
           train_dataset=dataset["train"],
           eval_dataset=dataset["test"],
           processing_class=tokenizer,
           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. 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 in Artifact Registry zu übertragen.

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

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

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

Fine-Tuning-Arbeitslast 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 < 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 nach dem Ausführen der Pods die Logs des Jobs:

   kubectl logs job.batch/finetune-job -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 Fine-Tuning-Job 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 auf dem angegebenen A4-Cluster sollte etwa 2,5 Stunden dauern.