O Google usa tecnologia de IA na tradução de conteúdos para seu idioma de preferência. As traduções com IA podem ter erros.

Ajustar o Gemma 4 em um cluster do GKE A4

GPUs em cluster

Neste tutorial, você vai aprender a ajustar um modelo de linguagem grande (LLM) Gemma 4 em um cluster do Google Kubernetes Engine (GKE) com vários hosts e GPUs no Google Cloud. Esse cluster usa uma única instância de máquina virtual (VM) A4 com 8 GPUs NVIDIA B200 anexadas.

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

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.
Depois que o cluster for implantado, execute um job de ajuste 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 4 usando o Hugging Face.
Prepare seu ambiente.
Crie e implante um cluster do GKE A4.
Ajuste o modelo Gemma 4 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

Ative as APIs necessárias, caso alguma ainda não esteja ativada:

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 pelo papel de Proprietário (roles/owner). Caso contrário, é possível receber essa permissão pelo 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
```

Ative a conta de serviço padrão para seu projeto 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

Conceda os papéis que a conta de serviço padrão precisa para executar a carga de trabalho de ajuste refinado:
```
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
```
Observação: esta etapa concede apenas os papéis predefinidos necessários para este tutorial. Conceder um papel básico, como Editor (roles/editor), dá à conta de serviço muito mais permissões do que ela precisa, e as políticas da sua organização podem bloquear a concessão. Para mais informações, consulte Práticas recomendadas para trabalhar com contas de serviço.
Crie as credenciais de autenticação local para sua conta de usuário:
```
gcloud auth application-default login
```
Observação: se você usa um shell local e um provedor de identidade (IdP) externo e encontrar um erro de autenticação após executar o comando anterior, faça login na CLI gcloud com sua identidade federada.

Ative o Login do SO no seu projeto:

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

Para conseguir as permissões necessárias a fim de concluir o tutorial, peça ao administrador para conceder a você os seguintes papéis do IAM no projeto:

Administrador do Kubernetes Engine (roles/container.admin)
Administrador do Compute (roles/compute.admin)
Administrador do Storage (roles/storage.admin)
Administrador do Artifact Registry (roles/artifactregistry.admin)
Editor do Cloud Build (roles/cloudbuild.builds.editor)
Usuário da conta de serviço (roles/iam.serviceAccountUser)
Administrador da conta de serviço (roles/iam.serviceAccountAdmin)
Administrador de projetos do IAM (roles/resourcemanager.projectIamAdmin)
Administrador do Service Usage (roles/serviceusage.serviceUsageAdmin)

Para mais informações sobre a concessão de papéis, consulte Gerenciar o acesso a projetos, pastas e organizações.

Também é possível conseguir as permissões necessárias usando papéis personalizados ou outros papéis predefinidos.

Acessar o Gemma 4 usando o Hugging Face

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

Fazer login no Hugging Face
Crie um token de acesso write do Hugging Face.
Clique em Seu perfil > Configurações > Tokens de acesso > +Criar novo token
Copie e salve o valor do token write access. Você vai usá-lo 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="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}"

Substitua:

YOUR_PROJECT_ID: o ID do Google Cloud projeto em que você quer criar o cluster do GKE.
YOUR_CLUSTER_NAME: o nome do cluster do GKE a ser criado.
YOUR_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á.
YOUR_RESERVATION_NAME: o identificador da capacidade reservada.
YOUR_HF_TOKEN: o token de acesso do Hugging Face que você criou na seção anterior.
YOUR_ARTIFACT_REGISTRY_LOCATION: a Google Cloud região em que você quer criar o repositório do Artifact Registry. Para minimizar a latência de extração de imagens, recomendamos usar a mesma região usada para YOUR_REGION.

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

A criação do cluster do GKE pode levar algum tempo. Para verificar se o Google Cloud concluiu a criação do 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:

Configure kubectl para se comunicar com o cluster do GKE:

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

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:

Criar scripts de carga de trabalho.
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:

Crie um diretório para os scripts de carga de trabalho. Use este diretório como seu diretório de trabalho.
```
mkdir llm-finetuning-gemma
cd llm-finetuning-gemma
```

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

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 && \
    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

Crie o arquivo accel_fsdp_gemma4_config.yaml. Esse arquivo de configuração direciona o Hugging Face Accelerate para dividir o job de ajuste nas oito GPUs locais no seu host único usando o 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

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_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

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

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

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

No diretório llm-finetuning-gemma criado 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}"
```

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

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

Iniciar a carga de trabalho de ajuste refinado

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

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.
Monitore o job executando o seguinte comando:
```
watch kubectl get pods
```
Depois que os pods estiverem em execução, verifique os registros do job:
```
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 for concluído, o processo de ajuste fino levará aproximadamente duas horas e 30 minutos.

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 2 horas e 30 minutos para ser concluído no cluster A4 especificado.

Limpar

Para evitar cobranças extras, exclua os recursos criados durante este tutorial.

Excluir os recursos

Para excluir seu job de ajuste refinado, execute o seguinte comando:
```
kubectl delete job finetune-job
```

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

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