Usa Ray para ajustar Gemma 3 para tareas de visión en GKE

En este instructivo, se muestra cómo ajustar un modelo de Gemma 3 con el framework de Ray en un clúster de GKE de varios nodos. El clúster usa dos instancias de máquina virtual (VM) A4, cada una con ocho GPUs NVIDIA B200 conectadas.

El contenido de este instructivo se divide en dos partes:

  1. Preparar el clúster de Ray sobre un clúster de Autopilot de GKE
  2. Ejecuta un trabajo de entrenamiento distribuido que utiliza 2 instancias A4, cada una con 8 GPUs B200.

Este instructivo está dirigido a ingenieros e investigadores de aprendizaje automático (AA), administradores y operadores de plataformas, y especialistas en datos y en IA que deseen distribuir una carga de trabajo de IA en varios nodos y GPUs.

Objetivos

  • Accede a un modelo de Gemma 3 con Hugging Face.

  • Prepara tu entorno.

  • Crea un clúster de GKE Autopilot con el operador de Ray instalado.

  • Configura el clúster de Ray en el clúster de GKE para que acepte trabajos de Ray.

  • Configurar y ejecutar un trabajo de Ray que ajuste el modelo Gemma 3 en función de la entrada visual

  • Supervisa tu carga de 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

  1. Accede a tu cuenta de Google Cloud . Si eres nuevo en Google Cloud, crea una cuenta para evaluar el rendimiento de nuestros productos en situaciones reales. Los clientes nuevos también obtienen $300 en créditos gratuitos para ejecutar, probar y, además, implementar cargas de trabajo.

  2. Instala Google Cloud CLI.

  3. Si usas un proveedor de identidad externo (IdP), primero debes acceder a la gcloud CLI con tu identidad federada.

  4. Para inicializar gcloud CLI, ejecuta el siguiente comando: 

    gcloud init

  5. Crea o selecciona un Google Cloud proyecto.

    Roles necesarios para seleccionar o crear un proyecto

    • Selecciona un proyecto: Para seleccionar un proyecto, no se requiere un rol de IAM específico. Puedes seleccionar cualquier proyecto en el que se te haya otorgado un rol.
    • Crear un proyecto: Para crear un proyecto, necesitas el rol de Creador de proyectos (roles/resourcemanager.projectCreator), que contiene el permiso resourcemanager.projects.create. Obtén más información para otorgar roles.

    • Crea un proyecto de Google Cloud :

      gcloud projects create PROJECT_ID

      Reemplaza PROJECT_ID por un nombre para el proyecto Google Cloud que estás creando.

    • Selecciona el proyecto Google Cloud que creaste:

      gcloud config set project PROJECT_ID

      Reemplaza PROJECT_ID por el nombre de tu proyecto de Google Cloud .

  6. Verifica que la facturación esté habilitada para tu proyecto de Google Cloud .

  7. Habilita la API necesaria:

    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 logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com container.googleapis.com

  8. Instala Google Cloud CLI.

  9. Si usas un proveedor de identidad externo (IdP), primero debes acceder a la gcloud CLI con tu identidad federada.

  10. Para inicializar gcloud CLI, ejecuta el siguiente comando: 

    gcloud init

  11. Crea o selecciona un Google Cloud proyecto.

    Roles necesarios para seleccionar o crear un proyecto

    • Selecciona un proyecto: Para seleccionar un proyecto, no se requiere un rol de IAM específico. Puedes seleccionar cualquier proyecto en el que se te haya otorgado un rol.
    • Crear un proyecto: Para crear un proyecto, necesitas el rol de Creador de proyectos (roles/resourcemanager.projectCreator), que contiene el permiso resourcemanager.projects.create. Obtén más información para otorgar roles.

    • Crea un proyecto de Google Cloud :

      gcloud projects create PROJECT_ID

      Reemplaza PROJECT_ID por un nombre para el proyecto Google Cloud que estás creando.

    • Selecciona el proyecto Google Cloud que creaste:

      gcloud config set project PROJECT_ID

      Reemplaza PROJECT_ID por el nombre de tu proyecto de Google Cloud .

  12. Verifica que la facturación esté habilitada para tu proyecto de Google Cloud .

  13. Habilita la API necesaria:

    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 logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com container.googleapis.com
    14.

  14. Otorga roles a tu cuenta de usuario. Ejecuta el siguiente comando una vez para cada uno de los siguientes roles de IAM: roles/compute.admin, roles/iam.serviceAccountUser, roles/file.editor, roles/storage.admin, roles/container.clusterAdmin, roles/serviceusage.serviceUsageAdmin 

    gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

    Reemplaza lo siguiente:

    • PROJECT_ID: ID del proyecto
    • USER_IDENTIFIER: Es el identificador de tu cuenta de usuario de . Por ejemplo, myemail@example.com.
    • ROLE: Es el rol de IAM que otorgas a tu cuenta de usuario.
  15. Habilita la cuenta de servicio predeterminada para tu proyecto Google Cloud : 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --project=PROJECT_ID

    Reemplaza PROJECT_NUMBER por el número del proyecto. Para revisar el número de tu proyecto, consulta Cómo obtener un proyecto existente.

  16. Otorga el rol de editor (roles/editor) a la cuenta de servicio predeterminada: 
    gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \
    --role=roles/editor
  17. Crea credenciales de autenticación locales para tu cuenta de usuario: 
    gcloud auth application-default login
  18. Accede a tu cuenta de Hugging Face o crea una.

Accede a Gemma 3 con Hugging Face

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

  1. Firma el acuerdo de consentimiento para usar Gemma 3.

  2. Crea un token de Hugging Face read access.

  3. Copia y guarda el valor del token read access. La usarás más adelante en este instructivo.

Prepara el entorno

Prepara tu entorno configurando los parámetros necesarios y estableciendo las variables de entorno.

Ejecuta lo siguiente:

export PROJECT_ID="YOUR_PROJECT_ID"
export RESERVATION="RESERVATION_ID"
export REGION="YOUR_REGION"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export NETWORK="default"
export RAY_SA="YOUR_RAY_SA"
export GSA_NAME="YOUR_GSA_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"

gcloud config set project $PROJECT_ID
gcloud config set billing/quota_project $PROJECT_ID

Reemplaza lo siguiente:

  • YOUR_PROJECT_ID: Es el nombre del Google Cloud proyecto en el que deseas crear el clúster de GKE.
  • YOUR_RESERVATION_ID: Es la URL de la reserva que deseas usar para crear tu clúster. Según el proyecto en el que existe la reserva, especifica uno de los siguientes valores:
    • La reserva existe en tu proyecto: Especifica el nombre de la reserva (p.ej., my-reservation).
    • La reserva existe en otro proyecto: Especifica la ruta completa en el formato projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME.
  • 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_CLUSTER_NAME: Es el nombre del clúster de GKE que se creará.
  • YOUR_HF_TOKEN: El token de Hugging Face que creaste en un paso anterior
  • YOUR_RAY_SA: Es el nombre de la cuenta de servicio dentro del clúster de Kubernetes.
  • YOUR_GSA_NAME: Es el nombre de la cuenta de servicio de Google.
  • YOUR_GCS_BUCKET: Es el nombre del bucket en el que almacenas los resultados del punto de control del entrenamiento.

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 \
    --enable-ray-operator \
    --enable-ray-cluster-monitoring \
    --enable-ray-cluster-logging \
    --location=$REGION

La creación del clúster de GKE puede tardar un tiempo en completarse. Para verificar si Google Cloud terminó 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

En Cloud Shell, para crear un secreto de Kubernetes para las credenciales de Hugging Face, haz lo siguiente:

  1. Configura kubectl para conectarte a tu clúster:

    gcloud container clusters get-credentials $CLUSTER_NAME \
    --region=$REGION

  2. Crea un Secret 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 -

Crea el bucket de Cloud Storage

Si deseas usar un bucket nuevo para almacenar tus artefactos de entrenamiento, ejecuta el siguiente comando:

gcloud storage buckets create gs://$GCS_BUCKET --location=$REGION

Si quieres usar un bucket existente, puedes omitir este paso. Sin embargo, debes asegurarte de que tu bucket esté en la misma región que tu clúster.

Crea una cuenta de servicio de IAM

En Cloud Shell, para crear una cuenta de servicio de IAM (también conocida como cuenta de servicio de Google o GSA) y otorgarle permisos para acceder a tu bucket de Cloud Storage, haz lo siguiente:

  1. Crea una cuenta de servicio de IAM:

    gcloud iam service-accounts create $GSA_NAME \
    --project=$PROJECT_ID \
    --description="Account for Ray to save Gemma training results" \
    --display-name="Ray Gemma Training SA"

  2. Otorga el rol de administrador de almacenamiento (roles/storage.admin) a tu cuenta de servicio de IAM para tu bucket de Cloud Storage:

    gcloud storage buckets add-iam-policy-binding gs://${GCS_BUCKET#gs://} \
    --member="serviceAccount:$GSA_NAME@$PROJECT_ID.iam.gserviceaccount.com" \
    --role="roles/storage.admin"

Crea una cuenta de servicio de Kubernetes

En Cloud Shell, crea una cuenta de servicio de Kubernetes y configura Workload Identity para otorgar acceso a los Pods de Ray a los recursos de Google Cloud :

  1. Crea una cuenta de servicio de Kubernetes: 

    kubectl create serviceaccount $RAY_SA \
    --dry-run=client -o yaml | kubectl apply -f -

  2. Vincula tu cuenta de servicio de Kubernetes a la cuenta de servicio de IAM que creaste anteriormente para habilitar Workload Identity: 

    gcloud iam service-accounts add-iam-policy-binding $GSA_NAME@$PROJECT_ID.iam.gserviceaccount.com \
    --project=$PROJECT_ID \
    --role="roles/iam.workloadIdentityUser" \
    --member="serviceAccount:$PROJECT_ID.svc.id.goog[default/$RAY_SA]"

  3. Anota tu cuenta de servicio de Kubernetes con la dirección de correo electrónico de la cuenta de servicio de IAM: 

    kubectl annotate serviceaccount $RAY_SA iam.gke.io/gcp-service-account=$GSA_NAME@$PROJECT_ID.iam.gserviceaccount.com --overwrite

Guarda tu código de entrenamiento como un ConfigMap

Para evitar la necesidad de incorporar tu secuencia de comandos de entrenamiento en una imagen de contenedor, la almacenas como un ConfigMap en tu clúster. Este ConfigMap se activa en los sistemas de archivos de Pod, lo que te permite actualizar la secuencia de comandos de entrenamiento sin tener que volver a crear todo el clúster de Ray.

Para almacenar tu secuencia de comandos de entrenamiento como un ConfigMap en tu clúster, completa los siguientes pasos:

  1. Crea un directorio con el nombre code y, en ese directorio, crea un archivo con el nombre vision_train.py.

  2. Copia el siguiente código en el archivo vision_train.py:

    import argparse
import datetime
import ray
import ray.train.huggingface.transformers
import torch
from PIL import Image
from datasets import load_dataset
from peft import LoraConfig
from ray.train import ScalingConfig, RunConfig
from ray.train.torch import TorchTrainer
from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig
from trl import SFTConfig
from trl import SFTTrainer

# System message for the assistant
system_message = "You are an expert product description writer for Amazon."

# User prompt that combines the user query and the schema
user_prompt = """Create a Short Product description based on the provided <PRODUCT> and <CATEGORY> and image.
Only return description. The description should be SEO optimized and for a better mobile search experience.

<PRODUCT>
{product}
</PRODUCT>

<CATEGORY>
{category}
</CATEGORY>
"""

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_id", type=str, default="google/gemma-3-4b-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("--dataset_name", type=str, default="philschmid/amazon-product-descriptions-vlm", help="Hugging Face dataset name")
    parser.add_argument("--output_dir", type=str, default="gemma-3-4b-seo-optimized", help="Directory to save model checkpoints")
    parser.add_argument("--gcs_bucket", type=str, required=True, help="storage bucket name used to synchronize tasks and save checkpoints")
    parser.add_argument("--push_to_hub", help="Push model to Hugging Face hub", action="store_true")

    # 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=1, help="Batch size per device during training")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=4, help="Gradient accumulation steps")
    parser.add_argument("--learning_rate", type=float, default=2e-4, 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="epoch", help="Checkpoint save strategy")
    parser.add_argument("--save_steps", type=int, default=100, help="Save checkpoint every X steps")

    return parser.parse_args()

# Convert dataset to OAI messages
def format_data(sample):
    return {
        "messages": [
            {
                "role": "system",
                "content": [{"type": "text", "text": system_message}],
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": user_prompt.format(
                            product=sample["Product Name"],
                            category=sample["Category"],
                        ),
                    },
                    {
                        "type": "image",
                        "image": sample["image"],
                    },
                ],
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": sample["description"]}],
            },
        ],
    }

def process_vision_info(messages: list[dict]) -> list[Image.Image]:
    image_inputs = []
    # Iterate through each conversation
    for msg in messages:
        # Get content (ensure it's a list)
        content = msg.get("content", [])
        if not isinstance(content, list):
            content = [content]

        # Check each content element for images
        for element in content:
            if isinstance(element, dict) and ("image" in element or element.get("type") == "image"):
                # Get the image and convert to RGB
                if "image" in element:
                    image = element["image"]
                else:
                    image = element
                image_inputs.append(image.convert("RGB"))
    return image_inputs

def train(args):
    # Load dataset from the hub
    dataset = load_dataset(args.dataset_name, split="train")

    # Convert dataset to OAI messages
    # need to use list comprehension to keep Pil.Image type, .mape convert image to bytes
    dataset = [format_data(sample) for sample in dataset]

    # Hugging Face model id
    model_id = args.model_id

    # Check if GPU benefits from bfloat16
    if torch.cuda.get_device_capability()[0] < 8:
        raise ValueError("GPU does not support bfloat16, please use a GPU that supports bfloat16.")

    # Define model init arguments
    model_kwargs = dict(
        attn_implementation="eager",  # Use "flash_attention_2" when running on Ampere or newer GPU
        torch_dtype=torch.bfloat16,  # What torch dtype to use, defaults to auto
        # device_map="auto",  # Let torch decide how to load the model
    )

    # BitsAndBytesConfig int-4 config
    model_kwargs["quantization_config"] = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=model_kwargs["torch_dtype"],
        bnb_4bit_quant_storage=model_kwargs["torch_dtype"],
    )

    # Load model and tokenizer
    model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs)
    processor = AutoProcessor.from_pretrained(model_id, use_fast=True)

    peft_config = LoraConfig(
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        r=args.lora_r,
        bias="none",
        target_modules="all-linear",
        task_type="CAUSAL_LM",
        modules_to_save=[
            "lm_head",
            "embed_tokens",
        ],
    )

    args = SFTConfig(
        output_dir=args.output_dir,  # directory to save and repository id
        num_train_epochs=args.num_train_epochs,  # number of training epochs
        per_device_train_batch_size=args.per_device_train_batch_size,  # batch size per device during training
        gradient_accumulation_steps=args.gradient_accumulation_steps,  # number of steps before performing a backward/update pass
        gradient_checkpointing=True,  # use gradient checkpointing to save memory
        optim="adamw_torch_fused",  # use fused adamw optimizer
        logging_steps=args.logging_steps,  # log every N steps
        save_strategy=args.save_strategy,  # save checkpoint every epoch
        learning_rate=args.learning_rate,  # learning rate, based on QLoRA paper
        bf16=True,  # use bfloat16 precision
        max_grad_norm=0.3,  # max gradient norm based on QLoRA paper
        warmup_ratio=0.03,  # warmup ratio based on QLoRA paper
        lr_scheduler_type="constant",  # use constant learning rate scheduler
        push_to_hub=args.push_to_hub,  # push model to hub
        report_to="tensorboard",  # report metrics to tensorboard
        gradient_checkpointing_kwargs={
            "use_reentrant": False
        },  # use reentrant checkpointing
        dataset_text_field="",  # need a dummy field for collator
        dataset_kwargs={"skip_prepare_dataset": True},  # important for collator
    )
    args.remove_unused_columns = False  # important for collator

    # Create a data collator to encode text and image pairs
    def collate_fn(examples):
        texts = []
        images = []
        for example in examples:
            image_inputs = process_vision_info(example["messages"])
            text = processor.apply_chat_template(
                example["messages"], add_generation_prompt=False, tokenize=False
            )
            texts.append(text.strip())
            images.append(image_inputs)

        # Tokenize the texts and process the images
        batch = processor(text=texts, images=images, return_tensors="pt", padding=True)

        # The labels are the input_ids, and we mask the padding tokens and image tokens in the loss computation
        labels = batch["input_ids"].clone()

        # Mask image tokens
        boi_token = processor.tokenizer.special_tokens_map.get("boi_token")
        image_token_id = processor.tokenizer.convert_tokens_to_ids(boi_token) if boi_token else None

        # Mask tokens for not being used in the loss computation
        labels[labels == processor.tokenizer.pad_token_id] = -100
        if image_token_id is not None:
            labels[labels == image_token_id] = -100
        labels[labels == 262144] = -100

        batch["labels"] = labels
        return batch

    trainer = SFTTrainer(
        model=model,
        args=args,
        train_dataset=dataset,
        peft_config=peft_config,
        processing_class=processor,
        data_collator=collate_fn,
    )

    callback = ray.train.huggingface.transformers.RayTrainReportCallback()
    trainer.add_callback(callback)
    trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)

    # Start training, the model will be automatically saved to the Hub and the output directory
    trainer.train()

    # Save the final model again to the Hugging Face Hub
    trainer.save_model()

if __name__ == "__main__":
    args = get_args()
    print("Starting training task!")
    training_name = f"gemma_vision_train_{datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')}"

    gcs_bucket = args.gcs_bucket
    if not gcs_bucket.startswith("gs://"):
        gcs_bucket = "gs://" + gcs_bucket

    run_config = RunConfig(
        storage_path=gcs_bucket,
        name=training_name,
    )
    scaling_config = ScalingConfig(num_workers=16, use_gpu=True, accelerator_type="B200")
    ray_trainer = TorchTrainer(train, train_loop_config=args, scaling_config=scaling_config, run_config=run_config)
    print("Commencing training!")
    result = ray_trainer.fit()

  3. Guarda el archivo.

  4. Crea un objeto ConfigMap en tu clúster:

    kubectl create cm ray-job-cm --from-file=code -o yaml --dry-run=client | kubectl apply -f -
    Para actualizar el script de entrenamiento, vuelve a ejecutar el comando anterior. Es posible que tarde un minuto en propagarse a todos los pods.

Configura el clúster de Ray

  1. Para crear un clúster de Ray en tu clúster de GKE, guarda el siguiente archivo YAML como un archivo con el nombre ray_cluster.yaml.

    apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: gemma3-tuning
spec:
  rayVersion: '2.48.0'
  headGroupSpec:
    rayStartParams:
      dashboard-host: '0.0.0.0'
    template:
      metadata:
      spec:
        serviceAccountName: $RAY_SA
        containers:
        - name: ray-head
          image: rayproject/ray:2.48.0
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          resources:
            limits:
              cpu: "24"
              ephemeral-storage: "9Gi"
              memory: "64Gi"
            requests:
              cpu: "24"
              ephemeral-storage: "9Gi"
              memory: "64Gi"
          env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
          volumeMounts:
            - name: job-code
              mountPath: /code/
            - mountPath: /mnt/local-ssd/
              name: local-storage
        volumes:
          - name: job-code
            configMap:
              name: ray-job-cm
          - name: local-storage
            emptyDir: { }
  workerGroupSpecs:
  - replicas: 2
    minReplicas: 1
    maxReplicas: 5
    groupName: gpu-group
    rayStartParams: {}
    template:
      spec:
        serviceAccountName: $RAY_SA
        containers:
        - name: ray-worker
          image: rayproject/ray:2.48.0-gpu
          resources:
            limits:
              nvidia.com/gpu: "8"
              cpu: "80"
              memory: "320Gi"
            requests:
              nvidia.com/gpu: "8"
              cpu: "80"
              memory: "320Gi"
          env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
          volumeMounts:
            - name: job-code
              mountPath: /code/
            - mountPath: /mnt/local-ssd/
              name: local-storage
        volumes:
          - name: job-code
            configMap:
              name: ray-job-cm
          - name: local-storage
            emptyDir: { }
        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

  2. Aplica esta definición de YAML a tu clúster con el siguiente comando:

    envsubst < ray_cluster.yaml | kubectl apply -f -
    La marca $RESERVATION se reemplaza automáticamente por el nombre que configuraste como variable de entorno.

    Ray Operator crea los Pods de raylet, lo que, a su vez, activa el ajuste de escala automático del clúster para proporcionar a esos Pods los nodos adecuados. Se crean tres pods en tu clúster: un nodo principal y dos nodos trabajadores. Los nodos trabajadores están equipados con las GPU B200.

  3. Para verificar que los tres Pods estén listos, ejecuta el siguiente comando:

    kubectl get pods
    La lista de pods de un clúster de Ray listo es similar a la siguiente:

    NAME                                   READY   STATUS    RESTARTS   AGE
gemma3-tuning-gpu-group-worker-s4h8f   2/2     Running   0          16m
gemma3-tuning-gpu-group-worker-stg5f   2/2     Running   0          5m34s
gemma3-tuning-head-zbdvp               2/2     Running   0          16m

Programa un trabajo de entrenamiento

  1. Guarda lo siguiente como un archivo ray_job.yaml:

    apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: test-ray-job
spec:
  entrypoint: python /code/vision_train.py --gcs_bucket $GCS_BUCKET
  runtimeEnvYAML: |
    pip:
      - torch==2.8.0
      - torchvision==0.23.0
      - ray==2.48.0
      - transformers==4.55.2
      - datasets==4.0.0
      - evaluate==0.4.5
      - accelerate==1.10.0
      - pillow==11.3.0
      - bitsandbytes==0.47.0
      - trl==0.21.0
      - peft==0.17.0
  clusterSelector:
    ray.io/cluster: gemma3-tuning

  2. Envía la definición de RayJob a tu RayCluster:

    envsubst < ray_job.yaml | kubectl apply -f -

  3. Verifica que haya un Pod nuevo en tu clúster:

    kubectl get pods

    Toma nota del nombre completo del Pod test-ray-job- que ves en el resultado. Este nombre es único para tu trabajo.

  4. Inspecciona el progreso de tu entrenamiento. Reemplaza gemma-training-ray-job-UNIQUE_ID por el nombre único del Pod que anotaste en el paso anterior.

    kubectl logs -f <test-ray-job-UNIQUE_ID>

    El resultado que ves es similar al siguiente:

    2025-08-20 08:29:34,966 INFO cli.py:41 -- Job submission server address: http://gemma3-tuning-head-svc.default.svc.cluster.local:8265
2025-08-20 08:29:34,991 SUCC cli.py:65 -- -----------------------------------------------
2025-08-20 08:29:34,991 SUCC cli.py:66 -- Job 'test-ray-job-82mm7' submitted successfully
2025-08-20 08:29:34,991 SUCC cli.py:67 -- -----------------------------------------------
2025-08-20 08:29:34,992 INFO cli.py:291 -- Next steps
2025-08-20 08:29:34,992 INFO cli.py:292 -- Query the logs of the job:
2025-08-20 08:29:34,992 INFO cli.py:294 -- ray job logs test-ray-job-82mm7
2025-08-20 08:29:34,992 INFO cli.py:296 -- Query the status of the job:
2025-08-20 08:29:34,992 INFO cli.py:298 -- ray job status test-ray-job-82mm7
2025-08-20 08:29:34,992 INFO cli.py:300 -- Request the job to be stopped:
2025-08-20 08:29:34,992 INFO cli.py:302 -- ray job stop test-ray-job-82mm7
2025-08-20 08:29:35,003 INFO cli.py:312 -- Tailing logs until the job exits (disable with --no-wait):
2025-08-20 08:29:34,982 INFO job_manager.py:531 -- Runtime env is setting up.
Starting training task!
Commencing training!
2025-08-20 08:30:08,498 INFO worker.py:1606 -- Using address 10.76.0.17:6379 set in the environment variable RAY_ADDRESS
2025-08-20 08:30:08,506 INFO worker.py:1747 -- Connecting to existing Ray cluster at address: 10.76.0.17:6379...
2025-08-20 08:30:08,527 INFO worker.py:1918 -- Connected to Ray cluster. View the dashboard at 10.76.0.17:8265
2025-08-20 08:30:08,701 INFO tune.py:253 -- Initializing Ray automatically. For cluster usage or custom Ray initialization, call `ray.init(...)` before `<FrameworkTrainer>(...)`.
2025-08-20 08:30:08,951 WARNING tune_controller.py:2132 -- The maximum number of pending trials has been automatically set to the number of available cluster CPUs, which is high (519 CPUs/pending trials). If you're running an experiment with a large number of trials, this could lead to scheduling overhead. In this case, consider setting the `TUNE_MAX_PENDING_TRIALS_PG` environment variable to the desired maximum number of concurrent pending trials.
2025-08-20 08:30:08,953 WARNING tune_controller.py:2132 -- The maximum number of pending trials has been automatically set to the number of available cluster CPUs, which is high (519 CPUs/pending trials). If you're running an experiment with a large number of trials, this could lead to scheduling overhead. In this case, consider setting the `TUNE_MAX_PENDING_TRIALS_PG` environment variable to the desired maximum number of concurrent pending trials.

View detailed results here: YOUR_GCS_BUCKET/gemma_vision_train_2025_08_20_08_30_07
To visualize your results with TensorBoard, run: `tensorboard --logdir /tmp/ray/session_2025-08-20_04-43-14_215096_1/artifacts/2025-08-20_08-30-08/gemma_vision_train_2025_08_20_08_30_07/driver_artifacts`

Training started with configuration:
╭──────────────────────────────────────────────────────────────────────╮
│ Training config                                                      │
├──────────────────────────────────────────────────────────────────────┤
│ train_loop_config/dataset_name                  ...-descriptions-vlm │
│ train_loop_config/gcs_bucket                    ...-bucket-yooo-west │
│ train_loop_config/gradient_accumulation_steps                      4 │
│ train_loop_config/learning_rate                               0.0002 │
│ train_loop_config/logging_steps                                   10 │
│ train_loop_config/lora_alpha                                      16 │
│ train_loop_config/lora_dropout                                  0.05 │
│ train_loop_config/lora_r                                          16 │
│ train_loop_config/max_seq_length                                 512 │
│ train_loop_config/model_id                      google/gemma-3-4b-it │
│ train_loop_config/num_train_epochs                                 3 │
│ train_loop_config/output_dir                    ...-4b-seo-optimized │
│ train_loop_config/per_device_train_batch_size                      1 │
│ train_loop_config/push_to_hub                                  False │
│ train_loop_config/save_steps                                     100 │
│ train_loop_config/save_strategy                                epoch │
╰──────────────────────────────────────────────────────────────────────╯
(RayTrainWorker pid=45455, ip=10.76.0.71) Setting up process group for: env:// [rank=0, world_size=16]
(TorchTrainer pid=45197, ip=10.76.0.71) Started distributed worker processes:
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45455) world_rank=0, local_rank=0, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45450) world_rank=1, local_rank=1, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45454) world_rank=2, local_rank=2, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45448) world_rank=3, local_rank=3, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45453) world_rank=4, local_rank=4, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45452) world_rank=5, local_rank=5, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45451) world_rank=6, local_rank=6, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=4c934ab2f646a578b03cc335586f30b943e811b645526a74c50bfca1, ip=10.76.0.71, pid=45449) world_rank=7, local_rank=7, node_rank=0
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45729) world_rank=8, local_rank=0, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45726) world_rank=9, local_rank=1, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45728) world_rank=10, local_rank=2, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45727) world_rank=11, local_rank=3, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45725) world_rank=12, local_rank=4, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45724) world_rank=13, local_rank=5, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45723) world_rank=14, local_rank=6, node_rank=1
(TorchTrainer pid=45197, ip=10.76.0.71) - (node_id=c0db52b44f891f3d6a1cedcbea4c6beb2c8434c66ef414dc15e65743, ip=10.76.0.135, pid=45722) world_rank=15, local_rank=7, node_rank=1

...

Training finished iteration 3 at 2025-08-20 08:40:43. Total running time: 10min 34s
╭─────────────────────────────────────────╮
│ Training result                         │
├─────────────────────────────────────────┤
│ checkpoint_dir_name   checkpoint_000002 │
│ time_this_iter_s               152.6374 │
│ time_total_s                  525.88585 │
│ training_iteration                    3 │
│ epoch                           2.75294 │
│ grad_norm                      47.27161 │
│ learning_rate                    0.0002 │
│ loss                            22.5275 │
│ mean_token_accuracy             0.90325 │
│ num_tokens                     1583017. │
│ step                                 60 │
╰─────────────────────────────────────────╯

...

Training completed after 3 iterations at 2025-08-20 08:40:52. Total running time: 10min 43s
2025-08-20 08:40:53,113 INFO tune.py:1009 -- Wrote the latest version of all result files and experiment state to 'YOUR_GCS_BUCKET/gemma_vision_train_2025_08_20_08_30_07' in 0.1663s.

2025-08-20 08:40:58,304 SUCC cli.py:65 -- ----------------------------------
2025-08-20 08:40:58,305 SUCC cli.py:66 -- Job 'test-ray-job-82mm7' succeeded
2025-08-20 08:40:58,305 SUCC cli.py:67 -- ----------------------------------

    Supervisa tu carga de trabajo

Puedes usar el panel en Ray para supervisar las cargas de trabajo programadas en tu clúster.

Para acceder a este panel, debes configurar el reenvío de puertos a tu clúster ejecutando el siguiente comando en una ventana de terminal nueva:

kubectl port-forward service/gemma3-tuning-head-svc 8265:8265 > fwd.log 2>&1 &

  1. Abre el siguiente vínculo en tu navegador: http://localhost:8265.

  2. De manera opcional, si usas Cloud Shell, después de ejecutar el comando del paso anterior, puedes hacer clic en el botón Vista previa en la Web.

    Selecciona la opción Cambiar puerto, ingresa 8265 y, luego, haz clic en Cambiar y obtener vista previa. El panel de Ray se abrirá en una pestaña nueva.

Realiza una limpieza

Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos usados en este instructivo, borra el proyecto que contiene los recursos o conserva el proyecto y borra los recursos individuales.

Borra tus recursos

  1. Para borrar el clúster de Ray y liberar el nodo potenciado por GPU, ejecuta el siguiente comando:

    kubectl delete -f ray_cluster.yaml
    GKE reduce automáticamente la escala de tu clúster y libera las máquinas A4 que usa Ray.

  2. Para borrar todo el clúster de GKE, ejecuta el siguiente comando:

    gcloud container clusters delete $CLUSTER_NAME \
    --region=$REGION

  3. Para borrar todo el bucket de Cloud Storage con todo su contenido, ejecuta el siguiente comando:

    gcloud storage rm -r gs://$GCS_BUCKET

Borra tu proyecto

Borra un Google Cloud proyecto:

gcloud projects delete PROJECT_ID

¿Qué sigue?