שימוש ב-Ray כדי לבצע כוונון עדין של Gemma 3 למשימות ראייה ב-GKE

במדריך הזה נסביר איך לכוונן מודל Gemma 3 באמצעות מסגרת Ray באשכול GKE מרובה צמתים. האשכול משתמש בשתי מכונות וירטואליות (VM) מסוג A4, שלכל אחת מהן מצורפים שמונה מעבדי GPU מסוג NVIDIA B200.

התוכן של המדריך הזה מחולק לשני חלקים:

  1. הכנת אשכול Ray על גבי אשכול GKE Autopilot.
  2. הפעלת משימת אימון מבוזרת, באמצעות 2 מופעי A4, עם 8 יחידות GPU מסוג B200 בכל אחד.

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

מטרות

  • גישה למודל Gemma 3 באמצעות Hugging Face.

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

  • יוצרים אשכול GKE Autopilot עם Ray Operator מותקן.

  • מגדירים את אשכול Ray באשכול GKE כך שיקבל משימות Ray.

  • הגדרה והרצה של משימת Ray שמכוונת את מודל Gemma 3 על סמך קלט ויזואלי.

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

  • לפנות.

עלויות

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

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

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

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

  1. נכנסים לחשבון Google Cloud . אם אתם משתמשים חדשים ב- Google Cloud, צרו חשבון כדי שתוכלו להעריך את הביצועים של המוצרים שלנו בתרחישים מהעולם האמיתי. לקוחות חדשים מקבלים בחינם גם קרדיט בשווי 300$ להרצה, לבדיקה ולפריסה של עומסי העבודה.
  2. התקינו את ה-CLI של Google Cloud.

  3. אם אתם משתמשים בספק זהויות חיצוני (IdP), קודם אתם צריכים להיכנס ל-CLI של gcloud באמצעות המאגר המאוחד לניהול זהויות.

  4. כדי לאתחל את ה-CLI של gcloud, הריצו את הפקודה הבאה:

    gcloud init
  5. יוצרים או בוחרים Google Cloud פרויקט.

    תפקידים שנדרשים כדי לבחור או ליצור פרויקט

    • Select a project: כדי לבחור פרויקט לא צריך תפקיד IAM ספציפי – אפשר לבחור כל פרויקט שקיבלתם בו תפקיד.
    • יצירת פרויקט: כדי ליצור פרויקט, צריך את התפקיד Project Creator (יצירת פרויקטים) (roles/resourcemanager.projectCreator), שכולל את ההרשאה resourcemanager.projects.create. איך מקצים תפקידים
    • יוצרים Google Cloud פרויקט:

      gcloud projects create PROJECT_ID

      מחליפים את PROJECT_ID בשם של פרויקט Google Cloud שיוצרים.

    • בוחרים את הפרויקט שיצרתם: Google Cloud

      gcloud config set project PROJECT_ID

      מחליפים את PROJECT_ID בשם הפרויקט ב- Google Cloud .

  6. מוודאים שהחיוב מופעל בפרויקט Google Cloud .

  7. מפעילים את ממשק ה-API הנדרש:

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

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

    gcloud services enable  compute.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com container.googleapis.com
  8. התקינו את ה-CLI של Google Cloud.

  9. אם אתם משתמשים בספק זהויות חיצוני (IdP), קודם אתם צריכים להיכנס ל-CLI של gcloud באמצעות המאגר המאוחד לניהול זהויות.

  10. כדי לאתחל את ה-CLI של gcloud, הריצו את הפקודה הבאה:

    gcloud init
  11. יוצרים או בוחרים Google Cloud פרויקט.

    תפקידים שנדרשים כדי לבחור או ליצור פרויקט

    • Select a project: כדי לבחור פרויקט לא צריך תפקיד IAM ספציפי – אפשר לבחור כל פרויקט שקיבלתם בו תפקיד.
    • יצירת פרויקט: כדי ליצור פרויקט, צריך את התפקיד Project Creator (יצירת פרויקטים) (roles/resourcemanager.projectCreator), שכולל את ההרשאה resourcemanager.projects.create. איך מקצים תפקידים
    • יוצרים Google Cloud פרויקט:

      gcloud projects create PROJECT_ID

      מחליפים את PROJECT_ID בשם של פרויקט Google Cloud שיוצרים.

    • בוחרים את הפרויקט שיצרתם: Google Cloud

      gcloud config set project PROJECT_ID

      מחליפים את PROJECT_ID בשם הפרויקט ב- Google Cloud .

  12. מוודאים שהחיוב מופעל בפרויקט Google Cloud .

  13. מפעילים את ממשק ה-API הנדרש:

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

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

    gcloud services enable  compute.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com container.googleapis.com
  14. מעניקים תפקידים לחשבון המשתמש. מריצים את הפקודה הבאה לכל אחד מהתפקידים הבאים ב-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

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

    • PROJECT_ID: מזהה הפרויקט.
    • USER_IDENTIFIER: המזהה של חשבון המשתמש . לדוגמה, myemail@example.com.
    • ROLE: תפקיד ה-IAM שאתם מקצים לחשבון המשתמש.
  15. מפעילים את חשבון השירות שמוגדר כברירת מחדל עבור Google Cloud הפרויקט:
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com \
        --project=PROJECT_ID

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

  16. מקצים את התפקיד 'עריכה' (roles/editor) לחשבון השירות שמוגדר כברירת מחדל:
    gcloud projects add-iam-policy-binding PROJECT_ID \
        --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \
        --role=roles/editor
  17. יצירת פרטי כניסה לאימות מקומי בחשבון המשתמש:
    gcloud auth application-default login
  18. נכנסים לחשבון Hugging Face או יוצרים חשבון.

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

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

  1. חתימה על הסכם ההסכמה לשימוש ב-Gemma 3

  2. יוצרים טוקן של Hugging Face read access.

  3. מעתיקים ושומרים את ערך הטוקן read access. תשתמשו בו בהמשך המדריך הזה.

הכנת הסביבה

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

מריצים את הפקודה הבאה:

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

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

  • YOUR_PROJECT_ID: השם של Google Cloud הפרויקט שבו רוצים ליצור את אשכול GKE.
  • YOUR_RESERVATION_ID: כתובת ה-URL של ההזמנה שבה רוצים להשתמש כדי ליצור את האשכול. בהתאם לפרויקט שבו קיימת ההזמנה, מציינים אחד מהערכים הבאים:
    • ההזמנה קיימת בפרויקט: מציינים את שם ההזמנה (למשל, my-reservation).
    • המקום השמור קיים בפרויקט אחר: מציינים את הנתיב המלא בפורמט projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME.
  • YOUR_REGION: האזור שבו רוצים ליצור את אשכול GKE. אפשר ליצור את האשכול רק באזור שבו קיימת ההזמנה.
  • YOUR_CLUSTER_NAME: השם של אשכול GKE שרוצים ליצור.
  • YOUR_HF_TOKEN: האסימון של Hugging Face שיצרתם בשלב קודם.
  • YOUR_RAY_SA: השם של חשבון השירות בתוך אשכול Kubernetes.
  • YOUR_GSA_NAME: השם של חשבון השירות של Google.
  • YOUR_GCS_BUCKET: השם של הקטגוריה שבה מאחסנים את התוצאות מנקודת הבדיקה של האימון.

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

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

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

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

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

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

  1. מגדירים את kubectl כך שיתחבר לאשכול:

    gcloud container clusters get-credentials $CLUSTER_NAME \
        --region=$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 -

יצירת קטגוריה של Cloud Storage

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

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

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

יצירת חשבון שירות ב-IAM

כדי ליצור ב-Cloud Shell חשבון שירות IAM (שנקרא גם חשבון שירות של Google או GSA) ולהעניק לו הרשאות גישה לקטגוריה של Cloud Storage, מבצעים את הפעולות הבאות:

  1. יוצרים חשבון שירות ב-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. מקצים לחשבון השירות שלכם ב-IAM את התפקיד 'אדמין של Storage' (roles/storage.admin) עבור הקטגוריה שלכם ב-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"

יצירת חשבון שירות של Kubernetes

ב-Cloud Shell, יוצרים חשבון שירות של Kubernetes ומגדירים Workload Identity כדי להעניק לקבוצות Pod של Ray גישה למשאבים: Google Cloud

  1. יוצרים חשבון שירות ב-Kubernetes:

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

  2. כדי להפעיל את Workload Identity, צריך לקשר את חשבון השירות של Kubernetes לחשבון השירות של IAM שיצרתם קודם:

    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. מוסיפים לחשבון השירות של Kubernetes את כתובת האימייל של חשבון השירות ב-IAM:

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

שמירת קוד האימון כ-ConfigMap

כדי להימנע מהצורך להטמיע את סקריפט ההדרכה בקובץ אימג' של קונטיינר, אפשר לאחסן אותו כ-ConfigMap באשכול. ה-ConfigMap הזה מותקן במערכות הקבצים של ה-Pod, וכך מאפשר לעדכן את סקריפט האימון בלי ליצור מחדש את כל אשכול Ray.

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

  1. יוצרים ספרייה בשם code, ובספרייה הזו יוצרים קובץ בשם vision_train.py.

  2. מעתיקים את הקוד הבא לקובץ 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. שומרים את הקובץ.

  4. יוצרים אובייקט ConfigMap באשכול:

    kubectl create cm ray-job-cm --from-file=code -o yaml --dry-run=client | kubectl apply -f -
    כדי לעדכן את סקריפט האימון, מריצים מחדש את הפקודה הקודמת. יכול להיות שיעבור זמן מה עד שהשינויים יתעדכנו בכל הפודים.

הגדרת אשכול Ray

  1. כדי ליצור אשכול Ray באשכול GKE, שומרים את קובץ ה-YAML הבא כקובץ בשם 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. מחילים את הגדרת ה-YAML הזו על האשכול באמצעות הפקודה הבאה:

    envsubst < ray_cluster.yaml | kubectl apply -f -
    הדגל $RESERVATION מוחלף אוטומטית בשם שהגדרתם כמשתנה סביבה.

    ‫Ray Operator יוצר את ה-pods של raylet, שבתורם מפעילים את ההתאמה האוטומטית לעומס (autoscaling) של האשכול כדי לספק לאותם pods את הצמתים המתאימים. שלושה פודים נוצרים באשכול: צומת ראש אחד ושני צומתי worker. צמתי העובדים מצוידים במעבדי B200 GPU.

  3. כדי לוודא שכל שלושת ה-pods מוכנים, מריצים את הפקודה הבאה:

    kubectl get pods
    רשימת ה-pods של אשכול Ray מוכן נראית בערך כך:

    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
    

תזמון משימת אימון

  1. שומרים את הקובץ הבא כקובץ 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. שולחים את הגדרת RayJob אל RayCluster:

    envsubst < ray_job.yaml | kubectl apply -f -
  3. בודקים אם יש Pod חדש באשכול:

    kubectl get pods

    רושמים בצד את השם המלא של ה-Pod‏ test-ray-job- שמופיע בפלט. השם הזה ייחודי לעבודה שלכם.

  4. בודקים את התקדמות האימון. מחליפים את gemma-training-ray-job-UNIQUE_ID בשם הייחודי של ה-Pod שרשמתם בשלב הקודם.

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

    הפלט שיוצג יהיה דומה לזה:

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

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

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

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

kubectl port-forward service/gemma3-tuning-head-svc 8265:8265 > fwd.log 2>&1 &
  1. פותחים את הקישור הבא בדפדפן: http://localhost:8265.

  2. אופציונלי: אם אתם משתמשים ב-Cloud Shell, אחרי שמריצים את הפקודה בשלב הקודם, אפשר ללחוץ על הלחצן תצוגה מקדימה באינטרנט.

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

הסרת המשאבים

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

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

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

    kubectl delete -f ray_cluster.yaml
    ‫GKE מקטין את האשכול באופן אוטומטי ומשחרר את מכונות A4 שנעשה בהן שימוש על ידי Ray.

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

    gcloud container clusters delete $CLUSTER_NAME \
        --region=$REGION
  3. כדי למחוק את כל הקטגוריה של Cloud Storage עם כל התוכן שבה, מריצים את הפקודה הבאה:

    gcloud storage rm -r gs://$GCS_BUCKET

מחיקת פרויקט

כדי למחוק פרויקט Google Cloud :

gcloud projects delete PROJECT_ID

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