שימוש ב-FSDP כדי לבצע כוונון עדין של Mixtral-8x7B באשכול A4 Slurm

במדריך הזה מוסבר איך לכוונן מודל mistralai/Mixtral-8x7B-v0.1 באשכול Slurm עם כמה צמתים וכמה GPU‏ Google Cloud. האשכול משתמש בשתי מכונות וירטואליות (VM) של a4-highgpu-8g, שלכל אחת מהן יש 8 יחידות GPU מסוג NVIDIA B200.

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

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

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

מטרות

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

עלויות

במסמך הזה משתמשים ברכיבים הבאים של 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, צריך את תפקיד ה-IAM 'אדמין של Service Usage' (roles/serviceusage.serviceUsageAdmin), שכולל את ההרשאה serviceusage.services.enable. איך מקצים תפקידים

    gcloud services enable compute.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com lustre.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, צריך את תפקיד ה-IAM 'אדמין של Service Usage' (roles/serviceusage.serviceUsageAdmin), שכולל את ההרשאה serviceusage.services.enable. איך מקצים תפקידים

    gcloud services enable compute.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com lustre.googleapis.com
  14. מעניקים תפקידים לחשבון המשתמש. מריצים את הפקודה הבאה לכל אחד מהתפקידים הבאים ב-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 שאתם מקצים לחשבון המשתמש.
  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. מפעילים את OS Login בפרויקט:
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  19. נכנסים לחשבון Hugging Face או יוצרים חשבון.
  20. מתקינים את יחסי התלות שנדרשים לשימוש ב-Cluster Toolkit.

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

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

  1. יוצרים טוקן של Hugging Face read access.
  2. מעתיקים ושומרים את ערך טוקן הגישה read. תצטרכו להשתמש בו בהמשך המדריך הזה.

הכנת הסביבה

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

  1. משכפלים את Google Cloud מאגר Cluster Toolkit:

    git clone https://github.com/GoogleCloudPlatform/cluster-toolkit.git
    
  2. יוצרים קטגוריה של Cloud Storage:

    export BUCKET_NAME="your-unique-bucket-name"
    gcloud storage buckets create gs://${BUCKET_NAME}
    

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

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

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

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

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

    פותחים את הקובץ a4high-slurm-deployment.yaml ועורכים אותו באופן הבא:

    terraform_backend_defaults:
      type: gcs
      configuration:
        bucket: BUCKET_NAME
    
    vars:
      deployment_name: DEPLOYMENT_NAME
      project_id: PROJECT_ID
      region: REGION
      zone: ZONE
      a4h_cluster_size: 2
      a4h_reservation_name: RESERVATION_NAME
    

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

    • BUCKET_NAME: : השם של הקטגוריה ב-Cloud Storage שיצרתם בקטע הקודם.
    • PROJECT_ID: המזהה של Google Cloud הפרויקט שבו קיים Cloud Storage ושבו רוצים ליצור את אשכול Slurm.
    • REGION: האזור שבו קיימת ההזמנה.
    • ZONE: האזור שבו קיימת ההזמנה.
    • A4h_reservation_name: משתמשים בשם של הזמנת A4.
  4. פותחים את הקובץ a4high-slurm-blueprint.yaml ועורכים אותו באופן הבא:

    • מסירים את המודול filestore_homefs.
    • מפעילים את המודולים lustrefs ו-private-service-access.
    • בבלוק vars, מגדירים את ההגדרות הבאות:
      1. Find slurm_vars ומגדירים את install_managed_lustre לערך true.
      2. מגדירים את הפרמטר per_unit_storage_throughput לערך 500.
      3. מגדירים את הפרמטר size_gib לערך 36000.
  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 דקות.
    • בשלב השני, המערכת פורסת את האשכול באמצעות התמונה המותאמת אישית. התהליך הזה אמור להסתיים מהר יותר מהשלב הראשון.

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

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

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

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

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

  4. התקנה של frameworks וכלים

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

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

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

    #!/bin/bash
    # This script sets a reliable environment for FSDP training.
    # It is meant to be run on a compute node.
    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
    
    # --- [STEP 2.4] Build Flash Attention from Source ---
    echo "--- Building flash-attn from source... This will take a while. ---"
    # Use all available CPU cores to speed up the build
    MAX_JOBS=$(nproc) pip install flash-attn --no-build-isolation
    
    # --- 3. Download the Model ---
    echo "--- [STEP 2.5] Downloading Mixtral 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 mistralai/Mixtral-8x7B-v0.1 --local-dir ~/Mixtral-8x7B-v0.1 --token $HF_TOKEN
    
    echo "--- Environment setup complete. ---"
    
  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
    protobuf==6.31.1
    
  3. מציינים את train-mixtral.py כסקריפט האימון הראשי:

    import torch
    from torch.distributed.fsdp import MixedPrecision
    from datasets import load_dataset
    import shutil
    import os
    import torch.distributed as dist
    
    from peft import LoraConfig, PeftModel, get_peft_model
    from transformers import (
        AutoModelForCausalLM,
        AutoTokenizer,
        TrainingArguments,
        HfArgumentParser,
    )
    
    from torch.distributed import get_rank, get_world_size
    
    from transformers.models.mixtral.modeling_mixtral import MixtralDecoderLayer
    from trl import SFTTrainer
    from dataclasses import dataclass, field
    from typing import Optional
    
    @dataclass
    class ScriptArguments:
        model_id: str = field(default="mistralai/Mixtral-8x7B-v0.1", 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=False, metadata={"help": "When using FSDP activation checkpointing, this must be set to False for Mixtral"})
    
    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": True}
    
        training_args.optim = "adamw_torch_fused"
    
        bf16_policy = MixedPrecision(
            param_dtype=torch.bfloat16,
            reduce_dtype=torch.bfloat16,
            buffer_dtype=torch.bfloat16,
        )
    
        training_args.fsdp = "full_shard"
        training_args.fsdp_config = {
            "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
            "fsdp_transformer_layer_cls_to_wrap": [MixtralDecoderLayer],
            "fsdp_state_dict_type": "SHARDED_STATE_DICT",
            "fsdp_offload_params": False,
            "fsdp_forward_prefetch": True,
            "fsdp_mixed_precision_policy": bf16_policy
        }
    
        tokenizer = AutoTokenizer.from_pretrained(script_args.model_id, trust_remote_code=True)
    
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.padding_side = "right"
    
        model = AutoModelForCausalLM.from_pretrained(
            script_args.model_id,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            attn_implementation="flash_attention_2",
        )
    
        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"],
        )
    
        model = get_peft_model(model, peft_config)
    
        data_splits = load_dataset(script_args.dataset_name)
    
        dataset = data_splits["train"]
        eval_dataset = data_splits["test"]
    
        if script_args.dataset_subset_size is not None:
            dataset = dataset.select(range(script_args.dataset_subset_size))
    
        dataset = dataset.shuffle(seed=training_args.seed)
    
        trainer = SFTTrainer(
            model=model,
            args=training_args,
            train_dataset=dataset,
            eval_dataset=eval_dataset,
            formatting_func=formatting_prompts_func,
            processing_class=tokenizer,
        )
    
        trainer.train()
    
        dist.barrier()
        if trainer.is_world_process_zero():
            best_model_path = trainer.state.best_model_checkpoint
    
            final_model_dir = os.path.join(training_args.output_dir, "final_best_model")
            print(f"Copying best model to: {final_model_dir}")
    
            if os.path.exists(final_model_dir):
                shutil.rmtree(final_model_dir)
            shutil.copytree(best_model_path, final_model_dir)
    
            if script_args.run_inference_after_training:
                del model, trainer
                torch.cuda.empty_cache()
                run_post_training_inference(script_args, final_model_dir, tokenizer)
    
    def run_post_training_inference(script_args, best_model_path, tokenizer):
        print("\n" + "="*50)
        print("=== RUNNING POST-TRAINING INFERENCE TEST ===")
        print("="*50 + "\n")
    
        base_model = AutoModelForCausalLM.from_pretrained(
            script_args.model_id,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            device_map="auto"
        )
        model = PeftModel.from_pretrained(base_model, best_model_path)
        model = model.merge_and_unload()
        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."
    
        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()
    
  4. כדי לציין את המשימות שהעבודות יפעילו באשכול Slurm, יוצרים את הקובץ train-mixtral.sh עם התוכן הבא:

    #!/bin/bash
    #SBATCH --job-name=mixtral-fsdp
    #SBATCH --nodes=2
    #SBATCH --ntasks-per-node=8
    #SBATCH --gpus-per-node=8
    #SBATCH --partition=a4high
    #SBATCH --output=mixtral-%j.out
    #SBATCH --error=mixtral-%j.err
    
    set -e
    set -x
    
    echo "--- Slurm Job Started ---"
    
    # --- Define Paths ---
    LOCAL_SSD_PATH="/mnt/localssd/job_${SLURM_JOB_ID}"
    VENV_PATH="${HOME}/.venv/venv-fsdp"
    MODEL_PATH="${HOME}/Mixtral-8x7B-v0.1"
    
    # --- 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 ${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}'/train-mixtral.py \
        --model_id="'${LOCAL_SSD_PATH}'/model/" \
        --output_dir="${HOME}/outputs/mixtral_job_${SLURM_JOB_ID}" \
        --dataset_name="philschmid/gretel-synthetic-text-to-sql" \
        --seed=900913 \
        --bf16=True \
        --num_train_epochs=3 \
        --per_device_train_batch_size=32 \
        --gradient_accumulation_steps=4 \
        --learning_rate=4e-5 \
        --logging_steps=3 \
        --lora_r=32 \
        --lora_alpha=32 \
        --lora_dropout=0.05 \
        --eval_strategy=steps \
        --eval_steps=10 \
        --save_strategy=steps \
        --save_steps=10 \
        --load_best_model_at_end=False \
        --metric_for_best_model=eval_loss \
        --run_inference_after_training \
        --dataset_subset_size=67000
    '
    
    # --- STAGE 3: 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. כדי לזהות את צומת הכניסה, מציגים ברשימה את כל מכונות ה-VM בפרויקט:

    gcloud compute instances list
    

    השם של צומת הכניסה דומה ל-a4-high-login-001.

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

    # Run this from your local machine where you created the files
    LOGIN_NODE_NAME="your-login-node-name" # e.g., a4high-login-001
    PROJECT_ID="your-gcp-project-id"
    ZONE="your-cluster-zone" # e.g., us-west4-a
    
    gcloud compute scp --project="$PROJECT_ID" --zone="$ZONE" --tunnel-through-iap \
      ./install_environment.sh \
      ./requirements-fsdp.txt \
      ./train-mixtral.py \
      ./train-mixtral.sh \
      "${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
    srun \
      --job-name=env-setup \
      --nodes=1 \
      --ntasks=1 \
      --gpus-per-node=1 \
      --partition=a4high \
      bash ./install_environment.sh
    

    הפקודה הזו מגדירה את הסביבה הווירטואלית, מתקינה את כל הרכיבים התלויים ומורידה את משקלי המודל של Mixtral אל ~/Mixtral-8x7B-v0.1. התהליך הזה יכול להימשך יותר מ-30 דקות.

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

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

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

    # On the login node
    sbatch train-mixtral.sh
    
  2. בצומת הכניסה באשכול Slurm, אפשר לעקוב אחרי ההתקדמות של העבודה על ידי בדיקת קובצי הפלט שנוצרו בספרייה home:

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

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

    העבודה מחולקת לשני שלבים עיקריים:

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

    ההרצה של כל העבודה נמשכת כ-40 דקות.

הסרת המשאבים

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

מחיקת אשכול Slurm

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

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

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

    ./gcluster destroy DEPLOYMENT_NAME --auto-approve
    

מחיקת פרויקט

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

gcloud projects delete PROJECT_ID

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