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

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

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

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

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

מטרות

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

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

  • ליצור ולפרוס אשכול Slurm מסוג A4 High-GPU ברמת ייצור.

  • הגדרת סביבה עם כמה צמתים לאימון מבוזר באמצעות FSDP.

  • כוונון עדין של מודל Llama 4 באמצעות Hugging Face trl.SFTTrainer.

  • העברת נתונים לכונני SSD מקומיים.

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

  • מנקים.

עלויות

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

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

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

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

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

  2. הגדירו שה-CLI של gcloud ישתמש בזהות המאוחדת שלכם.

    איך נכנסים ל-CLI של gcloud באמצעות הזהות המאוחדת?

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

    gcloud init

  4. יוצרים או בוחרים פרויקט ב- Google Cloud .

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

    • בחירת פרויקט: כדי לבחור פרויקט לא נדרש תפקיד IAM ספציפי – אתם יכולים לבחור כל פרויקט שקיבלתם בו תפקיד.
    • יצירת פרויקט: כדי ליצור פרויקט, צריך את התפקיד "יצירת פרויקטים" (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 .

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

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

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

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

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

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

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

  8. מפעילים את חשבון השירות שמוגדר כברירת מחדל בפרויקט ב- Google Cloud : 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@ 
    
    --project=PROJECT_ID

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

  9. מקצים את התפקיד "עריכה" (roles/editor) לחשבון השירות שמוגדר כברירת מחדל: 
    gcloud projects add-iam-policy-binding PROJECT_ID 
    
    --member="serviceAccount:PROJECT_NUMBER-compute@" 
    
    --role=roles/editor
  10. יצירת פרטי כניסה לאימות מקומי בחשבון המשתמש: 
    gcloud auth application-default login
  11. מפעילים את OS Login בפרויקט: 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  12. נכנסים לחשבון Hugging Face או יוצרים חשבון.
  13. מתקינים את יחסי התלות שנדרשים כדי להשתמש ב-Cluster Toolkit.

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

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

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

  2. יוצרים טוקן גישה ל-Hugging Face read.

    לוחצים על Your Profile (הפרופיל שלך) > Settings (הגדרות) > Access tokens (טוקני גישה) > +Create new token (יצירת טוקן חדש).

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

הכנת הסביבה

כדי להכין את הסביבה:

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

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

  2. יוצרים קטגוריה של 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:

  1. עוברים לספרייה cluster-toolkit:

    cd cluster-toolkit

  2. אם זו הפעם הראשונה שאתם משתמשים ב-Cluster Toolkit, אתם צריכים ליצור את הקובץ הבינארי gcluster:

    make

  3. עוברים לספרייה examples/machine-learning/a4-highgpu-8g:

    cd examples/machine-learning/a4-highgpu-8g/

  4. פותחים את הקובץ 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

  5. פורסים את האשכול:

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

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

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

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

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

  3. מתחברים לאשכול Slurm.

  4. התקנת Frameworks וכלים.

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

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

  1. כדי להגדיר את הסביבה הווירטואלית של 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, מתקין גרסת build אוטומטית של PyTorch ומוריד את המודל Llama 4.

  2. כדי לציין את יחסי התלות של 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

  3. מציינים את 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) ואת פורמט הנתונים.

  4. כדי לציין את המשימות שהמשימות יריצו באשכול 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 את הסקריפטים שיצרתם בקטע הקודם, פועלים לפי השלבים הבאים:

  1. כדי לזהות את צומת ההתחברות, מציגים את כל ה-A4 VM בפרויקט:

    gcloud compute instances list --filter="machineType:a4-highgpu-8g"

    השם של צומת ההתחברות דומה ל-a4-high-login-001.

  2. מעלים את הסקריפטים לספריית הבית של צומת ההתחברות:

    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 וכלים

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

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

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

  2. מריצים את סקריפט ההתקנה:

    # On the login node
chmod +x install_environment.sh
./install_environment.sh

    הפקודה הזו מגדירה סביבה וירטואלית עם כל יחסי התלות הנדרשים, ומורידה את משקלי המודל לקובץ ~/Llama-4-Scout-17B-16E-Instruct.

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

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

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

  1. שולחים את המשימה למתזמן של Slurm:

    sbatch submit.slurm

  2. בצומת ההתחברות באשכול Slurm, אפשר לעקוב אחרי התקדמות המשימה באמצעות בדיקה של קובצי הפלט שנוצרו בספרייה home:

    # On the login node
tail -f llama4-*.out

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

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

    • העתקת מודל הבסיס הגדול ל-SSD המקומי של כל צומת מחשוב.
    • משימת האימון, שמתחילה אחרי שהעתקת המודל מסתיימת. המשימה הזו נמשכת כ-35 דקות.

הסרת המשאבים

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

מחיקת הפרויקט

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

gcloud projects delete PROJECT_ID

מחיקת אשכול Slurm

כדי למחוק את אשכול Slurm, מבצעים את השלבים הבאים:

  1. עוברים לספרייה cluster-toolkit.

  2. משתמשים בפקודה gcluster destroy כדי להשמיד את קובץ Terraform ואת כל המשאבים שנוצרו:

    ./gcluster destroy a4-high --auto-approve

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