Google uses AI technology to translate content into your preferred language. AI translations can contain errors.

שימוש ב-FSDP כדי לבצע כוונון עדין של Llama 4 באשכול A4 Slurm

במדריך הזה נסביר איך לבצע כוונון עדין של Llama-4-Scout-17מודל שפה גדול (LLM) באשכול Slurm מרובה צמתים ומרובה GPU ב- Google Cloud. האשכול משתמש בשתי מכונות וירטואליות (VM) מסוג A4, שלכל אחת מהן יש 8 יחידות GPU מסוג NVIDIA B200.

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

פריסת אשכול Slurm ברמת ייצור עם ביצועים גבוהים באמצעות Cluster Toolkit. במסגרת הפריסה הזו, תיצרו תמונת מכונה וירטואלית (VM) מותאמת אישית עם התוכנה הנדרשת שכבר מותקנת. בנוסף, מגדירים רשת RDMA מהירה ויוצרים קטגוריה של Cloud Storage לאירוח מצב Terraform.
אחרי שהאשכול נפרס, מריצים עבודת כוונון עדין מבוזרת באמצעות קבוצת הסקריפטים שמצורפת למדריך הזה. העבודה מתבססת על PyTorch Fully Sharded Data Parallel ‏ (FSDP), שאפשר לגשת אליו דרך ספריית Transformer Reinforcement Learning של Hugging Face.

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

מטרות

אפשר לגשת אל Llama 4 באמצעות Hugging Face.
הכנת הסביבה.
ליצור ולפרוס אשכול Slurm ברמת ייצור עם GPU A4.
הגדרה של סביבה מרובת צמתים לאימון מבוזר באמצעות FSDP.
אפשר לבצע כוונון עדין של מודל Llama 4 באמצעות Hugging Face trl.SFTTrainer.
העברת נתונים זמנית ל-SSD מקומי.
עוקבים אחרי העבודה.
לפנות.

עלויות

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

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

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

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

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

התקינו את ה-CLI של Google Cloud.

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

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

gcloud init

יוצרים או בוחרים 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 .

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

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

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

כדי להפעיל ממשקי API, צריך את תפקיד ה-IAM 'אדמין של Service Usage' (roles/serviceusage.serviceUsageAdmin), שכולל את ההרשאה serviceusage.services.enable. איך מקצים תפקידים

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

התקינו את ה-CLI של Google Cloud.

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

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

gcloud init

יוצרים או בוחרים 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 .

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

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

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

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

מעניקים תפקידים לחשבון המשתמש. מריצים את הפקודה הבאה לכל אחד מהתפקידים הבאים ב-IAM: roles/compute.admin, roles/iam.serviceAccountUser, roles/file.editor, roles/storage.admin, 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 שאתם מקצים לחשבון המשתמש.
מפעילים את חשבון השירות שמוגדר כברירת מחדל עבור Google Cloud הפרויקט:
```
gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com 

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

מקצים את התפקיד 'עריכה' (roles/editor) לחשבון השירות שמוגדר כברירת מחדל:

gcloud projects add-iam-policy-binding PROJECT_ID 

    --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" 

    --role=roles/editor

יצירת פרטי כניסה לאימות מקומי בחשבון המשתמש:
```
gcloud auth application-default login
```
הערה: אם אתם משתמשים במעטפת מקומית ובספק זהויות חיצוני (IdP), ואחרי הפעלת הפקודה הקודמת מופיעה שגיאת אימות, אתם צריכים להיכנס ל-CLI של gcloud באמצעות המאגר המאוחד לניהול זהויות.

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

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

נכנסים לחשבון Hugging Face או יוצרים חשבון.
התקנת יחסי התלות שנדרשים לשימוש ב-Cluster Toolkit

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

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

חתימה על הסכם ההסכמה לשימוש ב-Llama 4
יוצרים אסימון גישה ל-Hugging Face read.

לוחצים על הפרופיל שלך > הגדרות > טוקנים של גישה > +יצירת טוקן חדש
מעתיקים ושומרים את ערך הטוקן read access. תשתמשו בו בהמשך המדריך הזה.

הכנת הסביבה

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

משכפלים את מאגר Cluster Toolkit ב-GitHub:

git clone https://github.com/GoogleCloudPlatform/cluster-toolkit.git

יוצרים קטגוריה של Cloud Storage לאירוח מצב Terraform:
```
gcloud storage buckets create gs://BUCKET_NAME \
    --project=PROJECT_ID
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫BUCKET_NAME: שם לקטגוריה שלכם ב-Cloud Storage שעומד בקריטריונים לשמות של קטגוריות.
- ‫PROJECT_ID: המזהה שלGoogle Cloud הפרויקט שבו רוצים ליצור את קטגוריה של Cloud Storage.

יצירת אשכול Slurm בגודל A4

כדי ליצור אשכול Slurm בגודל A4:

עוברים לספרייה cluster-toolkit:
```
cd cluster-toolkit
```
אם זו הפעם הראשונה שאתם משתמשים ב-Cluster Toolkit, אתם צריכים ליצור את קובץ ה-gcluster הבינארי:
```
make
```
עוברים לספרייה examples/machine-learning/a4-highgpu-8g:
```
cd examples/machine-learning/a4-highgpu-8g/
```
פותחים את הקובץ a4high-slurm-deployment.yaml ועורכים אותו באופן הבא:
```
terraform_backend_defaults:
  type: gcs
  configuration:
    bucket: BUCKET_NAME

vars:
  deployment_name: a4-high
  project_id: PROJECT_ID
  region: REGION
  zone: ZONE
  a4h_cluster_size: 2
  a4h_reservation_name: RESERVATION_URL
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫BUCKET_NAME: השם של הקטגוריה של Cloud Storage שיצרתם בקטע הקודם.
- ‫PROJECT_ID: המזהה שלGoogle Cloud הפרויקט שבו קיים המאגר ושבו רוצים ליצור את אשכול Slurm.
- ‫REGION: האזור שבו קיימת ההזמנה.
- ‫ZONE: האזור שבו קיימת ההזמנה.
- ‫RESERVATION_URL: כתובת ה-URL של ההזמנה שבה רוצים להשתמש כדי ליצור את אשכול Slurm. בהתאם לפרויקט שבו קיימת ההזמנה, מציינים אחד מהערכים הבאים:
  - ההזמנה קיימת בפרויקט: RESERVATION_NAME
  - ההזמנה קיימת בפרויקט אחר, והפרויקט שלכם יכול להשתמש בה: projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME
פורסים את האשכול:
```
./gcluster deploy -d examples/machine-learning/a4-highgpu-8g/a4high-slurm-deployment.yaml examples/machine-learning/a4-highgpu-8g/a4high-slurm-blueprint.yaml --auto-approve
```
הפקודה ./gcluster deploy היא תהליך דו-שלבי, שמתבצע באופן הבא:
- בשלב הראשון נוצרת תמונה מותאמת אישית עם כל התוכנות שמותקנות מראש, והתהליך הזה יכול להימשך עד 35 דקות.
- בשלב השני, המערכת פורסת את האשכול באמצעות התמונה המותאמת אישית. התהליך הזה אמור להסתיים מהר יותר מהשלב הראשון.
אם השלב הראשון מצליח אבל השלב השני נכשל, אפשר לנסות לפרוס מחדש את אשכול Slurm על ידי דילוג על השלב הראשון:
```
./gcluster deploy -d examples/machine-learning/a4-highgpu-8g/a4high-slurm-deployment.yaml examples/machine-learning/a4-highgpu-8g/a4high-slurm-blueprint.yaml --auto-approve --skip "image" -w
```

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

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

יצירת סקריפטים של עומסי עבודה.
העלאת סקריפטים לאשכול Slurm.
מתחברים לאשכול Slurm.
התקנה של frameworks וכלים

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

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

כדי להגדיר את הסביבה הווירטואלית של Python, יוצרים את הקובץ install_environment.sh עם התוכן הבא:

#!/bin/bash
# This script sets up a consistent environment for FSDP training.
# It is meant to be run once on the login node of your Slurm cluster
set -e

# --- 1. Create the Python virtual environment ---
VENV_PATH="$HOME/.venv/venv-fsdp"
if [ ! -d "$VENV_PATH" ]; then
  echo "--- Creating Python virtual environment at $VENV_PATH ---"
  python3 -m venv $VENV_PATH
else
  echo "--- Virtual environment already exists at $VENV_PATH ---"
fi

source $VENV_PATH/bin/activate

# --- 2. Install Dependencies ---
echo "--- [STEP 2.1] Upgrading build toolchain ---"
pip install --upgrade pip wheel packaging

echo "--- [STEP 2.2] Installing PyTorch Nightly ---"
pip install --force-reinstall --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128

echo "--- [STEP 2.3] Installing application dependencies ---"
if [ -f "requirements-fsdp.txt" ]; then
    pip install -r requirements-fsdp.txt
else
    echo "ERROR: requirements-fsdp.txt not found!"
    exit 1
fi

# --- 3. Download the Model ---
echo "--- [STEP 2.4] Downloading Llama4 model ---"
if [ -z "$HF_TOKEN" ]; then
  echo "ERROR: The HF_TOKEN environment variable is not set."; exit 1;
fi
pip install huggingface_hub[cli]

# Execute the CLI using its full, explicit path
$VENV_PATH/bin/huggingface-cli download meta-llama/Llama-4-Scout-17B-16E-Instruct --local-dir ~/Llama-4-Scout-17B-16E-Instruct --token $HF_TOKEN

echo "--- Environment setup complete. ---"

הסקריפט הזה מגדיר סביבה וירטואלית אמינה של Python, מתקין גרסת nightly build של PyTorch ומוריד את מודל Llama 4.

כדי לציין את יחסי התלות של Python בסקריפט ההדרכה, יוצרים קובץ requirements-fsdp.txt עם התוכן הבא:
```
transformers==4.55.0
datasets==4.0.0
peft==0.16.0
accelerate==1.9.0
trl==0.21.0

# Other dependencies
sentencepiece==0.2.0
```

מציינים את llama4-train-distributed.py כסקריפט האימון הראשי:

import torch
from datasets import load_dataset
from peft import LoraConfig, PeftModel
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    HfArgumentParser,
)

from torch.distributed import get_rank, get_world_size

from transformers.models.llama4.modeling_llama4 import Llama4TextDecoderLayer
from trl import SFTTrainer
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class ScriptArguments:
    model_id: str = field(metadata={"help": "Hugging Face model ID from the Hub"})
    dataset_name: str = field(default="philschmid/gretel-synthetic-text-to-sql", metadata={"help": "Dataset from the Hub"})
    run_inference_after_training: bool = field(default=False, metadata={"help": "Run sample inference on rank 0 after training"})
    dataset_subset_size: Optional[int] = field(default=None, metadata={"help": "Number of samples to use from the dataset for training. If None, uses the full dataset."})

@dataclass
class PeftArguments:
    lora_r: int = field(default=16, metadata={"help": "LoRA attention dimension"})
    lora_alpha: int = field(default=32, metadata={"help": "LoRA alpha scaling factor"})
    lora_dropout: float = field(default=0.05, metadata={"help": "LoRA dropout probability"})

@dataclass
class SftTrainingArguments(TrainingArguments):
    max_length: Optional[int] = field(default=2048, metadata={"help": "The maximum sequence length for SFTTrainer"})
    packing: Optional[bool] = field(default=False, metadata={"help": "Enable packing for SFTTrainer"})
    ddp_find_unused_parameters: Optional[bool] = field(default=True, metadata={"help": "When using FSDP activation checkpointing, this must be set to True"})

def formatting_prompts_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 = f"### SCHEMA:\n{example['sql_context']}\n\n### USER QUERY:\n{example['sql_prompt']}"
    response = f"\n\n### SQL QUERY:\n{example['sql']}"
    return f"{system_message}\n\n{user_prompt}{response}"

def main():
    parser = HfArgumentParser((ScriptArguments, PeftArguments, SftTrainingArguments))
    script_args, peft_args, training_args = parser.parse_args_into_dataclasses()

    training_args.gradient_checkpointing = True
    training_args.gradient_checkpointing_kwargs = {"use_reentrant": False}

    training_args.optim = "adamw_torch_fused"

    training_args.fsdp = "full_shard"
    training_args.fsdp_config = {
        "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
        "fsdp_transformer_layer_cls_to_wrap": [Llama4TextDecoderLayer],
        "fsdp_state_dict_type": "FULL_STATE_DICT",
        "fsdp_offload_params": False,
        "fsdp_forward_prefetch": True,
    }

    tokenizer = AutoTokenizer.from_pretrained(script_args.model_id, trust_remote_code=True)

    model = AutoModelForCausalLM.from_pretrained(
        script_args.model_id,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        attn_implementation="sdpa",
    )

    peft_config = LoraConfig(
        r=peft_args.lora_r,
        lora_alpha=peft_args.lora_alpha,
        lora_dropout=peft_args.lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    )
    rank = get_rank()
    world_size = get_world_size()

    dataset = load_dataset(script_args.dataset_name, split="train")

    if script_args.dataset_subset_size is not None:
        dataset = dataset.select(range(script_args.dataset_subset_size))
    else:
        print(f"Using the full dataset with {len(dataset)} samples.")

    dataset = dataset.shuffle(seed=training_args.seed)
    print(f"Dataset shuffled with seed: {training_args.seed}.")

    if world_size > 1:
        print(f"Sharding dataset for Rank {rank} of {world_size}.")
        dataset = dataset.shard(num_shards=world_size, index=rank)

    print("Initializing SFTTrainer...")
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        peft_config=peft_config,
        formatting_func=formatting_prompts_func,
        processing_class=tokenizer,
    )

    trainer.train()

    trainer.save_model(training_args.output_dir)

    if script_args.run_inference_after_training and trainer.is_world_process_zero():
        del model
        del trainer
        torch.cuda.empty_cache()
        run_post_training_inference(script_args, training_args, tokenizer)

def run_post_training_inference(script_args, training_args, tokenizer):
    """
    Loads the fine-tuned PEFT adapter from the local output directory and runs inference.
    This should only be called on rank 0 after training is complete.
    """
    print("\n" + "="*50)
    print("=== RUNNING POST-TRAINING INFERENCE TEST ===")
    print("="*50 + "\n")

    # Load the base model and merge the adapter.
    base_model = AutoModelForCausalLM.from_pretrained(
        script_args.model_id,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map="auto"
    )
    # Load the PEFT adapter and merge it into the base model
    model = PeftModel.from_pretrained(base_model, training_args.output_dir)
    model = model.merge_and_unload() # Merge weights for faster inference
    model.eval()

    # Define the test case
    schema = "CREATE TABLE artists (Name TEXT, Country TEXT, Genre TEXT)"
    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."
    question = "Show me all artists from the Country just north of the USA."

    # This must match the formatting_func exactly
    prompt = f"{system_message}\n\n### SCHEMA:\n{schema}\n\n### USER QUERY:\n{question}\n\n### SQL QUERY:\n"

    print(f"Test Prompt:\n{prompt}")

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    print("\n--- Generating SQL... ---")
    outputs = model.generate(
        **inputs,
        max_new_tokens=100,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=False,
        temperature=None,
        top_p=None,
    )

    generated_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)[len(prompt):].strip()

    print(f"\n--- Generated SQL Query ---")
    print(generated_sql)
    print("\n" + "="*50)
    print("=== INFERENCE TEST COMPLETE ===")
    print("="*50 + "\n")

if __name__ == "__main__":
    main()

הסקריפט הזה משתמש ב-TRL Supervised Fine-Tuning (SFT) Trainer כדי לנהל לולאות אימון של FSDP, הגדרת Low-Rank Adaptation (‏LoRA) ועיצוב נתונים.

כדי לציין את המשימות שהעבודות יפעילו באשכול Slurm, יוצרים את הקובץ submit.slurm עם התוכן הבא:

#!/bin/bash
#SBATCH --job-name=llama4-fsdp-fixed
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8
#SBATCH --partition=a4high
#SBATCH --output=llama4-%j.out
#SBATCH --error=llama4-%j.err

set -e
set -x

echo "--- Slurm Job Started ---"
echo "Job ID: $SLURM_JOB_ID"
echo "Node List: $SLURM_JOB_NODELIST"

# --- Define Paths ---
LOCAL_SSD_PATH="/mnt/localssd/job_${SLURM_JOB_ID}"
VENV_PATH="${HOME}/.venv/venv-fsdp"
MODEL_PATH="${HOME}/Llama-4-Scout-17B-16E-Instruct"

# --- STAGE 1: Stage Data to Local SSD on Each Node ---
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 bash -c "
  echo '--- Staging on node: $(hostname) ---'

  mkdir -p ${LOCAL_SSD_PATH}

  echo 'Copying virtual environment...'
  rsync -a -q ${VENV_PATH}/ ${LOCAL_SSD_PATH}/venv/

  echo 'Copying model weights...'
  rsync -a --info=progress2 ${MODEL_PATH}/ ${LOCAL_SSD_PATH}/model/

  mkdir -p ${LOCAL_SSD_PATH}/hf_cache

  echo '--- Staging on $(hostname) complete ---'
"
echo "--- Staging complete on all nodes ---"

# --- STAGE 2: Run the Training Job ---
echo "--- Launching Distributed Training with GIB NCCL Plugin ---"
nodes=( $( scontrol show hostnames "$SLURM_JOB_NODELIST" ) )
head_node=${nodes[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)

export MASTER_ADDR=$head_node_ip
export MASTER_PORT=29500

export NCCL_SOCKET_IFNAME=enp0s19
export NCCL_NET=gIB

# export NCCL_DEBUG=INFO # Un-comment to diagnose NCCL issues if needed

srun --cpu-bind=none --accel-bind=g bash -c '
  # Activate the environment from the local copy
  source '${LOCAL_SSD_PATH}'/venv/bin/activate

  # Point Hugging Face cache to the local SSD
  export HF_HOME='${LOCAL_SSD_PATH}'/hf_cache

  export RANK=$SLURM_PROCID
  export WORLD_SIZE=$SLURM_NTASKS
  export LOCAL_RANK=$SLURM_LOCALID

  export LD_LIBRARY_PATH=/usr/local/gib/lib64:$LD_LIBRARY_PATH
  source /usr/local/gib/scripts/set_nccl_env.sh

  # --- Launch the training ---
  python \
    '${SLURM_SUBMIT_DIR}'/llama4-train-distributed.py \
      --model_id="'${LOCAL_SSD_PATH}'/model/" \
      --output_dir="'${LOCAL_SSD_PATH}'/outputs/" \
      --dataset_name="philschmid/gretel-synthetic-text-to-sql" \
      --seed=900913 \
      --bf16=True \
      --num_train_epochs=1 \
      --per_device_train_batch_size=2 \
      --gradient_accumulation_steps=4 \
      --learning_rate=2e-5 \
      --logging_steps=10 \
      --lora_r=16 \
      --lora_alpha=32 \
      --lora_dropout=0.05 \
      --run_inference_after_training
'

# --- STAGE 3: Copy Final Results Back to Persistent Storage ---
echo "--- Copying final results from local SSD to shared storage ---"
PERSISTENT_OUTPUT_DIR="${HOME}/outputs/llama4_job_${SLURM_JOB_ID}"
mkdir -p "$PERSISTENT_OUTPUT_DIR"

# Only copy from the head node where trl has combined the results
srun --nodes=1 --ntasks=1 -w "$head_node" \
  rsync -a --info=progress2 "${LOCAL_SSD_PATH}/outputs/" "${PERSISTENT_OUTPUT_DIR}/"

# --- STAGE 4: Cleanup ---
echo "--- Cleaning up local SSD on all nodes ---"
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 bash -c "rm -rf ${LOCAL_SSD_PATH}"

echo "--- Slurm Job Finished ---"

העלאת סקריפטים לאשכול Slurm

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

כדי לזהות את צומת הכניסה, מציגים ברשימה את כל המכונות הווירטואליות מסוג A4 בפרויקט:
```
gcloud compute instances list --filter="machineType:a4-highgpu-8g"
```
השם של צומת הכניסה דומה ל-a4-high-login-001.

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

gcloud compute scp --project="$PROJECT_ID" --zone="$ZONE" --tunnel-through-iap \
  ./install_environment.sh \
  ./requirements-fsdp.txt \
  ./llama4-train-distributed.py \
  ./submit.slurm \
  "${LOGIN_NODE_NAME}":~/

מחליפים את LOGIN_NODE_NAME בשם של צומת הכניסה.

התחברות לאשכול Slurm

מתחברים לאשכול Slurm על ידי התחברות לצומת הכניסה באמצעות SSH:

gcloud compute ssh LOGIN_NODE_NAME \
    --project=PROJECT_ID \
    --tunnel-through-iap \
    --zone=ZONE

התקנת frameworks וכלים

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

מייצאים את הטוקן של Hugging Face:

# On the login node
export HF_TOKEN="hf_..." # Replace with your token

מריצים את סקריפט ההתקנה:
```
# On the login node
chmod +x install_environment.sh
./install_environment.sh
```
הפקודה הזו מגדירה סביבה וירטואלית עם כל הרכיבים התלויים הנדרשים, ומורידה את משקלי המודל לקובץ ~/Llama-4-Scout-17B-16E-Instruct.

ההורדה של המודל גדולה מאוד (כ-200GB), ולכן התהליך הזה נמשך כ-30 דקות, בהתאם לתנאי הרשת.

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

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

שולחים את העבודה למתזמן של Slurm:
```
sbatch submit.slurm
```
בצומת הכניסה באשכול Slurm, אפשר לעקוב אחרי התקדמות העבודה על ידי בדיקת קובצי הפלט שנוצרו בספרייה home:
```
# On the login node
tail -f llama4-*.out
```
אם העבודה מתחילה בהצלחה, בקובץ .err מוצג סרגל התקדמות שמתעדכן ככל שהעבודה מתקדמת.

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

הסרת המשאבים

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

מחיקת פרויקט

זהירות: כשמוחקים פרויקט:

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

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

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

gcloud projects delete PROJECT_ID

מחיקת אשכול Slurm

כדי למחוק את אשכול Slurm:

עוברים לספרייה cluster-toolkit.
מכבים באופן סופי את קובץ Terraform ואת כל המשאבים שנוצרו:
```
./gcluster destroy a4-high --auto-approve
```