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

Clustered GPUs

במדריך הזה מוסבר איך לבצע כוונון עדין של מודל 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.

חשוב: כדי להשלים את המדריך הזה, צריך לשריין קיבולת ליצירת שתי מכונות וירטואליות מסוג A4. מידע נוסף על האפשרויות להזמנת קיבולת ב-AI Hypercomputer לתאריך ולשעה עתידיים זמין במאמר בחירת אפשרות צריכה.

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

מטרות

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

עלויות

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

• Compute Engine
• Cloud Storage
• Cloud Logging
• Google Cloud Managed Lustre

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

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

gcloud services enable compute.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com lustre.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

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

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

1. מעניקים תפקידים לחשבון המשתמש. מריצים את הפקודה הבאה לכל אחד מהתפקידים הבאים ב-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: המזהה של חשבון המשתמש . לדוגמה, [e-mail].
   • ‫ROLE: תפקיד ה-IAM שאתם מקצים לחשבון המשתמש.
2. מפעילים את חשבון השירות שמוגדר כברירת מחדל עבור Google Cloud הפרויקט:

   gcloud iam service-accounts enable [e-mail] \
       --project=PROJECT_ID

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

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

   gcloud projects add-iam-policy-binding PROJECT_ID \
     --member="serviceAccount:[e-mail]" \
     --role=roles/editor

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

   gcloud auth application-default login

   הערה: אם אתם משתמשים במעטפת מקומית ובספק זהויות חיצוני (IdP), ואחרי הפעלת הפקודה הקודמת מופיעה שגיאת אימות, אתם צריכים להיכנס ל-CLI של Google Cloud באמצעות הזהות המאוחדת שלכם.
5. מפעילים את OS Login בפרויקט:

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

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

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

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

1. נכנסים לחשבון ב-Hugging Face ומתנסים במודל Mixtral.
2. יוצרים טוקן גישה ל-Hugging Face read.
3. מעתיקים את ערך הטוקן ושומרים אותו. תצטרכו להשתמש בו בהמשך המדריך הזה.

התקנת Cluster Toolkit

‫Cluster Toolkit הוא כלי בקוד פתוח שמפשט את הפריסה של עומסי עבודה של מחשוב עתיר ביצועים (HPC), בינה מלאכותית (AI) ולמידת מכונה (ML) ב- Google Cloud. למידע נוסף על השימוש ב-gcluster ועל ניהול אשכולות, אפשר לעיין במאמר סקירה כללית של Cluster Toolkit.

1. הכנת הגרסה של Cluster Toolkit:

   export CLUSTER_TOOLKIT_TAG=v1.97.0
   
   # Detect OS (linux or mac)
   case "$(uname -s)" in
     Linux*)     OS="linux" ;;
     Darwin*)    OS="mac" ;;
     *)          echo "Error: Unsupported operating system: $(uname -s)" >&2; exit 1 ;;
   esac
   
   # Detect Architecture (amd64 or arm64)
   case "$(uname -m)" in
     x86_64)     ARCH="amd64" ;;
     aarch64|arm64) ARCH="arm64" ;;
     *)          echo "Error: Unsupported architecture: $(uname -m)" >&2; exit 1 ;;
   esac

2. מורידים את הגרסה:

   # Download and extract the platform-specific bundle
   curl -LO "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_TAG}/gcluster_bundle_${OS}_${ARCH}.zip"
   unzip "gcluster_bundle_${OS}_${ARCH}.zip" -d cluster-toolkit/
   rm -f "gcluster_bundle_${OS}_${ARCH}.zip"

3. מגדירים את הנתיב gcluster:

   export CLUSTER_TOOLKIT_PATH="$(pwd)/cluster-toolkit"
   export PATH="${CLUSTER_TOOLKIT_PATH}:${PATH}"
   gcluster --version

הכנת הסביבה

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

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

   export PROJECT_ID="YOUR_PROJECT_ID"
   export ZONE="YOUR_ZONE"
   export REGION="YOUR_REGION"
   export RESERVATION_NAME="YOUR_RESERVATION_NAME"
   export DEPLOYMENT_NAME="YOUR_CLUSTER_NAME"
   export BUCKET_NAME="YOUR_GCS_BUCKET"
   export HF_TOKEN="YOUR_HF_TOKEN"

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

   • ‫YOUR_PROJECT_ID: השם של Google Cloud הפרויקט שבו רוצים ליצור את אשכול Slurm.
   • ‫YOUR_ZONE: האזור שבו קיימת ההזמנה.
   • ‫YOUR_REGION: האזור שבו קיימת ההזמנה.
   • ‫YOUR_RESERVATION_NAME: כתובת ה-URL או השם של ההזמנה שרוצים להשתמש בה כדי ליצור את אשכול Slurm.
   • ‫YOUR_CLUSTER_NAME: השם של הפריסה. משתמשים בשם קצר שמכיל רק אותיות ומספרים (לדוגמה, a4high). השם הזה מוקצה גם לאשכול Slurm שנוצר על ידי הפריסה.
   • ‫YOUR_GCS_BUCKET: השם של הקטגוריה שבה מאחסנים את התוצאות מנקודת הבדיקה של האימון. מציינים דלי קיים או יוצרים דלי חדש. לפני שיוצרים קטגוריה, חשוב לקרוא את הדרישות לבחירת שמות לקטגוריות.
   • ‫YOUR_HF_TOKEN: האסימון של Hugging Face שיצרתם בשלב קודם.

2. יוצרים קטגוריה של Cloud Storage:

   gcloud storage buckets create "gs://${BUCKET_NAME}" \
     --project="${PROJECT_ID}"

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

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

1. כדי להחליף את הקובץ a4high-slurm-deployment.yaml בספרייה examples/machine-learning/a4-highgpu-8g במשתני הסביבה שלכם, מריצים את הפקודה הבאה:

   MANIFEST_PATH="${CLUSTER_TOOLKIT_PATH}/examples/machine-learning/a4-highgpu-8g"
   cat <<EOF > "${MANIFEST_PATH}/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}
   EOF

2. פותחים את הקובץ a4high-slurm-blueprint.yaml בספרייה examples/machine-learning/a4-highgpu-8g ועורכים אותו כדי להשתמש ב-Managed Lustre בספרייה המשותפת /home באופן הבא:

   • מסירים את בלוק המודול homefs שמוגדר כברירת מחדל באמצעות source: modules/file-system/filestore.
   • מפעילים את המודולים lustrefs (חסימת homefs עם remote_mount: lustrefs) ו-private-service-access.
   • בבלוק vars, מגדירים את ההגדרות הבאות:
     1. משנים את הערך של install_managed_lustre ב-/var/tmp/slurm_vars.json ל-true.
     2. מגדירים את הפרמטר per_unit_storage_throughput לערך 500.
     3. מגדירים את הפרמטר lustre_size_gib לערך 36000.
     4. ביטול ההערה של lustre_instance_id: lustre-instance.
     5. מוסיפים הערה או מסירים את filestore_ip_range שלא בשימוש.

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

   ./gcluster deploy \
     -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
     "${MANIFEST_PATH}/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 --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. ---"
   # Bypass PyTorch CUDA version mismatch check for extension builds
   python3 -c "
   import torch.utils.cpp_extension as ce
   p = ce.__file__
   content = open(p).read().replace('raise RuntimeError(CUDA_MISMATCH_MESSAGE', 'pass # raise RuntimeError')
   open(p, 'w').write(content)
   "
   
   # 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/hf download mistralai/Mixtral-8x7B-v0.1 --local-dir ~/Mixtral-8x7B-v0.1 --token $HF_TOKEN
   
   echo "--- Environment setup complete. ---"

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

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. מגדירים את המשתנה LOGIN_NODE על ידי אחזור השם של צומת הכניסה לאשכול:

   LOGIN_NODE="$(gcloud compute instances list \
                   --project="${PROJECT_ID}" \
                   --filter="labels.ghpc_deployment='${DEPLOYMENT_NAME}' AND labels.slurm_instance_role='login'" \
                   --format="value(name)" | head -n 1)"

   המשתנה LOGIN_NODE מאחסן ערך דומה ל-${DEPLOYMENT_NAME}-login-001.

2. יוצרים כלל לחומת האש:

   gcloud compute firewall-rules create allow-ssh-ingress-from-iap \
     --project="${PROJECT_ID}" \
     --network="${CLUSTER_NETWORK}" \
     --direction=INGRESS \
     --action=allow \
     --rules=tcp:22 \
     --source-ranges=35.235.240.0/20 \
     --description="Allow SSH ingress from Google Cloud Identity-Aware Proxy (IAP)"

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

   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}":~/

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

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

gcloud compute ssh "${LOGIN_NODE}" \
    --project="${PROJECT_ID}" \
    --tunnel-through-iap \
    --zone="${ZONE}"
    -- -t "export HF_TOKEN='${HF_TOKEN}'; bash -l"

התקנת frameworks וכלים

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

# 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 דקות.

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

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

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

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

   # On the login node
   sbatch train-mixtral.sh

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

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

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

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

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

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

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

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

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

   https://console.cloud.google.com/monitoring/metrics-explorer?project=YOUR_PROJECT_ID&pageState=%7B%22xyChart%22%3A%7B%22dataSets%22%3A%5B%7B%22timeSeriesFilter%22%3A%7B%22filter%22%3A%22metric.type%3D%5C%22agent.googleapis.com%2Fgpu%2Futilization%5C%22%20resource.type%3D%5C%22gce_instance%5C%22%22%2C%22perSeriesAligner%22%3A%22ALIGN_MEAN%22%7D%2C%22plotType%22%3A%22LINE%22%7D%5D%7D%7D
   

   אפשר גם להזין את הפקודה ישירות בטרמינל:

   open "https://console.cloud.google.com/monitoring/metrics-explorer?project=${PROJECT_ID}&pageState=%7B%22xyChart%22%3A%7B%22dataSets%22%3A%5B%7B%22timeSeriesFilter%22%3A%7B%22filter%22%3A%22metric.type%3D%5C%22agent.googleapis.com%2Fgpu%2Futilization%5C%22%20resource.type%3D%5C%22gce_instance%5C%22%22%2C%22perSeriesAligner%22%3A%22ALIGN_MEAN%22%7D%2C%22plotType%22%3A%22LINE%22%7D%5D%7D%7D"

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

   • השימוש ב-GPU: כדי שהאימון יהיה תקין, אפשר לצפות שהשימוש בכל 16 יחידות ה-GPU (שמונה יחידות GPU לכל מכונה וירטואלית באשכול) יעלה ויתייצב ברמה מסוימת לאורך האימון.

   • משך העבודה: העבודה צפויה להימשך כשעה.

הסרת המשאבים

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

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

מריצים את הפקודות הבאות כדי למחוק את המשאבים שיצרתם במדריך הזה:

1. כדי למחוק את אשכול Slurm, עוברים לספרייה cluster-toolkit ומריצים את הפקודה הבאה:

   ./gcluster destroy "${DEPLOYMENT_NAME}" --auto-approve --robust

2. מוחקים את הקטגוריה:

   gcloud storage rm --recursive "gs://${BUCKET_NAME}" --quiet || true

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

   http://console.cloud.google.com/compute/images

מחיקת פרויקט

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

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

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

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

gcloud projects delete PROJECT_ID

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

• פריסה מחדש של אשכול Slurm
• בדיקת ביצועי הרשת באשכול Slurm
• מעקב אחרי מכונות וירטואליות באשכול Slurm
• יצירת נקודת קצה למילוי בקשות: אחרי שיש לכם מודל שעבר כוונון עדין, אתם יכולים לפרוס אותו בנקודת קצה למילוי בקשות באמצעות Google Kubernetes Engine‏ (GKE) או Vertex AI, כדי שיהיה נגיש להסקת מסקנות.
• מידע על שירותי אחסון מומלצים לעומסי עבודה של AI ו-ML