Utiliser Ray pour affiner Gemma 3 pour les tâches de vision sur GKE

Ce tutoriel explique comment affiner un modèle Gemma 3 à l'aide du framework Ray sur un cluster GKE à nœuds multiples. Le cluster utilise deux instances de machines virtuelles A4, chacune dotée de huit GPU NVIDIA B200.

Le contenu de ce tutoriel est divisé en deux parties :

  1. Préparer le cluster Ray sur un cluster GKE Autopilot.
  2. Exécution d'un job d'entraînement distribué, utilisant deux instances A4, avec huit GPU B200 chacune.

Ce tutoriel s'adresse aux ingénieurs et chercheurs en machine learning (ML), aux administrateurs et opérateurs de plate-forme, ainsi qu'aux spécialistes des données et de l'IA qui souhaitent distribuer une charge de travail d'IA sur plusieurs nœuds et GPU.

Objectifs

  • Accédez à un modèle Gemma 3 en utilisant Hugging Face.

  • Préparez votre environnement.

  • Créez un cluster GKE Autopilot sur lequel l'opérateur Ray est installé.

  • Configurez le cluster Ray sur le cluster GKE pour qu'il accepte les tâches Ray.

  • Configurez et exécutez un job Ray qui ajuste le modèle Gemma 3 en fonction de l'entrée visuelle.

  • Surveillez votre charge de travail.

  • Effectuer un nettoyage.

Coûts

Dans ce document, vous utilisez les composants facturables suivants de Google Cloud :

Pour obtenir une estimation des coûts en fonction de votre utilisation prévue, utilisez le simulateur de coût.

Les nouveaux utilisateurs de Google Cloud peuvent bénéficier d'un essai sans frais.

Avant de commencer

  1. Connectez-vous à votre compte Google Cloud . Si vous débutez sur Google Cloud, créez un compte pour évaluer les performances de nos produits en conditions réelles. Les nouveaux clients bénéficient également de 300 $ de crédits sans frais pour exécuter, tester et déployer des charges de travail.

  2. Installez la Google Cloud CLI.

  3. Si vous utilisez un fournisseur d'identité (IdP) externe, vous devez d'abord vous connecter à la gcloud CLI avec votre identité fédérée.

  4. Pour initialiser la gcloud CLI, exécutez la commande suivante : 

    gcloud init

  5. Créez ou sélectionnez un projet Google Cloud .

    Rôles requis pour sélectionner ou créer un projet

    • Sélectionnez un projet : la sélection d'un projet ne nécessite pas de rôle IAM spécifique. Vous pouvez sélectionner n'importe quel projet pour lequel un rôle vous a été attribué.
    • Créer un projet : pour créer un projet, vous devez disposer du rôle Créateur de projet (roles/resourcemanager.projectCreator), qui contient l'autorisation resourcemanager.projects.create. Découvrez comment attribuer des rôles.

    • Créez un projet Google Cloud  :

      gcloud projects create PROJECT_ID

      Remplacez PROJECT_ID par le nom du projet Google Cloud que vous créez.

    • Sélectionnez le projet Google Cloud que vous avez créé :

      gcloud config set project PROJECT_ID

      Remplacez PROJECT_ID par le nom de votre projet Google Cloud .

  6. Vérifiez que la facturation est activée pour votre projet Google Cloud .

  7. Activez l'API requise :

    Rôles requis pour activer les API

    Pour activer les API, vous devez disposer de l'autorisation serviceusage.services.enable. Si vous avez créé le projet, vous disposez probablement déjà de cette autorisation grâce au rôle Propriétaire (roles/owner). Sinon, vous pouvez obtenir cette autorisation grâce au rôle Administrateur Service Usage (roles/serviceusage.serviceUsageAdmin). Découvrez comment attribuer des rôles. 

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

  8. Installez la Google Cloud CLI.

  9. Si vous utilisez un fournisseur d'identité (IdP) externe, vous devez d'abord vous connecter à la gcloud CLI avec votre identité fédérée.

  10. Pour initialiser la gcloud CLI, exécutez la commande suivante : 

    gcloud init

  11. Créez ou sélectionnez un projet Google Cloud .

    Rôles requis pour sélectionner ou créer un projet

    • Sélectionnez un projet : la sélection d'un projet ne nécessite pas de rôle IAM spécifique. Vous pouvez sélectionner n'importe quel projet pour lequel un rôle vous a été attribué.
    • Créer un projet : pour créer un projet, vous devez disposer du rôle Créateur de projet (roles/resourcemanager.projectCreator), qui contient l'autorisation resourcemanager.projects.create. Découvrez comment attribuer des rôles.

    • Créez un projet Google Cloud  :

      gcloud projects create PROJECT_ID

      Remplacez PROJECT_ID par le nom du projet Google Cloud que vous créez.

    • Sélectionnez le projet Google Cloud que vous avez créé :

      gcloud config set project PROJECT_ID

      Remplacez PROJECT_ID par le nom de votre projet Google Cloud .

  12. Vérifiez que la facturation est activée pour votre projet Google Cloud .

  13. Activez l'API requise :

    Rôles requis pour activer les API

    Pour activer les API, vous devez disposer de l'autorisation serviceusage.services.enable. Si vous avez créé le projet, vous disposez probablement déjà de cette autorisation grâce au rôle Propriétaire (roles/owner). Sinon, vous pouvez obtenir cette autorisation grâce au rôle Administrateur Service Usage (roles/serviceusage.serviceUsageAdmin). Découvrez comment attribuer des rôles. 

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

  14. Attribuez des rôles à votre compte utilisateur. Exécutez la commande suivante une fois pour chacun des rôles IAM suivants : 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

    Remplacez les éléments suivants :

    • PROJECT_ID : ID de votre projet
    • USER_IDENTIFIER : identifiant de votre compte utilisateur . Par exemple, myemail@example.com.
    • ROLE : rôle IAM que vous accordez à votre compte utilisateur.
  15. Activez le compte de service par défaut pour votre projet Google Cloud  : 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --project=PROJECT_ID

    Remplacez PROJECT_NUMBER par votre numéro de projet. Pour consulter le numéro de votre projet, consultez Obtenir un projet existant.

  16. Attribuez le rôle Éditeur (roles/editor) au compte de service par défaut : 
    gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \
    --role=roles/editor
  17. Créez des identifiants d'authentification locaux pour votre compte utilisateur : 
    gcloud auth application-default login
  18. Connectez-vous à votre compte Hugging Face ou créez-en un.

Accéder à Gemma 3 à l'aide de Hugging Face

Pour utiliser Hugging Face afin d'accéder à Gemma 3, procédez comme suit :

  1. Signez le contrat de consentement pour utiliser Gemma 3.

  2. Créez un jeton read access Hugging Face.

  3. Copiez et enregistrez la valeur du jeton read access. Vous en aurez besoin dans la suite de ce tutoriel.

Préparer votre environnement

Préparez votre environnement en configurant les paramètres nécessaires et en définissant les variables d'environnement.

Exécutez la commande suivante :

export PROJECT_ID="YOUR_PROJECT_ID"
export RESERVATION="RESERVATION_ID"
export REGION="YOUR_REGION"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export HF_TOKEN="HUGGING_FACE_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

Remplacez les éléments suivants :

  • YOUR_PROJECT_ID : nom du Google Cloud projet dans lequel vous souhaitez créer le cluster GKE.
  • YOUR_RESERVATION_ID : URL de la réservation que vous souhaitez utiliser pour créer votre cluster. En fonction du projet dans lequel la réservation existe, spécifiez l'une des valeurs suivantes :
    • La réservation existe dans votre projet : spécifiez le nom de la réservation (par exemple, my-reservation).
    • La réservation existe dans un autre projet : spécifiez le chemin d'accès complet au format projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME.
  • YOUR_REGION : région dans laquelle vous souhaitez créer votre cluster GKE. Vous ne pouvez créer le cluster que dans la région où se trouve votre réservation.
  • YOUR_CLUSTER_NAME : nom du cluster GKE à créer.
  • HUGGING_FACE_TOKEN : jeton Hugging Face que vous avez créé à une étape précédente.
  • YOUR_RAY_SA : nom du compte de service dans le cluster Kubernetes.
  • YOUR_GSA_NAME : nom du compte de service Google.
  • YOUR_GCS_BUCKET : nom du bucket dans lequel vous stockez les résultats du point de contrôle de l'entraînement.

Créer un cluster GKE en mode Autopilot

Pour créer un cluster GKE en mode Autopilot, exécutez la commande suivante :

gcloud container clusters create-auto $CLUSTER_NAME \
    --enable-ray-operator \
    --enable-ray-cluster-monitoring \
    --enable-ray-cluster-logging \
    --location=$REGION

La création du cluster GKE peut prendre un certain temps. Pour vérifier si Google Cloud a terminé de créer votre cluster, accédez à Clusters Kubernetes dans la console Google Cloud .

Créer un secret Kubernetes pour les identifiants Hugging Face

Dans Cloud Shell, pour créer un secret Kubernetes pour les identifiants Hugging Face, procédez comme suit :

  1. Configurez kubectl pour vous connecter à votre cluster :

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

  2. Créez un secret Kubernetes pour stocker votre jeton Hugging Face :

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

Créer le bucket Cloud Storage

Si vous souhaitez utiliser un nouveau bucket pour stocker vos artefacts d'entraînement, exécutez la commande suivante :

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

Si vous souhaitez utiliser un bucket existant, vous pouvez ignorer cette étape. Toutefois, vous devez vous assurer que votre bucket se trouve dans la même région que votre cluster.

Créer un compte de service IAM

Dans Cloud Shell, pour créer un compte de service IAM (également appelé compte de service Google ou GSA) et lui accorder des autorisations d'accès à votre bucket Cloud Storage, procédez comme suit :

  1. Créez un compte de service 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. Attribuez le rôle Administrateur de l'espace de stockage (roles/storage.admin) à votre compte de service IAM pour votre bucket 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"

Créer un compte de service Kubernetes

Dans Cloud Shell, créez un compte de service Kubernetes et configurez Workload Identity pour accorder aux pods Ray l'accès aux ressources Google Cloud  :

  1. Créez un compte de service Kubernetes : 

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

  2. Associez votre compte de service Kubernetes au compte de service IAM que vous avez créé précédemment pour activer 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. Annotez votre compte de service Kubernetes avec l'adresse e-mail du compte de service IAM : 

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

Enregistrer votre code d'entraînement en tant que ConfigMap

Pour éviter d'intégrer votre script d'entraînement dans une image de conteneur, stockez-le en tant que ConfigMap dans votre cluster. Cette ConfigMap est montée sur les systèmes de fichiers du pod, ce qui vous permet de mettre à jour le script d'entraînement sans avoir à recréer l'intégralité du cluster Ray.

Pour stocker votre script d'entraînement en tant que ConfigMap dans votre cluster, procédez comme suit :

  1. Créez un répertoire nommé code, puis créez un fichier nommé vision_train.py dans ce répertoire.

  2. Copiez le code suivant dans le fichier 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. Enregistrez le fichier.

  4. Créez un objet ConfigMap dans votre cluster :

    kubectl create cm ray-job-cm --from-file=code -o yaml --dry-run=client | kubectl apply -f -
    Pour mettre à jour le script d'entraînement, réexécutez la commande précédente. Il peut s'écouler une minute avant que les modifications ne soient propagées à tous les pods.

Configurer un cluster Ray

  1. Pour créer un cluster Ray dans votre cluster GKE, enregistrez le fichier YAML suivant sous le nom 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. Appliquez cette définition YAML à votre cluster à l'aide de la commande suivante :

    envsubst < ray_cluster.yaml | kubectl apply -f -
    L'indicateur $RESERVATION est automatiquement remplacé par le nom que vous avez configuré en tant que variable d'environnement.

    L'opérateur Ray crée les pods raylet, ce qui déclenche l'autoscaling du cluster pour fournir à ces pods les nœuds appropriés. Trois pods sont créés dans votre cluster : un nœud principal et deux nœuds de calcul. Les nœuds de calcul sont équipés de GPU B200.

  3. Pour vérifier que les trois pods sont prêts, exécutez la commande suivante :

    kubectl get pods
    La liste des pods d'un cluster Ray prêt ressemble à ce qui suit :

    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

Planifier un job d'entraînement

  1. Enregistrez le code suivant dans un fichier 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. Envoyez la définition RayJob à votre RayCluster :

    envsubst < ray_job.yaml | kubectl apply -f -

  3. Vérifiez qu'un nouveau pod se trouve dans votre cluster :

    kubectl get pods

    Notez le nom complet du pod test-ray-job- qui s'affiche dans le résultat. Ce nom est propre à votre job.

  4. Inspectez la progression de votre entraînement. Remplacez gemma-training-ray-job-UNIQUE_ID par le nom unique du pod que vous avez noté à l'étape précédente.

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

    Le résultat renvoyé ressemble à ceci :

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

    Surveiller votre charge de travail

Vous pouvez utiliser le tableau de bord de Ray pour surveiller les charges de travail planifiées dans votre cluster.

Pour accéder à ce tableau de bord, vous devez configurer le transfert de port vers votre cluster en exécutant la commande suivante dans une nouvelle fenêtre de terminal :

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

  1. Ouvrez le lien suivant dans votre navigateur : http://localhost:8265.

  2. Si vous utilisez Cloud Shell, vous pouvez éventuellement cliquer sur le bouton Aperçu sur le Web après avoir exécuté la commande de l'étape précédente.

    Sélectionnez l'option Modifier le port, saisissez 8265, puis cliquez sur Modifier et prévisualiser. Le tableau de bord Ray s'ouvre dans un nouvel onglet.

Effectuer un nettoyage

Pour éviter que les ressources utilisées lors de ce tutoriel soient facturées sur votre compte Google Cloud, supprimez le projet contenant les ressources, ou conservez le projet et supprimez les ressources individuelles.

Supprimer vos ressources

  1. Pour supprimer le cluster Ray et libérer le nœud équipé d'un GPU, exécutez la commande suivante :

    kubectl delete -f ray_cluster.yaml
    GKE réduit automatiquement la taille de votre cluster et libère les machines A4 utilisées par Ray.

  2. Pour supprimer l'intégralité du cluster GKE, exécutez la commande suivante :

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

  3. Pour supprimer l'intégralité du bucket Cloud Storage et de son contenu, exécutez la commande suivante :

    gcloud storage rm -r gs://$GCS_BUCKET

Supprimer votre projet

Supprimer un projet Google Cloud  :

gcloud projects delete PROJECT_ID

Étapes suivantes