כוונון עדין של Gemma 4 באשכול A4 GKE

במדריך הזה נסביר איך לכוונן מודל שפה גדול (LLM) של Gemma 4 באשכול Google Kubernetes Engine ‏ (GKE) עם כמה יחידות GPU במארח יחיד ב- Google Cloud. האשכול הזה משתמש במכונה וירטואלית (VM) אחת מסוג A4 עם 8 מעבדי GPU של NVIDIA B200 שמצורפים אליה.

שני התהליכים העיקריים שמתוארים במדריך הזה הם:

  1. פריסת אשכול GKE עם ביצועים גבוהים באמצעות GKE Autopilot. כחלק מהפריסה הזו, יוצרים קובץ אימג' של מכונה וירטואלית בהתאמה אישית עם התוכנה הנדרשת שכבר מותקנת.
  2. אחרי פריסת האשכול, מריצים עבודת כוונון עדין מבוזרת באמצעות קבוצת הסקריפטים שמצורפת למדריך הזה. העבודה מתבססת על ספריית Hugging Face Accelerate.

המדריך הזה מיועד למהנדסי למידת מכונה (ML), לחוקרים, לאדמינים ולמפעילים של פלטפורמות ולמומחים בתחום הנתונים וה-AI שרוצים לפרוס אשכולות GKE ב- Google Cloud כדי לאמן מודלים גדולים של שפה (LLM).

מטרות

  • אפשר לגשת למודל Gemma 4 באמצעות Hugging Face.

  • מכינים את הסביבה.

  • יצירה ופריסה של אשכול GKE מסוג A4.

  • אפשר לבצע כוונון עדין של מודל Gemma 4 באמצעות ספריית Hugging Face Accelerate עם נתונים מקבילים שחולקו באופן מלא (FSDP).

  • עוקבים אחרי העבודה.

  • לפנות.

עלויות

במסמך הזה משתמשים ברכיבים הבאים של Google Cloud, והשימוש בהם כרוך בתשלום:

כדי ליצור הערכת עלויות בהתאם לשימוש החזוי, אתם יכולים להשתמש במחשבון התמחור.

משתמשים חדשים של Google Cloud ? יכול להיות שאתם זכאים לתקופת ניסיון בחינם.

לפני שמתחילים

  1. מפעילים את ממשקי ה-API הנדרשים, אם יש כאלה שלא מופעלים כבר:

    תפקידים שנדרשים להפעלת ממשקי API

    כדי להפעיל ממשקי API, צריך את ההרשאה serviceusage.services.enable. אם יצרתם את הפרויקט, סביר להניח שכבר יש לכם את ההרשאה הזו דרך התפקיד 'בעלים' (roles/owner). אחרת, תוכלו לקבל את ההרשאה הזו דרך התפקיד 'אדמין בממשק Service Usage' (roles/serviceusage.serviceUsageAdmin). איך מקצים תפקידים 

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

  2. מפעילים את חשבון השירות שמוגדר כברירת מחדל עבור 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

  3. מקצים את התפקידים שחשבון השירות שמוגדר כברירת מחדל צריך כדי להריץ את עומס העבודה של הכוונון העדין:

    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

  4. יצירת פרטי כניסה לאימות מקומי בחשבון המשתמש:

    ‫
    gcloud auth application-default login

  5. מפעילים את OS Login בפרויקט:

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

כדי לקבל את ההרשאות שדרושות להשלמת המדריך הזה, צריך לבקש מהאדמין להקצות לכם בפרויקט את תפקידי ה-IAM הבאים:

להסבר על מתן תפקידים, ראו איך מנהלים את הגישה ברמת הפרויקט, התיקייה והארגון.

יכול להיות שאפשר לקבל את ההרשאות הנדרשות גם באמצעות תפקידים בהתאמה אישית או תפקידים מוגדרים מראש.

גישה ל-Gemma 4 באמצעות Hugging Face

כדי להשתמש ב-Hugging Face כדי לגשת ל-Gemma 4:

  1. כניסה ל-Hugging Face
  2. יוצרים טוקן גישה ל-Hugging Face write.
    לוחצים על Your Profile > Settings > Access tokens > +Create new token
  3. מעתיקים את ערך האסימון write access ושומרים אותו. תשתמשו בו בהמשך המדריך הזה.

הכנת הסביבה

כדי להכין את הסביבה, מגדירים את הפרטים הבאים:

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

מחליפים את מה שכתוב בשדות הבאים:

  • ‫YOUR_PROJECT_ID: המזהה של Google Cloud הפרויקט שבו רוצים ליצור את אשכול GKE.

  • ‫YOUR_CLUSTER_NAME: השם של אשכול GKE שרוצים ליצור.

  • ‫YOUR_REGION: האזור שבו רוצים ליצור את אשכול GKE. אפשר ליצור את האשכול רק באזור שבו קיימת ההזמנה.

  • ‫YOUR_RESERVATION_NAME: המזהה של הקיבולת המוזמנת.

  • ‫YOUR_HF_TOKEN: טוקן הגישה ל-Hugging Face שיצרתם בקטע הקודם.

  • ‫YOUR_ARTIFACT_REGISTRY_LOCATION: Google Cloud האזור שבו רוצים ליצור את מאגר Artifact Registry. כדי לצמצם את זמן האחזור של שליפת התמונות, מומלץ להשתמש באותו אזור שבו השתמשתם בשביל YOUR_REGION.

יצירת אשכול GKE במצב אוטומטי

כדי ליצור אשכול GKE במצב אוטומטי, מריצים את הפקודה הבאה:

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

יצירת אשכול GKE עשויה להימשך זמן מה. כדי לוודא ש- Google Cloud סיים ליצור את האשכול, עוברים אל Kubernetes clusters במסוף Google Cloud .

יצירת סוד ב-Kubernetes לפרטי הכניסה של Hugging Face

כדי ליצור סוד של Kubernetes לפרטי הכניסה של Hugging Face, פועלים לפי השלבים הבאים:

  1. מגדירים את kubectl לתקשורת עם אשכול GKE:

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

  2. יוצרים סוד של Kubernetes לאחסון הטוקן של Hugging Face:

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

הכנת עומס העבודה

כדי להכין את עומס העבודה, מבצעים את הפעולות הבאות:

  1. יצירת סקריפטים של עומסי עבודה.

  2. שימוש ב-Docker וב-Cloud Build כדי ליצור קונטיינר לכוונון עדין.

יצירת סקריפטים של עומסי עבודה

כדי ליצור את הסקריפטים שעומס העבודה של הכוונון העדין משתמש בהם, מבצעים את הפעולות הבאות:

  1. יוצרים ספרייה לסקריפטים של עומס העבודה. משתמשים בספרייה הזו כספריית העבודה.

    mkdir llm-finetuning-gemma
cd llm-finetuning-gemma

  2. יוצרים את הקובץ cloudbuild.yaml כדי להשתמש ב-Google Cloud Build. הקובץ הזה יוצר את קונטיינר עומס העבודה ושומר אותו ב-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

  3. יוצרים קובץ Dockerfile כדי להפעיל את משימת הכוונון העדין:

    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

  4. יוצרים את קובץ ה-accel_fsdp_gemma4_config.yaml. קובץ ההגדרות הזה מכוון את Hugging Face Accelerate לפצל את משימת הכוונון בין שמונה מעבדי ה-GPU המקומיים במארח היחיד שלכם באמצעות 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

  5. יוצרים את הקובץ 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

  6. יוצרים את הקובץ 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()

שימוש ב-Docker וב-Cloud Build ליצירת קונטיינר לכוונון עדין

  1. יוצרים מאגר Docker ב-Artifact Registry:

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

  2. בספרייה llm-finetuning-gemma שיצרתם בשלב קודם, מריצים את הפקודה הבאה כדי ליצור את תמונת הכוונון העדין ולהעביר אותה אל Artifact Registry.

    gcloud builds submit . \
    --substitutions=_ARTIFACT_REPO_LOCATION="${ARTIFACT_REPO_LOCATION}"

  3. מייצאים את כתובת ה-URL של התמונה. תשתמשו בו בשלב מאוחר יותר במדריך הזה:

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

התחלת עומס העבודה של הכוונון העדין

כדי להתחיל את עומס העבודה של הכוונון העדין:

  1. מחילים את מניפסט הכוונון העדין כדי ליצור את משימת הכוונון העדין:

    envsubst < finetune.yaml | kubectl apply -f -

    מכיוון שאתם משתמשים באשכולות במצב GKE Autopilot, יכול להיות שיעברו כמה דקות עד שהצומת עם ה-GPU יתחיל לפעול.

  2. מריצים את הפקודה הבאה כדי לעקוב אחרי העבודה:

    watch kubectl get pods

  3. אחרי שה-Pods יפעלו, בודקים את היומנים של העבודה:

    kubectl logs job.batch/finetune-job -f

    משאב המשימה מוריד את נתוני המודל ואז מבצע כוונון עדין של המודל בכל שמונת המעבדים הגרפיים. ההורדה נמשכת כחמש דקות. אחרי שההורדה מסתיימת, תהליך הכוונון העדין נמשך כשעתיים וחצי.

מעקב אחר עומס העבודה

אתם יכולים לעקוב אחר השימוש במעבדי ה-GPU באשכול GKE כדי לוודא שעבודת הכוונון העדין פועלת ביעילות. כדי לעשות זאת, פותחים את הקישור הבא בדפדפן:

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

כשעוקבים אחרי עומס העבודה, אפשר לראות את הנתונים הבאים:

  • השימוש ביחידות GPU: כדי שהאימון יהיה תקין, צפוי שהשימוש בכל 8 יחידות ה-GPU יעלה ויתייצב ברמה גבוהה לאורך כל האימון.
  • משך העבודה: העבודה אמורה להימשך כשעתיים וחצי באשכול A4 שצוין.

הסרת המשאבים

כדי להימנע מחיובים נוספים, מוחקים את המשאבים שנוצרו במהלך המדריך הזה.

מחיקת המשאבים

  1. כדי למחוק את משימת הכוונון העדין, מריצים את הפקודה הבאה:

    kubectl delete job finetune-job

  2. כדי למחוק את אשכול GKE, מריצים את הפקודה הבאה:

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

המאמרים הבאים