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, mostramos como ajustar um Gemma 4 modelo de linguagem grande (LLM) em um cluster do Google Kubernetes Engine (GKE) de vários hosts e GPUs em 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ê usa o Cloud Build para criar uma imagem de contêiner personalizada com o software necessário pré-instalado.
Depois que o cluster é implantado, você executa um job de ajuste fino que usa o paralelismo de dados totalmente fragmentados (FSDP, na sigla em inglês) para fragmentar o modelo em todas as oito GPUs de um único nó A4 (a4-highgpu-8g). O job usa o conjunto de scripts que acompanham este tutorial e aproveita a biblioteca Accelerate do Hugging Face.

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

Objetivos

Acesse o modelo Gemma 4 usando o Hugging Face.
Prepare o 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).
Monitore o job.
Fazer a limpeza.

Custos

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

Para gerar uma estimativa de custo baseada na projeção de uso, 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, se alguma ainda não estiver ativada:

Funções necessárias para ativar APIs
Para ativar as APIs, é necessário ter a permissão serviceusage.services.enable. Se você criou o projeto, provavelmente já tem essa permissão pela função de proprietário (roles/owner). Caso contrário, você pode receber essa permissão pela função de administrador de uso do serviço (roles/serviceusage.serviceUsageAdmin). Saiba como conceder funções.
```
gcloud services enable compute.googleapis.com container.googleapis.com artifactregistry.googleapis.com cloudbuild.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com
```

Ative a conta de serviço padrão para o seu Google Cloud projeto:

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 as funções que a conta de serviço padrão precisa para executar a carga de trabalho de ajuste fino:

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

Para verificar se as funções foram concedidas, execute o seguinte 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)"

Crie credenciais de autenticação local para sua conta de usuário:
```
gcloud auth application-default login
```
Observação: se você usar 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 para seu projeto:

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

Para conseguir as permissões que você precisa para concluir este tutorial, peça ao administrador para conceder a você os seguintes papéis do IAM no seu projeto:

Administrador do Kubernetes Engine (roles/container.admin)
Administrador do Compute (roles/compute.admin)
Administrador de armazenamento (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 com papéis personalizados ou outros papéis predefinidos.

Acessar o Gemma 4 usando o Hugging Face

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

Faça 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ê o usará 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 existe.
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 Google Cloud terminou de criar o cluster, acesse Clusters do Kubernetes no Google Cloud console.

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

Crie 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 usados pela carga de trabalho de ajuste fino, faça o seguinte:

Crie um diretório para os scripts de carga de trabalho. Use esse diretório como 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 fino:

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 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 do Docker no 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 fino e enviá-la ao Artifact Registry.
```
gcloud builds submit . \
    --substitutions=_ARTIFACT_REPO_LOCATION="${ARTIFACT_REPO_LOCATION}"
```
Observação: a criação e o envio da imagem podem levar 20 minutos ou mais. Para conferir o progresso da criação, abra o link fornecido pelo comando com a saída Logs are available at.

Exporte o URL da imagem. Você o usará 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 fino

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

Aplique o manifesto de ajuste fino para criar o job de ajuste fino:
```
envsubst < finetune.yaml | kubectl apply -f -
```
Como você está usando clusters no modo Autopilot do GKE, pode levar alguns minutos para iniciar o nó ativado por 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 é concluído, o processo de ajuste fino leva aproximadamente 45 minutos. O job salva um checkpoint a cada 20 etapas e envia o adaptador de adaptação de classificação baixa (LoRA) ajustado para um repositório privado do Hugging Face chamado gemma-31b-text-to-sql.

Depois que o job for concluído, acesse o URL do modelo ajustado no Hugging Face e abra o URL no navegador:

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"

Monitorar a carga de trabalho

É possível monitorar o uso das GPUs no cluster do GKE para verificar se o job de ajuste fino está sendo executado com eficiência. 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 a carga de trabalho, você pode conferir o seguinte:

Uso de GPUs: para um job de ajuste fino normal, você pode esperar que o uso de todas as 8 GPUs aumente e se estabilize em um nível alto durante todo o treinamento.
Duração do job: o job leva aproximadamente 45 minutos para ser concluído no cluster A4 especificado.

Limpar

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

Excluir os recursos

Para excluir o job de ajuste fino, 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}"

Para excluir o repositório do Artifact Registry, execute o seguinte comando:

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