Ottimizza Gemma 3 su un cluster GKE A4

Questo tutorial mostra come ottimizzare un modello linguistico di grandi dimensioni (LLM) Gemma 3 su un cluster GKE multi-nodo e multi-GPU su Google Cloud. Questo cluster utilizza un'istanza di macchina virtuale (VM) A4 con 8 GPU NVIDIA B200.

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 3 utilizzando Hugging Face.

  • Prepara l'ambiente.

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

  • Ottimizza il modello Gemma 3 utilizzando la libreria Hugging Face Accelerate con il parallelismo dei dati completamente partizionato (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. Accedi al tuo account Google Cloud . Se non conosci Google Cloud, crea un account per valutare le prestazioni dei nostri prodotti in scenari reali. I nuovi clienti ricevono anche 300 $di crediti senza costi per l'esecuzione, il test e il deployment dei workload.

  2. Installa Google Cloud CLI.

  3. Se utilizzi un provider di identità (IdP) esterno, devi prima accedere a gcloud CLI con la tua identità federata.

  4. Per inizializzare gcloud CLI, esegui questo comando: 

    gcloud init

  5. Crea o seleziona un Google Cloud progetto.

    Ruoli richiesti per selezionare o creare un progetto

    • Seleziona un progetto: la selezione di un progetto non richiede un ruolo IAM specifico. Puoi selezionare qualsiasi progetto per il quale ti è stato concesso un ruolo.
    • Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l'autorizzazione resourcemanager.projects.create. Scopri come concedere i ruoli.

    • Creare un progetto Google Cloud :

      gcloud projects create PROJECT_ID

      Sostituisci PROJECT_ID con un nome per il progetto Google Cloud che stai creando.

    • Seleziona il progetto Google Cloud che hai creato:

      gcloud config set project PROJECT_ID

      Sostituisci PROJECT_ID con il nome del progetto Google Cloud .

  6. Verifica che la fatturazione sia abilitata per il tuo progetto Google Cloud .

  7. Abilita l'API richiesta:

    Ruoli richiesti per abilitare le API

    Per abilitare le API, devi disporre del ruolo IAM Amministratore utilizzo dei servizi (roles/serviceusage.serviceUsageAdmin), che include l'autorizzazione serviceusage.services.enable. Scopri come concedere i ruoli. 

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

  8. Installa Google Cloud CLI.

  9. Se utilizzi un provider di identità (IdP) esterno, devi prima accedere a gcloud CLI con la tua identità federata.

  10. Per inizializzare gcloud CLI, esegui questo comando: 

    gcloud init

  11. Crea o seleziona un Google Cloud progetto.

    Ruoli richiesti per selezionare o creare un progetto

    • Seleziona un progetto: la selezione di un progetto non richiede un ruolo IAM specifico. Puoi selezionare qualsiasi progetto per il quale ti è stato concesso un ruolo.
    • Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l'autorizzazione resourcemanager.projects.create. Scopri come concedere i ruoli.

    • Creare un progetto Google Cloud :

      gcloud projects create PROJECT_ID

      Sostituisci PROJECT_ID con un nome per il progetto Google Cloud che stai creando.

    • Seleziona il progetto Google Cloud che hai creato:

      gcloud config set project PROJECT_ID

      Sostituisci PROJECT_ID con il nome del progetto Google Cloud .

  12. Verifica che la fatturazione sia abilitata per il tuo progetto Google Cloud .

  13. Abilita l'API richiesta:

    Ruoli richiesti per abilitare le API

    Per abilitare le API, devi disporre del ruolo IAM Amministratore utilizzo dei servizi (roles/serviceusage.serviceUsageAdmin), che include l'autorizzazione serviceusage.services.enable. Scopri come concedere i ruoli. 

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

  14. Concedi ruoli al tuo account utente. Esegui il seguente comando una volta per ciascuno dei seguenti ruoli IAM: 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

    Sostituisci quanto segue:

    • PROJECT_ID: il tuo ID progetto.
    • USER_IDENTIFIER: l'identificatore del tuo account utente . Ad esempio: myemail@example.com.
    • ROLE: il ruolo IAM che concedi al tuo account utente.
  15. Abilita il account di servizio predefinito per il tuo progetto Google Cloud : 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --project=PROJECT_ID

    Sostituisci PROJECT_NUMBER con il numero del progetto. Per rivedere il numero del progetto, consulta Recuperare un progetto esistente.

  16. Concedi il ruolo Editor (roles/editor) al service account predefinito: 
    gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \
    --role=roles/editor
  17. Crea le credenziali di autenticazione locale per il tuo account utente: 
    gcloud auth application-default login
  18. Attiva OS Login per il tuo progetto: 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  19. Accedi o crea un account Hugging Face.

Accedere a Gemma 3 utilizzando Hugging Face

Per utilizzare Hugging Face per accedere a Gemma 3:

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

prepara l'ambiente

Per preparare l'ambiente, imposta quanto segue:

gcloud config set project PROJECT_NAME
gcloud config set billing/quota_project PROJECT_NAME
export RESERVATION=YOUR_RESERVATION_ID
export PROJECT_ID=$(gcloud config get project)
export REGION=CLUSTER_REGION
export CLUSTER_NAME=CLUSTER_NAME
export HF_TOKEN=YOUR_TOKEN
export NETWORK=default

Sostituisci quanto segue:

  • PROJECT_NAME: il nome del Google Cloud progetto in cui vuoi creare il cluster GKE.

  • YOUR_RESERVATION_ID: l'identificatore della capacità riservata.

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

  • CLUSTER_NAME: il nome del cluster GKE da creare.

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

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=${REGION} \
    --release-channel=rapid

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=$REGION

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

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
    --location=${REGION}
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 di workload.

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

Crea script del workload

Per creare gli script utilizzati dal carico di lavoro di fine tuning:

  1. Crea una directory per gli script del carico di lavoro. Utilizza 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', 'us-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-gpu:1.0.0', '.' ]
images:
- 'us-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-gpu:1.0.0'

  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
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. Crea il file accel_fsdp_gemma3_config.yaml. Questo file di configurazione indica a Hugging Face Accelerate di dividere il job di ottimizzazione su più GPU.

    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: 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_gemma3_config.yaml"
        - "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
        volumeMounts:
        - mountPath: /dev/shm
          name: dshm
      volumes:
      - name: dshm
        emptyDir:
          medium: Memory
      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, 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-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=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-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=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 ---
    # The `create_conversation` function is no longer needed.
    # The SFTTrainer will use the `formatting_func` to apply the chat template.
    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 = (
        ""
        ""
    )
    tokenizer.chat_template = gemma_chat_template
    # --- 4. Define the Formatting Function ---
    # This function will be used by the SFTTrainer to format each sample
    # from the dataset into the correct chat template format.
    def formatting_func(example):
        # The create_conversation logic is now implicitly handled by this.
        # We need to construct the messages list here.
        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"][0], context=example["sql_context"][0])},
            {"role": "assistant", "content": example["sql"][0]}
        ]
        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
    # We'll be loading this model full precision because we're planning to do FSDP
    # Load the base model with quantization
    print("Loading base model...")
    model = AutoModelForCausalLM.from_pretrained(
        args.model_id,
        config=config,
        attn_implementation="eager",
        torch_dtype=torch_dtype_obj,
    )

    # Prepare the model for k-bit training
    model = prepare_model_for_kbit_training(model)
    # Configure LoRA.
    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",
    )
    # Apply the PEFT config to the model
    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,
        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_ratio=0.03,
        lr_scheduler_type="constant",
        push_to_hub=True,
        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}")
    model.cpu()
    trainer.save_model(args.output_dir)
    torch.distributed.destroy_process_group()
if __name__ == "__main__":
    main()

Utilizza Docker e Cloud Build per creare un container di perfezionamento

  1. Crea un repository Docker di Artifact Registry:

    gcloud artifacts repositories create gemma  \
    --project=${PROJECT_ID} \
    --repository-format=docker \
    --location=us \
    --description="Gemma Repo"

  2. Nella directory llm-finetuning-gemma creata in un passaggio precedente, esegui questo comando per creare il container di perfezionamento e inviarlo ad Artifact Registry.

     gcloud builds submit .

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

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

Avvia il carico di lavoro 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 cluster in modalità GKE Autopilot, potrebbero essere necessari alcuni minuti per avviare il nodo abilitato per la GPU.

  2. Monitora il job eseguendo questo comando:

    ewatch kubectl get pods

  3. Controlla i log del job eseguendo questo comando:

    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 e otto le GPU. Il completamento del download richiede circa 5 minuti. Al termine del download, la procedura di perfezionamento richiede circa due ore e 30 minuti.

Monitorare 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:

https://console.cloud.google.com/kubernetes/clusters/details/us-central1/[CLUSTER_NAME]/observability?mods=monitoring_api_prod&project=[YOUR_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 messa a punto 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 completamento del job dovrebbe richiedere circa 10 minuti sul cluster A4 specificato.

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questo tutorial, elimina il progetto che contiene le risorse oppure mantieni il progetto ed elimina le singole risorse.

Elimina il progetto

Elimina un progetto Google Cloud :

gcloud projects delete PROJECT_ID

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=$REGION

Passaggi successivi