Google utiliza tecnología de IA para traducir contenido a tu idioma preferido. Las traducciones realizadas con IA pueden contener errores.

Ajusta Gemma 4 en un clúster de GKE A4

GPUs agrupadas

En este instructivo, se muestra cómo ajustar un modelo de lenguaje grande (LLM) Gemma 4 en un clúster de Google Kubernetes Engine (GKE) de varios nodos y varias GPU en Google Cloud. Este clúster usa una sola instancia de máquina virtual (VM) A4 con 8 GPUs NVIDIA B200 conectadas.

Los dos procesos principales que se describen en este instructivo son los siguientes:

Implementa un clúster de GKE de alto rendimiento con GKE Autopilot. Como parte de esta implementación, crearás una imagen de VM personalizada con el software necesario preinstalado.
Después de que se implemente el clúster, ejecutarás un trabajo de ajuste distribuido con el conjunto de secuencias de comandos que acompañan a este instructivo. El trabajo aprovecha la biblioteca Hugging Face Accelerate.

Este instructivo está dirigido a ingenieros, investigadores, administradores y operadores de plataformas de aprendizaje automático (AA), y a especialistas en datos y en IA que deseen implementar clústeres de GKE en Google Cloud para entrenar LLMs.

Objetivos

Accede al modelo de Gemma 4 con Hugging Face.
Prepara tu entorno.
Crea e implementa un clúster de GKE de A4.
Ajusta el modelo Gemma 4 con la biblioteca Accelerate de Hugging Face y el paralelismo de datos completamente fragmentado (FSDP).
Supervisar tu trabajo
Realizar una limpieza

Costos

En este documento, usarás los siguientes componentes facturables de Google Cloud:

Para generar una estimación de costos en función del uso previsto, usa la calculadora de precios.

Es posible que los usuarios nuevos de Google Cloud cumplan con los requisitos para acceder a una prueba gratuita.

Antes de comenzar

Habilita las APIs requeridas, si alguna aún no está habilitada:

Roles necesarios para habilitar las APIs
Para habilitar APIs, necesitas el permiso serviceusage.services.enable. Si creaste el proyecto, es probable que ya tengas este permiso a través del rol de propietario (roles/owner). De lo contrario, puedes obtener este permiso a través del rol de administrador de Service Usage (roles/serviceusage.serviceUsageAdmin). Obtén más información para otorgar roles.
```
gcloud services enable compute.googleapis.com container.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com
```

Habilita la cuenta de servicio predeterminada para tu proyecto de 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

Otorga los roles que necesita la cuenta de servicio predeterminada para ejecutar la carga de trabajo de ajuste:
```
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
```
Nota: Este paso solo otorga los roles predefinidos que requiere este instructivo. Otorgar un rol básico, como Editor (roles/editor), le da a la cuenta de servicio muchos más permisos de los que necesita, y es posible que las políticas de tu organización bloqueen el otorgamiento. Si deseas obtener más información, consulta Prácticas recomendadas para trabajar con cuentas de servicio.
Crea credenciales de autenticación locales para tu cuenta de usuario:
```
gcloud auth application-default login
```
Nota: Si usas un shell local y un proveedor de identidad (IdP) externo, y encuentras un error de autenticación después de ejecutar el comando anterior, accede a la gcloud CLI con tu identidad federada.

Habilita el Acceso al SO para tu proyecto:

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

Si quieres obtener los permisos que necesitas para completar este instructivo, pídele a tu administrador que te otorgue los siguientes roles de IAM en tu proyecto:

Administrador de Kubernetes Engine (roles/container.admin)
Administrador de Compute (roles/compute.admin)
Administrador de almacenamiento (roles/storage.admin)
Administrador de Artifact Registry (roles/artifactregistry.admin)
Editor de Cloud Build (roles/cloudbuild.builds.editor)
Usuario de la cuenta de servicio (roles/iam.serviceAccountUser)
Administrador de cuenta de servicio (roles/iam.serviceAccountAdmin)
Administrador de IAM de proyecto (roles/resourcemanager.projectIamAdmin)
Administrador de Service Usage ()roles/serviceusage.serviceUsageAdmin

Para obtener más información sobre cómo otorgar roles, consulta Administra el acceso a proyectos, carpetas y organizaciones.

También puedes obtener los permisos necesarios a través de roles personalizados o cualquier otro rol predefinido.

Accede a Gemma 4 con Hugging Face

Para usar Hugging Face y acceder a Gemma 4, haz lo siguiente:

Accede a Hugging Face
Crea un token de acceso de Hugging Face write.
Haz clic en Tu perfil > Configuración > Tokens de acceso > +Crear token nuevo.
Copia y guarda el valor del token write access. La usarás más adelante en este instructivo.

Prepara el entorno

Para preparar tu entorno, establece lo siguiente:

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

Reemplaza lo siguiente:

YOUR_PROJECT_ID: Es el ID del Google Cloud proyecto en el que deseas crear el clúster de GKE.
YOUR_CLUSTER_NAME: Es el nombre del clúster de GKE que se creará.
YOUR_REGION: Es la región en la que deseas crear tu clúster de GKE. Solo puedes crear el clúster en la región en la que existe tu reserva.
YOUR_RESERVATION_NAME: Es el identificador de tu capacidad reservada.
YOUR_HF_TOKEN: El token de acceso de Hugging Face que creaste en la sección anterior.
YOUR_ARTIFACT_REGISTRY_LOCATION: Es la Google Cloud región en la que deseas crear tu repositorio de Artifact Registry. Para minimizar la latencia de extracción de imágenes, te recomendamos que uses la misma región que usaste para YOUR_REGION.

Crea un clúster de GKE en modo Autopilot

Para crear un clúster de GKE en modo Autopilot, ejecuta el siguiente comando:

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

La creación del clúster de GKE puede tardar un tiempo en completarse. Para verificar que Google Cloud haya terminado de crear tu clúster, ve a Clústeres de Kubernetes en la consola de Google Cloud .

Crea un secreto de Kubernetes para las credenciales de Hugging Face

Para crear un secreto de Kubernetes para las credenciales de Hugging Face, sigue estos pasos:

Configura kubectl para comunicarse con tu clúster de GKE:

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

Crea un secreto de Kubernetes para almacenar tu token de Hugging Face:

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

Prepara tu carga de trabajo

Para preparar tu carga de trabajo, haz lo siguiente:

Crea secuencias de comandos de carga de trabajo.
Usa Docker y Cloud Build para crear un contenedor de ajuste.

Crea secuencias de comandos de cargas de trabajo

Para crear las secuencias de comandos que usa tu carga de trabajo de ajuste, haz lo siguiente:

Crea un directorio para las secuencias de comandos de la carga de trabajo. Usa este directorio como tu directorio de trabajo.
```
mkdir llm-finetuning-gemma
cd llm-finetuning-gemma
```

Crea el archivo cloudbuild.yaml para usar Google Cloud Build. Este archivo crea el contenedor de tu carga de trabajo y lo almacena en Artifact Registry:

steps:
- name: 'gcr.io/cloud-builders/docker'
  args:
  - 'build'
  - '-t'
  - '$_ARTIFACT_REPO_LOCATION-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-gpu:2.0.0'
  - '.'
images:
- '$_ARTIFACT_REPO_LOCATION-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-gpu:2.0.0'
options:
  logging: CLOUD_LOGGING_ONLY

Crea un archivo Dockerfile para ejecutar el trabajo 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

Crea el archivo accel_fsdp_gemma4_config.yaml. Este archivo de configuración indica a Hugging Face Accelerate que divida el trabajo de ajuste en las ocho GPUs locales de tu único host con FSDP:

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
enable_cpu_affinity: false
fsdp_config:
  fsdp_activation_checkpointing: false
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_cpu_ram_efficient_loading: false
  fsdp_offload_params: false
  fsdp_reshard_after_forward: true
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: Gemma4TextDecoderLayer
  fsdp_version: 2
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

Crea el archivo finetune.yaml:

apiVersion: batch/v1
kind: Job
metadata:
  name: finetune-job
  namespace: default
spec:
  backoffLimit: 2
  template:
    metadata:
      annotations:
        kubectl.kubernetes.io/default-container: finetuner
    spec:
      terminationGracePeriodSeconds: 600
      containers:
      - name: finetuner
        image: $IMAGE_URL
        command: ["accelerate","launch"]
        args:
        - "--config_file"
        - "accel_fsdp_gemma4_config.yaml"
        - "finetune.py"
        - "--model_id"
        - "google/gemma-4-31B-it"
        - "--output_dir"
        - "gemma-31b-text-to-sql"
        - "--per_device_train_batch_size"
        - "2"
        - "--gradient_accumulation_steps"
        - "32"
        - "--num_train_epochs"
        - "3"
        - "--learning_rate"
        - "1e-5"
        - "--save_strategy"
        - "steps"
        - "--save_steps"
        - "20"
        - "--push_to_hub"
        resources:
          limits:
            nvidia.com/gpu: "8"
            memory: "1200Gi"
            ephemeral-storage: "350Gi"
          requests:
            nvidia.com/gpu: "8"
            memory: "1200Gi"
            ephemeral-storage: "350Gi"
        env:
        - name: HF_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        volumeMounts:
        - mountPath: /dev/shm
          name: dshm
      volumes:
      - name: dshm
        emptyDir:
          medium: Memory
          sizeLimit: 64Gi
      nodeSelector:
        cloud.google.com/gke-accelerator: nvidia-b200
        cloud.google.com/reservation-name: $RESERVATION
        cloud.google.com/reservation-affinity: "specific"
        cloud.google.com/gke-gpu-driver-version: latest
      restartPolicy: OnFailure

Crea el archivo 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()

Usa Docker y Cloud Build para crear un contenedor de ajuste

Crea un repositorio de Docker de Artifact Registry:

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

En el directorio llm-finetuning-gemma que creaste en un paso anterior, ejecuta el siguiente comando para crear la imagen de ajuste y enviarla a Artifact Registry.
```
gcloud builds submit . \
    --substitutions=_ARTIFACT_REPO_LOCATION="${ARTIFACT_REPO_LOCATION}"
```

Exporta la URL de la imagen. Lo usarás en un paso posterior de este instructivo:

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

Inicia tu carga de trabajo de ajuste

Para iniciar tu carga de trabajo de ajuste, haz lo siguiente:

Aplica el manifiesto de ajuste para crear el trabajo de ajuste:
```
envsubst < finetune.yaml | kubectl apply -f -
```
Como usas clústeres en el modo Autopilot de GKE, es posible que el inicio del nodo habilitado para GPU tarde unos minutos.
Supervisa el trabajo ejecutando el siguiente comando:
```
watch kubectl get pods
```
Después de que los Pods se ejecuten, verifica los registros del trabajo:
```
kubectl logs job.batch/finetune-job -f
```
El recurso de trabajo descarga los datos del modelo y, luego, ajusta el modelo en las ocho GPUs. La descarga tarda alrededor de cinco minutos en completarse. Una vez que se completa la descarga, el proceso de ajuste tarda aproximadamente dos horas y media en completarse.

Supervisa tu carga de trabajo

Puedes supervisar el uso de las GPUs en tu clúster de GKE para verificar que tu trabajo de ajuste fino se ejecute de manera eficiente. Para ello, abre el siguiente vínculo en tu 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\"))"

Cuando supervisas tu carga de trabajo, puedes ver lo siguiente:

Uso de las GPUs: Para un trabajo de ajuste fino correcto, puedes esperar ver que el uso de las 8 GPUs aumente y se estabilice en un nivel alto durante todo el entrenamiento.
Duración del trabajo: El trabajo debería tardar aproximadamente dos horas y 30 minutos en completarse en el clúster A4 especificado.

Realiza una limpieza

Para evitar que se generen cargos adicionales, borra los recursos que creaste durante este instructivo.

Borra tus recursos

Para borrar tu trabajo de ajuste, ejecuta el siguiente comando:
```
kubectl delete job finetune-job
```

Para borrar tu clúster de GKE, ejecuta el siguiente comando:

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