Ottimizzare Gemma 4 su un cluster GKE A4

Questo tutorial mostra come ottimizzare un modello linguistico di grandi dimensioni (LLM) Gemma 4 su un cluster Google Kubernetes Engine (GKE) multihost e multi-GPU su Google Cloud. Questo cluster utilizza una singola istanza di macchina virtuale (VM) A4 con 8 GPU NVIDIA B200 collegate.

Le due procedure principali descritte in questo tutorial sono le seguenti:

  1. Esegui il deployment di un cluster GKE ad alte prestazioni utilizzando GKE Autopilot. Nell'ambito di questo deployment, crei un'immagine VM personalizzata con il software necessario preinstallato.
  2. Dopo il deployment del cluster, esegui un job di ottimizzazione distribuita utilizzando il set di script che accompagnano questo tutorial. Il job utilizza la libreria Hugging Face Accelerate.

Questo tutorial è rivolto a machine learning engineer, ricercatori, amministratori e operatori di piattaforme e a specialisti di dati e AI interessati al deployment di cluster GKE su Google Cloud per addestrare LLM.

Obiettivi

  • Accedi al modello Gemma 4 utilizzando Hugging Face.

  • Prepara l'ambiente.

  • Crea ed esegui il deployment di un cluster GKE A4.

  • Ottimizza il modello Gemma 4 utilizzando la libreria Hugging Face Accelerate con il parallelismo dei dati completamente partizionati (FSDP).

  • Monitorare il job.

  • Eseguire la pulizia.

Costi

In questo documento vengono utilizzati i seguenti componenti fatturabili di Google Cloud:

Per generare una stima dei costi in base all'utilizzo previsto, utilizza il calcolatore prezzi.

I nuovi utenti di Google Cloud potrebbero avere diritto a una prova senza costi.

Prima di iniziare

  1. Abilita le API richieste, se non sono già abilitate:

    Ruoli richiesti per abilitare le API

    Per abilitare le API, devi disporre dell'autorizzazione serviceusage.services.enable. Se hai creato il progetto, probabilmente disponi già di questa autorizzazione tramite il ruolo Proprietario (roles/owner). In caso contrario, puoi ottenere questa autorizzazione tramite il ruolo Amministratore utilizzo servizi (roles/serviceusage.serviceUsageAdmin). Scopri come concedere i ruoli. 

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

  2. Abilita il account di servizio predefinito per il tuo progetto Google Cloud :

    export PROJECT_NUMBER="$(gcloud projects describe "YOUR_PROJECT_ID" --format "value(project_number)")"
gcloud iam service-accounts enable "${PROJECT_NUMBER}-compute@developer.gserviceaccount.com" \
    --project=YOUR_PROJECT_ID

  3. Concedi i ruoli necessari al account di servizio predefinito per eseguire il carico di lavoro di ottimizzazione:

    ROLES=("roles/aiplatform.user"
       "roles/container.developer"
       "roles/storage.objectUser")
for role in "${ROLES[@]}"; do
  gcloud projects add-iam-policy-binding YOUR_PROJECT_ID \
      --member="serviceAccount:${PROJECT_NUMBER}-compute@developer.gserviceaccount.com" \
      --role="${role}"
done
unset ROLES

  4. Crea credenziali di autenticazione locali per il tuo account utente:

    gcloud auth application-default login

  5. Attiva OS Login per il tuo progetto:

    gcloud compute project-info add-metadata \
    --metadata=enable-oslogin=TRUE \
    --project=YOUR_PROJECT_ID

Per ottenere le autorizzazioni necessarie per completare questo tutorial, chiedi all'amministratore di concederti i seguenti ruoli IAM nel progetto:

Per saperne di più sulla concessione dei ruoli, consulta Gestisci l'accesso a progetti, cartelle e organizzazioni.

Potresti anche riuscire a ottenere le autorizzazioni richieste tramite i ruoli personalizzati o altri ruoli predefiniti.

Accedere a Gemma 4 utilizzando Hugging Face

Per utilizzare Hugging Face per accedere a Gemma 4:

  1. Accedere a Hugging Face
  2. Crea un token di accesso write di Hugging Face.
    Fai clic su Il tuo profilo > Impostazioni > Token di accesso > +Crea nuovo token.
  3. Copia e salva il valore del token write access. Lo utilizzerai più avanti in questo tutorial.

prepara l'ambiente

Per preparare l'ambiente, imposta quanto segue:

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"

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

Sostituisci quanto segue:

  • YOUR_PROJECT_ID: l'ID del progetto Google Cloud in cui vuoi creare il cluster GKE.

  • YOUR_CLUSTER_NAME: il nome del cluster GKE da creare.

  • YOUR_REGION: la regione in cui vuoi creare il cluster GKE. Puoi creare il cluster solo nella regione in cui esiste la prenotazione.

  • YOUR_RESERVATION_NAME: l'identificatore della capacità riservata.

  • YOUR_HF_TOKEN: il token di accesso a Hugging Face che hai creato nella sezione precedente.

  • YOUR_ARTIFACT_REGISTRY_LOCATION: la Google Cloud regione in cui vuoi creare il repository Artifact Registry. Per ridurre al minimo la latenza di pull delle immagini, ti consigliamo di utilizzare la stessa regione utilizzata per YOUR_REGION.

Crea un cluster GKE in modalità Autopilot

Per creare un cluster GKE in modalità Autopilot, esegui questo comando:

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

La creazione del cluster GKE potrebbe richiedere un po' di tempo. Per verificare che Google Cloud abbia terminato la creazione del cluster, vai a Cluster Kubernetes nella console Google Cloud .

Crea un secret Kubernetes per le credenziali di Hugging Face

Per creare un secret Kubernetes per le credenziali di Hugging Face:

  1. Configura kubectl per comunicare con il cluster GKE:

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

  2. Crea un secret Kubernetes per archiviare il token Hugging Face:

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

Prepara il workload

Per preparare il workload:

  1. Crea script del workload.

  2. Utilizza Docker e Cloud Build per creare un container di perfezionamento.

Crea script del workload

Per creare gli script utilizzati dal tuo workload di fine tuning:

  1. Crea una directory per gli script del workload. Usa questa directory come directory di lavoro.

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

  2. Crea il file cloudbuild.yaml per utilizzare Google Cloud Build. Questo file crea il container del carico di lavoro e lo archivia in Artifact Registry:

    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. Crea un file Dockerfile per eseguire il job di perfezionamento:

    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. Crea il file accel_fsdp_gemma4_config.yaml. Questo file di configurazione indica a Hugging Face Accelerate di dividere il job di ottimizzazione tra le otto GPU locali sul singolo host utilizzando FSDP:

    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: false
  fsdp_offload_params: false
  fsdp_reshard_after_forward: true
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: Gemma4TextDecoderLayer
  fsdp_version: 2
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. Crea il file 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"
        - "20"
        - "--push_to_hub"
        resources:
          limits:
            nvidia.com/gpu: "8"
            memory: "1200Gi"
            ephemeral-storage: "350Gi"
          requests:
            nvidia.com/gpu: "8"
            memory: "1200Gi"
            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. Crea il file finetune.py:

    import torch
import argparse
import subprocess
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    AutoConfig,
)
from peft import LoraConfig
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=lambda x: str(x).lower() == "true",
        default=False,
        help="Trust remote code when loading tokenizer")
    parser.add_argument(
        "--use_fast", type=lambda x: str(x).lower() == "true",
        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=lambda x: str(x).lower() == "true",
        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,
    )

    peft_config = LoraConfig(
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        r=args.lora_r,
        bias="none",
        # Train only the text decoder projections; skip vision/audio modules.
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ],
        exclude_modules=(
            r".*(vision_tower|audio_tower|embed_vision|embed_audio).*"),
        task_type="CAUSAL_LM",
    )
    # --- 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,
        hub_private_repo=args.hub_private_repo,
        report_to="tensorboard",
    )
    # --- 7. Create Trainer and Start Training ---
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        peft_config=peft_config,
        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}")
    if trainer.is_fsdp_enabled:
        trainer.accelerator.state.fsdp_plugin.set_state_dict_type(
            "FULL_STATE_DICT")
    trainer.save_model(args.output_dir)
    if torch.distributed.is_initialized():
        torch.distributed.destroy_process_group()

if __name__ == "__main__":
    main()

Utilizza Docker e Cloud Build per creare un container di fine tuning

  1. Crea un repository Docker Artifact Registry:

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

  2. Nella directory llm-finetuning-gemma creata in un passaggio precedente, esegui questo comando per creare l'immagine di perfezionamento e inviarla ad Artifact Registry.

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

  3. Esporta l'URL dell'immagine. Lo utilizzerai in un passaggio successivo di questo tutorial:

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

Avvia il workload di perfezionamento

Per avviare il workload di perfezionamento:

  1. Applica il manifest di ottimizzazione per creare il job di ottimizzazione:

    envsubst < finetune.yaml | kubectl apply -f -

    Poiché utilizzi i cluster in modalità GKE Autopilot, l'avvio del nodo abilitato per la GPU potrebbe richiedere alcuni minuti.

  2. Monitora il job eseguendo questo comando:

    watch kubectl get pods

  3. Dopo l'esecuzione dei pod, controlla i log del job:

    kubectl logs job.batch/finetune-job -f

    La risorsa job scarica i dati del modello e poi esegue il fine-tuning del modello su tutte le otto GPU. Il completamento del download richiede circa cinque minuti. Al termine del download, la procedura di perfezionamento richiede circa 2 ore e 30 minuti.

Monitora il workload

Puoi monitorare l'utilizzo delle GPU nel tuo cluster GKE per verificare che il job di ottimizzazione venga eseguito in modo efficiente. Per farlo, apri il seguente link nel 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\"))"

Quando monitori il carico di lavoro, puoi visualizzare quanto segue:

  • Utilizzo delle GPU: per un job di perfezionamento corretto, puoi aspettarti di vedere l'utilizzo di tutte le 8 GPU aumentare e stabilizzarsi a un livello elevato durante l'addestramento.
  • Durata del job: il job dovrebbe richiedere circa 2 ore e 30 minuti per essere completato sul cluster A4 specificato.

Esegui la pulizia

Per evitare addebiti aggiuntivi, elimina le risorse create durante questo tutorial.

Eliminare le risorse

  1. Per eliminare il job di perfezionamento, esegui questo comando:

    kubectl delete job finetune-job

  2. Per eliminare il cluster GKE, esegui questo comando:

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

Passaggi successivi