Google utilizza la tecnologia AI per tradurre i contenuti nella tua lingua preferita. Le traduzioni generate dall'AI potrebbero contenere errori.

Ottimizzare Gemma 4 su un cluster GKE A4

GPU cluster

Questo tutorial mostra come ottimizzare un Gemma 4 Gemma 4 su un cluster Google Kubernetes Engine (GKE) multi-GPU a host singolo Google Cloud. Questo cluster utilizza una singola istanza di macchina virtuale (VM) A4 con 8 GPU NVIDIA B200 collegate.

I due processi principali descritti in questo tutorial sono i seguenti:

Esegui il deployment di un cluster GKE ad alte prestazioni utilizzando GKE Autopilot. Nell'ambito di questo deployment, utilizzi Cloud Build per creare un'immagine container personalizzata con il software necessario preinstallato.
Dopo il deployment del cluster, esegui un job di ottimizzazione che utilizza Fully Sharded Data Parallel (FSDP) per partizionare il modello su tutte le otto GPU di un singolo nodo A4 (a4-highgpu-8g). Il job utilizza l'insieme di script che accompagnano questo tutorial e sfrutta la libreria Hugging Face Accelerate.

Questo tutorial è destinato a engineer di machine learning (ML), ricercatori, amministratori e operatori di piattaforme e a specialisti di dati e AI interessati al deployment di cluster GKE on Google Cloud to train 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 Fully Sharded Data Parallel (FSDP).
Monitora il job.
Esegui 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 Google Cloud utenti potrebbero avere diritto a una prova senza costi.

Prima di iniziare

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 hai già questa autorizzazione tramite il ruolo Proprietario (roles/owner). In caso contrario, puoi ottenere questa autorizzazione tramite il ruolo Amministratore Service Usage (roles/serviceusage.serviceUsageAdmin). Scopri come concedere i ruoli.
```
gcloud services enable compute.googleapis.com container.googleapis.com artifactregistry.googleapis.com cloudbuild.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com
```

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

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

Concedi i ruoli necessari al account di servizio predefinito per eseguire il workload di ottimizzazione:

ROLES=("roles/artifactregistry.writer"
       "roles/cloudbuild.builds.builder"
       "roles/logging.logWriter"
       "roles/monitoring.metricWriter"
       "roles/monitoring.viewer"
       "roles/stackdriver.resourceMetadata.writer"
       "roles/storage.objectViewer")
for role in "${ROLES[@]}"; do
  gcloud projects add-iam-policy-binding YOUR_PROJECT_ID \
      --member="serviceAccount:${PROJECT_NUMBER}-compute@developer.gserviceaccount.com" \
      --role="${role}" 1>/dev/null
done
unset ROLES

Per verificare che i ruoli siano stati concessi, esegui questo comando:

gcloud projects get-iam-policy YOUR_PROJECT_ID \
    --flatten="bindings[].members" \
    --filter="bindings.members:serviceAccount:${PROJECT_NUMBER}-compute@developer.gserviceaccount.com" \
    --format="table(bindings.role)"

Crea credenziali di autenticazione locali per il tuo account utente:
```
gcloud auth application-default login
```
Nota: se utilizzi una shell locale e un provider di identità (IdP) esterno e si verifica un errore di autenticazione dopo aver eseguito il comando precedente, accedi a gcloud CLI con la tua identità federata.

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:

Amministratore Kubernetes Engine (roles/container.admin)
Amministratore Compute (roles/compute.admin)
Amministratore Storage (roles/storage.admin)
Amministratore Artifact Registry (roles/artifactregistry.admin)
Editor Cloud Build (roles/cloudbuild.builds.editor)
Utente dell'account di servizio (roles/iam.serviceAccountUser)
Amministratore Service Account (roles/iam.serviceAccountAdmin)
Amministratore IAM progetto (roles/resourcemanager.projectIamAdmin)
Amministratore Service Usage (roles/serviceusage.serviceUsageAdmin)

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.

Accedi a Gemma 4 utilizzando Hugging Face

Per utilizzare Hugging Face per accedere a Gemma 4:

Accedi a Hugging Face
Crea un token di accesso write di Hugging Face.
Fai clic su Il tuo profilo > Impostazioni > Token di accesso > +Crea nuovo token
Copia e salva il valore del token di 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 Google Cloud progetto 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 tuo 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 la creazione del cluster sia stata completata, vai a Cluster Kubernetes nella Google Cloud console.

Crea un secret di Kubernetes per le credenziali di Hugging Face

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

Configura kubectl per comunicare con il cluster GKE:

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

Crea un secret di Kubernetes per archiviare il token di 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:

Crea script di workload.
Utilizza Docker e Cloud Build per creare un container di ottimizzazione.

Crea script di workload

Per creare gli script utilizzati dal workload di ottimizzazione:

Crea una directory per gli script del workload. Utilizza questa directory come directory di lavoro.
```
mkdir llm-finetuning-gemma
cd llm-finetuning-gemma
```

Crea il file cloudbuild.yaml da utilizzare con Google Cloud Build. Questo file crea il container del workload 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

Crea un file Dockerfile per eseguire il job di ottimizzazione:

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

Crea il file accel_fsdp_gemma4_config.yaml. Questo file di configurazione indica a Hugging Face Accelerate di suddividere 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

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

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 ottimizzazione

Crea un repository Docker di Artifact Registry:

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

Nella directory llm-finetuning-gemma creata in un passaggio precedente, esegui questo comando per creare l'immagine di ottimizzazione ed eseguirne il push su Artifact Registry.
```
gcloud builds submit . \
    --substitutions=_ARTIFACT_REPO_LOCATION="${ARTIFACT_REPO_LOCATION}"
```
Nota: la creazione e il push dell'immagine potrebbero richiedere 20 minuti o più. Per visualizzare lo stato di avanzamento della build, apri il link fornito dal comando con l'output Logs are available at.

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 ottimizzazione

Per avviare il workload di ottimizzazione:

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.
Monitora il job eseguendo questo comando:
```
watch kubectl get pods
```
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 lo ottimizza su tutte le otto GPU. Il download richiede circa cinque minuti. Al termine del download, il processo di ottimizzazione richiede circa 45 minuti. Il job salva un checkpoint ogni 20 passaggi ed esegue il push dell'adattatore Low-Rank Adaptation (LoRA) ottimizzato in un repository Hugging Face privato denominato gemma-31b-text-to-sql.

Al termine del job, recupera l'URL del modello ottimizzato su Hugging Face e poi aprilo nel browser:

HF_USERNAME=$(curl -s -H "Authorization: Bearer ${HF_TOKEN}" \
    https://huggingface.co/api/whoami-v2 | jq -r .name)
echo "https://huggingface.co/${HF_USERNAME}/gemma-31b-text-to-sql"

Monitora il workload

Puoi monitorare l'utilizzo delle GPU nel cluster GKE per verificare che il job di ottimizzazione sia in esecuzione in modo efficiente. Per farlo, apri questo 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 workload, puoi vedere quanto segue:

Utilizzo GPU: per un job di ottimizzazione integro, puoi prevedere che l'utilizzo di tutte le 8 GPU aumenti e si stabilizzi a un livello elevato durante l'addestramento.
Durata del job: il completamento del job dovrebbe richiedere circa 45 minuti sul cluster A4 specificato.

Esegui la pulizia

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

Elimina le risorse

Per eliminare il job di ottimizzazione, esegui questo comando:
```
kubectl delete job finetune-job
```

Per eliminare il cluster GKE, esegui questo comando:

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

Per eliminare il repository Artifact Registry, esegui questo comando:

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