Ajustar o Gemma 3 em um cluster do GKE A4

Neste tutorial, você vai aprender a ajustar um modelo de linguagem grande (LLM) Gemma 3 em um cluster do GKE com vários nós e GPUs no Google Cloud. Esse cluster usa uma instância de máquina virtual (VM) A4, que tem 8 GPUs NVIDIA B200.

Os dois processos principais descritos neste tutorial são os seguintes:

  1. Implante um cluster do GKE de alta performance usando o Autopilot do GKE. Como parte dessa implantação, você cria uma imagem de VM personalizada com o software necessário pré-instalado.
  2. Depois que o cluster for implantado, execute um job de ajuste refinado distribuído usando o conjunto de scripts que acompanham este tutorial. O job usa a biblioteca Accelerate do Hugging Face.

Este tutorial é destinado a engenheiros de machine learning (ML), pesquisadores, administradores e operadores de plataforma, além de especialistas em dados e IA interessados em implantar clusters do GKE em Google Cloud para treinar LLMs.

Objetivos

  • Acesse o modelo Gemma 3 usando o Hugging Face.

  • Prepare seu ambiente.

  • Crie e implante um cluster do GKE A4.

  • Ajuste o modelo Gemma 3 usando a biblioteca Accelerate do Hugging Face com paralelismo de dados totalmente fragmentados (FSDP).

  • Monitorar o job.

  • Fazer a limpeza.

Custos

Neste documento, você vai usar os seguintes componentes faturáveis do Google Cloud:

Para gerar uma estimativa de custo baseada na projeção de uso deste tutorial, use a calculadora de preços.

Novos usuários do Google Cloud podem estar qualificados para um teste sem custo financeiro.

Antes de começar

  1. Faça login na sua conta do Google Cloud . Se você começou a usar o Google Cloud, crie uma conta para avaliar o desempenho de nossos produtos em situações reais. Clientes novos também recebem US$ 300 em créditos para executar, testar e implantar cargas de trabalho.

  2. Instale a CLI do Google Cloud.

  3. Ao usar um provedor de identidade (IdP) externo, primeiro faça login na CLI gcloud com sua identidade federada.

  4. Para inicializar a CLI gcloud, execute o seguinte comando: 

    gcloud init

  5. Crie ou selecione um Google Cloud projeto.

    Funções necessárias para selecionar ou criar um projeto

    • Selecionar um projeto: não é necessário um papel específico do IAM para selecionar um projeto. Você pode escolher qualquer projeto em que tenha recebido um papel.
    • Criar um projeto: para criar um projeto, é necessário ter o papel de Criador de projetos (roles/resourcemanager.projectCreator), que contém a permissão resourcemanager.projects.create. Saiba como conceder papéis.

    • Crie um projeto do Google Cloud :

      gcloud projects create PROJECT_ID

      Substitua PROJECT_ID por um nome para o projeto Google Cloud que você está criando.

    • Selecione o projeto Google Cloud que você criou:

      gcloud config set project PROJECT_ID

      Substitua PROJECT_ID pelo nome do projeto do Google Cloud .

  6. Verifique se o faturamento está ativado para o projeto do Google Cloud .

  7. Ative as APIs necessárias:

    Funções necessárias para ativar APIs

    Para ativar APIs, você precisa da permissão serviceusage.services.enable. Se você criou o projeto, provavelmente já tem essa permissão com o papel de Proprietário (roles/owner). Caso contrário, é possível receber essa permissão com o papel de Administrador do Service Usage (roles/serviceusage.serviceUsageAdmin). Saiba como conceder papéis. 

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

  8. Instale a CLI do Google Cloud.

  9. Ao usar um provedor de identidade (IdP) externo, primeiro faça login na CLI gcloud com sua identidade federada.

  10. Para inicializar a CLI gcloud, execute o seguinte comando: 

    gcloud init

  11. Crie ou selecione um Google Cloud projeto.

    Funções necessárias para selecionar ou criar um projeto

    • Selecionar um projeto: não é necessário um papel específico do IAM para selecionar um projeto. Você pode escolher qualquer projeto em que tenha recebido um papel.
    • Criar um projeto: para criar um projeto, é necessário ter o papel de Criador de projetos (roles/resourcemanager.projectCreator), que contém a permissão resourcemanager.projects.create. Saiba como conceder papéis.

    • Crie um projeto do Google Cloud :

      gcloud projects create PROJECT_ID

      Substitua PROJECT_ID por um nome para o projeto Google Cloud que você está criando.

    • Selecione o projeto Google Cloud que você criou:

      gcloud config set project PROJECT_ID

      Substitua PROJECT_ID pelo nome do projeto do Google Cloud .

  12. Verifique se o faturamento está ativado para o projeto do Google Cloud .

  13. Ative as APIs necessárias:

    Funções necessárias para ativar APIs

    Para ativar APIs, você precisa da permissão serviceusage.services.enable. Se você criou o projeto, provavelmente já tem essa permissão com o papel de Proprietário (roles/owner). Caso contrário, é possível receber essa permissão com o papel de Administrador do Service Usage (roles/serviceusage.serviceUsageAdmin). Saiba como conceder papéis. 

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

  14. Atribua papéis à sua conta de usuário. Execute uma vez o seguinte comando para cada um dos seguintes papéis do 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

    Substitua:

    • PROJECT_ID: o ID do projeto.
    • USER_IDENTIFIER: o identificador da sua conta de usuário . Por exemplo, myemail@example.com.
    • ROLE: o papel do IAM concedido à sua conta de usuário.
  15. Ative a conta de serviço padrão para seu projeto do Google Cloud : 
      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. Conceda o papel de editor (roles/editor) à conta de serviço padrão: 
    gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \
    --role=roles/editor
  17. Crie as credenciais de autenticação local para sua conta de usuário: 
    gcloud auth application-default login
  18. Ative o Login do SO no seu projeto: 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  19. Faça login ou crie uma conta do Hugging Face.

Acessar o Gemma 3 usando o Hugging Face

Para usar o Hugging Face e acessar o Gemma 3, faça o seguinte:

  1. Fazer login no Hugging Face
  2. Crie um token de acesso write do Hugging Face.
    Clique em Seu perfil > Configurações > Tokens de acesso > +Criar novo token
  3. Copie e salve o valor do token write access. Você vai usar esse valor mais tarde neste tutorial.

Preparar o ambiente

Para preparar o ambiente, defina o seguinte:

export PROJECT_ID="YOUR_PROJECT_ID"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export CLUSTER_REGION="YOUR_REGION"
export RESERVATION="RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export ARTIFACT_REPO_LOCATION="ARTIFACT_REGISTRY_LOCATION"
export NETWORK="default"

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

Substitua:

  • PROJECT_ID: o nome do Google Cloud projeto em que você quer criar o cluster do GKE.

  • RESERVATION: o identificador da capacidade reservada.

  • CLUSTER_NAME: o nome do cluster do GKE a ser criado.

  • CLUSTER_REGION: a região em que você quer criar o cluster do GKE. Só é possível criar o cluster na região em que a reserva está.

  • HF_TOKEN: o token de acesso do Hugging Face que você criou na seção anterior.

  • ARTIFACT_REPO_LOCATION: o local em que você quer criar o repositório do Artifact Registry.

Criar um cluster do GKE no modo Autopilot

Para criar um cluster do GKE no modo Autopilot, execute o seguinte comando:

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

A criação do cluster do GKE pode levar algum tempo. Para verificar se o Google Cloud terminou de criar o cluster, acesse Clusters do Kubernetes no console Google Cloud .

Criar um secret do Kubernetes para as credenciais do Hugging Face

Para criar um secret do Kubernetes para as credenciais do Hugging Face, siga estas etapas:

  1. Configure kubectl para se comunicar com o cluster do GKE:

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

  2. Crie um secret do Kubernetes para armazenar seu token do Hugging Face:

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

Preparar sua carga de trabalho

Para preparar sua carga de trabalho, faça o seguinte:

  1. Crie scripts de carga de trabalho.

  2. Use o Docker e o Cloud Build para criar um contêiner de ajuste fino.

Criar scripts de carga de trabalho

Para criar os scripts que sua carga de trabalho de ajuste refinado usa, faça o seguinte:

  1. Crie um diretório para os scripts de carga de trabalho. Use esse diretório como seu diretório de trabalho.

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

  2. Crie o arquivo cloudbuild.yaml para usar o Google Cloud Build. Esse arquivo cria o contêiner de carga de trabalho e o armazena no Artifact Registry:

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

  3. Crie um arquivo Dockerfile para executar o job de ajuste:

    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. Crie o arquivo accel_fsdp_gemma3_config.yaml. Esse arquivo de configuração orienta o Hugging Face Accelerate a dividir o job de ajuste em várias GPUs.

    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. Crie o arquivo 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. Crie o arquivo 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 ---
    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"][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
    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,
        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()

Usar o Docker e o Cloud Build para criar um contêiner de ajuste fino

  1. Crie um repositório Docker do Artifact Registry:

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

  2. No diretório llm-finetuning-gemma que você criou em uma etapa anterior, execute o comando a seguir para criar a imagem de ajuste refinado e enviá-la ao Artifact Registry.

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

  3. Exporte o URL da imagem. Você vai usá-lo em uma etapa posterior deste tutorial:

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

Iniciar a carga de trabalho de ajuste refinado

Para iniciar sua carga de trabalho de ajuste refinado, faça o seguinte:

  1. Aplique o manifesto de ajuste para criar o job de ajuste:

    envsubst < finetune.yaml | kubectl apply -f -

    Como você está usando clusters no modo Autopilot do GKE, pode levar alguns minutos para iniciar o nó habilitado para GPU.

  2. Monitore o job executando o seguinte comando:

    watch kubectl get pods

  3. Verifique os registros do job executando o seguinte comando:

    kubectl logs job.batch/finetune-job -f

    O recurso de job faz o download dos dados do modelo e, em seguida, ajusta o modelo em todas as oito GPUs. O download leva cerca de cinco minutos para ser concluído. Depois que o download é concluído, o processo de ajuste fino leva aproximadamente duas horas e 30 minutos para ser concluído.

Monitore sua carga de trabalho

É possível monitorar o uso das GPUs no cluster do GKE para verificar se o job de ajuste refinado está sendo executado de maneira eficiente. Para fazer isso, abra o seguinte link no navegador:

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"))"

Ao monitorar sua carga de trabalho, você pode ver o seguinte:

  • Uso de GPUs: para um job de ajuste refinado saudável, espere ver o uso de todas as oito GPUs aumentar e se estabilizar em um nível alto durante todo o treinamento.
  • Duração do job: o job leva aproximadamente 10 minutos para ser concluído no cluster A4 especificado.

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.

Excluir os recursos

  1. Para excluir seu job de ajuste refinado, execute o seguinte comando:

    kubectl delete job finetune-job

  2. Para excluir o cluster do GKE, execute o seguinte comando:

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

Excluir o projeto

Excluir um projeto do Google Cloud :

gcloud projects delete PROJECT_ID

A seguir