Google utilise la technologie IA pour traduire le contenu dans votre langue préférée. Les traductions générées par IA peuvent contenir des erreurs.

Affiner Gemma 4 sur un cluster A4 GKE

GPU en cluster

Ce tutoriel explique comment affiner un grand modèle de langage (LLM) Gemma 4 sur un cluster Google Kubernetes Engine (GKE) à hôte unique et GPU multiples sur Google Cloud. Ce cluster utilise une seule instance de machine virtuelle (VM) A4 avec huit GPU NVIDIA B200 associés.

Les deux principaux processus décrits dans ce tutoriel sont les suivants :

Déployez un cluster GKE hautes performances à l'aide de GKE Autopilot. Dans le cadre de ce déploiement, vous allez créer une image de VM personnalisée avec le logiciel nécessaire préinstallé.
Une fois le cluster déployé, vous exécutez un job d'affinage distribué à l'aide de l'ensemble de scripts qui accompagnent ce tutoriel. Le job utilise la bibliothèque Hugging Face Accelerate.

Ce tutoriel est destiné aux ingénieurs en machine learning (ML), aux chercheurs, aux administrateurs et opérateurs de plate-forme, ainsi qu'aux spécialistes des données et de l'IA qui souhaitent déployer des clusters GKE sur Google Cloud pour entraîner des LLM.

Objectifs

Accédez au modèle Gemma 4 à l'aide de Hugging Face.
Préparez votre environnement.
Créez et déployez un cluster GKE A4.
Ajustez le modèle Gemma 4 à l'aide de la bibliothèque Hugging Face Accelerate avec le parallélisme des données entièrement segmentées (FSDP).
surveiller votre job ;
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

Activez les API requises, si certaines ne le sont pas déjà :

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

Activez le compte de service par défaut pour votre projet 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

Attribuez les rôles dont le compte de service par défaut a besoin pour exécuter la charge de travail d'affinage :
```
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
```
Remarque : Cette étape n'accorde que les rôles prédéfinis requis par ce tutoriel. Si vous attribuez un rôle de base tel qu'Éditeur (roles/editor), le compte de service dispose de beaucoup plus d'autorisations que nécessaire, et les règles de votre organisation peuvent bloquer l'attribution. Pour en savoir plus, consultez Bonnes pratiques d'utilisation des comptes de service.
Créez des identifiants d'authentification locaux pour votre compte utilisateur :
```
gcloud auth application-default login
```
Remarque : Si vous utilisez un shell local et un fournisseur d'identité (IdP) externe, et que vous rencontrez une erreur d'authentification après avoir exécuté la commande précédente, connectez-vous à la gcloud CLI avec votre identité fédérée.

Activez OS Login pour votre projet :

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

Pour obtenir les autorisations nécessaires pour suivre ce tutoriel, demandez à votre administrateur de vous accorder les rôles IAM suivants sur votre projet :

Administrateur Kubernetes Engine (roles/container.admin)
Administrateur de Compute (roles/compute.admin)
Administrateur de l'espace de stockage (roles/storage.admin)
Administrateur Artifact Registry (roles/artifactregistry.admin)
Éditeur Cloud Build (roles/cloudbuild.builds.editor)
Utilisateur du compte de service (roles/iam.serviceAccountUser)
Administrateur de compte de service (roles/iam.serviceAccountAdmin)
Administrateur de projet IAM (roles/resourcemanager.projectIamAdmin)
Administrateur Service Usage (roles/serviceusage.serviceUsageAdmin)

Pour en savoir plus sur l'attribution de rôles, consultez la page Gérer l'accès aux projets, aux dossiers et aux organisations.

Vous pouvez également obtenir les autorisations requises avec des rôles personnalisés ou d'autres rôles prédéfinis.

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

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

Se connecter à Hugging Face
Créez un jeton d'accès write Hugging Face.
Cliquez sur Votre profil > Paramètres > Jetons d'accès > + Créer un jeton.
Copiez et enregistrez la valeur du jeton write access. Vous en aurez besoin plus loin dans ce tutoriel.

Préparer votre environnement

Pour préparer votre environnement, définissez les éléments suivants :

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

Remplacez les éléments suivants :

YOUR_PROJECT_ID : ID du Google Cloud projet dans lequel vous souhaitez créer le cluster GKE.
YOUR_CLUSTER_NAME : nom du cluster GKE à créer.
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_RESERVATION_NAME : identifiant de votre capacité réservée.
YOUR_HF_TOKEN : jeton d'accès Hugging Face que vous avez créé dans la section précédente.
YOUR_ARTIFACT_REGISTRY_LOCATION : Google Cloud région dans laquelle vous souhaitez créer votre dépôt Artifact Registry. Pour minimiser la latence d'extraction des images, nous vous recommandons d'utiliser la même région que celle que vous avez utilisée pour YOUR_REGION.

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}" \
    --project="${PROJECT_ID}" \
    --location="${CLUSTER_REGION}"

La création du cluster GKE peut prendre un certain temps. Pour vérifier que 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

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

Configurez kubectl pour communiquer avec votre cluster GKE :

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

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 -

Préparer votre charge de travail

Pour préparer votre charge de travail, procédez comme suit :

Créez des scripts de charge de travail.
Utilisez Docker et Cloud Build pour créer un conteneur d'affinage.

Créer des scripts de charge de travail

Pour créer les scripts utilisés par votre charge de travail d'affinage, procédez comme suit :

Créez un répertoire pour les scripts de charge de travail. Utilisez ce répertoire comme répertoire de travail.
```
mkdir llm-finetuning-gemma
cd llm-finetuning-gemma
```

Créez le fichier cloudbuild.yaml pour utiliser Google Cloud Build. Ce fichier crée le conteneur de charge de travail et le stocke dans 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

Créez un fichier Dockerfile pour exécuter le job d'affinage :

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

Créez le fichier accel_fsdp_gemma4_config.yaml. Ce fichier de configuration indique à Hugging Face Accelerate de répartir le job d'optimisation sur les huit GPU locaux de votre hôte unique à l'aide de 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

Créez le fichier 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

Créez le fichier 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()

Créer un conteneur d'affinage à l'aide de Docker et Cloud Build

Créez un dépôt Docker Artifact Registry :

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

Dans le répertoire llm-finetuning-gemma que vous avez créé lors d'une étape précédente, exécutez la commande suivante pour créer l'image d'affinage et la transférer vers Artifact Registry.
```
gcloud builds submit . \
    --substitutions=_ARTIFACT_REPO_LOCATION="${ARTIFACT_REPO_LOCATION}"
```

Exportez l'URL de l'image. Vous l'utiliserez lors d'une prochaine étape de ce tutoriel :

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

Démarrer votre charge de travail d'affinage

Pour démarrer votre charge de travail d'affinage :

Appliquez le fichier manifeste d'affinage pour créer le job d'affinage :
```
envsubst < finetune.yaml | kubectl apply -f -
```
Comme vous utilisez des clusters en mode GKE Autopilot, le démarrage de votre nœud compatible avec les GPU peut prendre quelques minutes.
Surveillez le job en exécutant la commande suivante :
```
watch kubectl get pods
```
Une fois les pods en cours d'exécution, vérifiez les journaux du job :
```
kubectl logs job.batch/finetune-job -f
```
La ressource du job télécharge les données du modèle, puis affine le modèle sur les huit GPU. Le téléchargement prend environ cinq minutes. Une fois le téléchargement terminé, le processus d'affinage prend environ deux heures et demie.

Surveiller votre charge de travail

Vous pouvez surveiller l'utilisation des GPU dans votre cluster GKE pour vérifier que votre job de réglage fin s'exécute efficacement. Pour ce faire, ouvrez le lien suivant dans votre navigateur :

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\"))"

Lorsque vous surveillez votre charge de travail, vous pouvez voir les éléments suivants :

Utilisation des GPU : pour une tâche d'affinage saine, vous pouvez vous attendre à ce que l'utilisation de vos huit GPU augmente et se stabilise à un niveau élevé tout au long de votre entraînement.
Durée du job : l'exécution du job devrait prendre environ deux heures et 30 minutes sur le cluster A4 spécifié.

Effectuer un nettoyage

Pour éviter des frais supplémentaires, supprimez les ressources créées au cours de ce tutoriel.

Supprimer vos ressources

Pour supprimer votre job d'affinage, exécutez la commande suivante :
```
kubectl delete job finetune-job
```

Pour supprimer votre cluster GKE, exécutez la commande suivante :

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