אימון של Qwen2 באשכול A4 Slurm

במדריך הזה מוסבר איך לאמן מודל שפה גדול (LLM) באשכול Slurm עם כמה צמתים וכמה מעבדי GPU ב- Google Cloud. המודל שבו משתמשים במדריך הזה מבוסס על מודל Qwen2 עם 1.5 מיליארד פרמטרים. באשכול Slurm נעשה שימוש בשתי מכונות וירטואליות (VM) מסוג a4-highgpu-8g, שבכל אחת מהן יש 8 מעבדי NVIDIA B200 GPU.

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

  1. פריסת אשכול Slurm ברמת ייצור עם ביצועים גבוהים באמצעותGoogle Cloud Cluster Toolkit. כחלק מהפריסה הזו, אתם יוצרים קובץ אימג' של VM בהתאמה אישית עם התוכנה הנדרשת שמותקנת מראש. בנוסף, אתם מגדירים מכונה משותפת של Filestore ורשת RDMA במהירות גבוהה.
  2. אחרי פריסת האשכול, מריצים משימת אימון מראש מבוזרת באמצעות קבוצת הסקריפטים שמצורפת למדריך הזה. המשימה משתמשת בספריית Hugging Face Accelerate.

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

מטרות

  • ניגשים למודל Qwen2 באמצעות Hugging Face.
  • מכינים את הסביבה.
  • ליצור ולפרוס אשכול Slurm מסוג A4 ברמת ייצור.
  • מאמנים את המודל Qwen2 באמצעות הספרייה Accelerate.
  • עוקבים אחרי המשימה.
  • מנקים.

עלויות

במסמך הזה משתמשים ברכיבים הבאים של 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 file.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, roles/compute.osAdminLogin, roles/iap.tunnelResourceAccessor 

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

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

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

  1. חותמים על הסכם ההסכמה לשימוש ב-Qwen 2 1.5B.

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

התקנת Cluster Toolkit

‫Cluster Toolkit הוא כלי בקוד פתוח שמפשט את הפריסה של עומסי עבודה של מחשוב עתיר ביצועים (HPC), בינה מלאכותית (AI) ולמידת מכונה (ML) ב- Google Cloud. למידע נוסף על השימוש ב-CLI של 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 CLUSTER_NAME="YOUR_CLUSTER_NAME"
export ZONE="YOUR_ZONE"
export REGION="YOUR_REGION"
export RESERVATION_URL="YOUR_RESERVATION_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"
export HF_TOKEN="YOUR_HF_TOKEN"

gcloud config set project "${PROJECT_ID}"
gcloud config set billing/quota_project "${PROJECT_ID}"

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

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

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

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

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

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

  1. יוצרים את הקובץ a4high-slurm-deployment.yaml:

    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: ${GCS_BUCKET}

vars:
  deployment_name: ${CLUSTER_NAME}
  project_id: ${PROJECT_ID}
  region: ${REGION}
  zone: ${ZONE}
  a4h_cluster_size: 2
  a4h_reservation_name: ${RESERVATION_URL}
EOF

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

    gcluster create \
  -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
  "${MANIFEST_PATH}/a4high-slurm-blueprint.yaml" \
  -o "${CLUSTER_NAME}"

  3. מתקנים את המניפסטים:

    sed -i '/deletion_protection = {/,/}/ { s/enabled = true/enabled = false/; /reason  = "Avoid data loss"/d; }' "${CLUSTER_NAME}/${CLUSTER_NAME}/cluster-env/main.tf"

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

    gcluster deploy "${CLUSTER_NAME}/${CLUSTER_NAME}" --auto-approve

    הפקודה gcluster deploy היא תהליך דו-שלבי:

    • בשלב הראשון נוצרת תמונה מותאמת אישית עם כל התוכנות שמותקנות מראש. השלב הזה יכול להימשך עד 50 דקות.

    • בשלב השני, האשכול נפרס באמצעות קובץ האימג' בהתאמה אישית. התהליך הזה בדרך כלל נמשך פחות זמן מהשלב הראשון.

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

    gcluster deploy "${CLUSTER_NAME}" --auto-approve --skip "image" -w

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

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

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

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

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

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

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

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

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

    #!/bin/bash
# This script should be run ONCE on the login node to set up the
# shared Python virtual environment.

set -e
echo "--- Creating Python virtual environment in /home ---"
python3 -m venv ~/.venv
echo "--- Activating virtual environment ---"
source ~/.venv/bin/activate

echo "--- Installing build dependencies ---"
pip install --upgrade pip wheel packaging

echo "--- Installing PyTorch for CUDA 12.8 ---"
pip install torch --index-url https://download.pytorch.org/whl/cu128

echo "--- Installing application requirements ---"
pip install -r requirements.txt

echo "--- Environment setup complete. You can now submit jobs with sbatch. ---"

  2. כדי לציין את ההגדרות של משימת הכוונון, יוצרים את הקובץ accelerate_config.yaml עם התוכן הבא:

    compute_environment: "LOCAL_MACHINE"
distributed_type: "FSDP"
downcast_bf16: "no"
fsdp_config:
  fsdp_auto_wrap_policy: "TRANSFORMER_BASED_WRAP"
  fsdp_backward_prefetch: "BACKWARD_PRE"
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: false
  fsdp_sharding_strategy: "FULL_SHARD"
  fsdp_state_dict_type: "SHARDED_STATE_DICT"
  fsdp_transformer_layer_cls_to_wrap: "Qwen2DecoderLayer"
machine_rank: 0
main_training_function: "main"
mixed_precision: "bf16"
num_machines: 2
num_processes: 16
rdzv_backend: "static"
same_network: true
tpu_env: []
use_cpu: false

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

    #SBATCH --job-name=qwen2-pretrain-smollm-fineweb
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1  # 1 main srun task on the node that manages accelerate
#SBATCH --gpus-per-node=8    # access to all 8 GPUs on the node
#SBATCH --partition=a4high
#SBATCH --output=logs/slurm-%j.out
#SBATCH --error=logs/slurm-%j.err

set -euo pipefail
echo "--- Slurm Job Started ---"

# --- STAGE 1: Setup environment and pre-process data on each node's local SSD ---
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 --gpu-bind=none bash -c '
  set -e
  echo "Setting up local environment on $(hostname)..."
  LOCAL_VENV="/mnt/localssd/venv_job_${SLURM_JOB_ID}"
  LOCAL_CACHE="/mnt/localssd/hf_cache_job_${SLURM_JOB_ID}"
  PROCESSED_DATA_DIR="/mnt/localssd/processed_data_${SLURM_JOB_ID}"
  LOCAL_TMP="/mnt/localssd/tmp_job_${SLURM_JOB_ID}"

  rsync -a --info=progress2 ~/.venv/ ${LOCAL_VENV}/
  mkdir -p ${LOCAL_CACHE} ${PROCESSED_DATA_DIR} ${LOCAL_TMP}

  echo "Pre-processing data on $(hostname)..."
  source ${LOCAL_VENV}/bin/activate
  export TMPDIR="${LOCAL_TMP}"
  export TEMP="${LOCAL_TMP}"
  export TMP="${LOCAL_TMP}"
  export HF_HOME="${LOCAL_CACHE}"
  export HF_DATASETS_CACHE="${LOCAL_CACHE}"
  export HF_MODULES_CACHE="${LOCAL_CACHE}/modules"
  export HF_METRICS_CACHE="${LOCAL_CACHE}/metrics"

  python "${HOME}/preprocess_data.py" \
    --dataset_name "HuggingFaceFW/fineweb-edu" \
    --dataset_config "CC-MAIN-2024-10" \
    --tokenizer_id "Qwen/Qwen2-1.5B" \
    --max_seq_length 1024 \
    --output_path ${PROCESSED_DATA_DIR}

  echo "Setup on $(hostname) complete."
'

# --- STAGE 2: Run the Training Job using the Local Environment ---
echo "--- Starting Training ---"

LOCAL_VENV="/mnt/localssd/venv_job_${SLURM_JOB_ID}"
PROCESSED_DATA_DIR="/mnt/localssd/processed_data_${SLURM_JOB_ID}"
LOCAL_OUTPUT_DIR="/mnt/localssd/outputs_${SLURM_JOB_ID}"

export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=29505

# Network and initialization debugging configurations
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=INIT,COLL
export TORCH_DISTRIBUTED_DEBUG=INFO
export NCCL_IB_DISABLE=0

# Launching with full GPU access enabled for accelerate
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 --gpu-bind=none bash -c "
  mkdir -p ${LOCAL_OUTPUT_DIR}
  source ${LOCAL_VENV}/bin/activate

  # Reset the forced Slurm isolation to expose all 8 GPUs to the processes
  unset CUDA_VISIBLE_DEVICES

  # Retrieve the default network device.
  DETECTED_IFACE=\$(ip route show | grep default | awk '{print \$5}' | head -n 1)
  echo \"[INFO] Automatically detected cluster network interface: \$DETECTED_IFACE\"

  # Dynamic injection of the detected interface into the network configuration
  export NCCL_SOCKET_IFNAME=\"\${DETECTED_IFACE},gpu*\"
  export TP_SOCKET_IFNAME=\"\${DETECTED_IFACE}\"
  export GLOO_SOCKET_IFNAME=\"\${DETECTED_IFACE}\"

  accelerate launch \
    --config_file ~/accelerate_config.yaml \
    --num_machines \$SLURM_NNODES \
    --num_processes \$((SLURM_NNODES * 8)) \
    --machine_rank \$SLURM_NODEID \
    --main_process_ip \$MASTER_ADDR \
    --main_process_port \$MASTER_PORT \
    train.py \
      --model_config_id 'Qwen/Qwen2-1.5B' \
      --preprocessed_data_path ${PROCESSED_DATA_DIR} \
      --output_dir ${LOCAL_OUTPUT_DIR} \
      --per_device_train_batch_size 4 \
      --gradient_accumulation_steps 4 \
      --max_steps 10000 \
      --learning_rate 5e-5 \
      --save_strategy steps \
      --save_steps 500 \
      --logging_steps 1
"

# --- STAGE 3: Copy Final Model from Local SSD to Home Directory ---
echo "--- Copying final model from local SSD to /home ---"
mkdir -p ~/qwen2-from-scratch-on-smollm-fineweb/

srun --nodes=1 --ntasks=1 --ntasks-per-node=1 bash -c "
  rsync -a --info=progress2 ${LOCAL_OUTPUT_DIR}/ ~/qwen2-from-scratch-on-smollm-fineweb/
"

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

  4. כדי לציין את יחסי התלות של משימת הכוונון, יוצרים קובץ requirements.txt עם התוכן הבא:

    # Hugging Face Libraries (Pinned to recent, stable versions for reproducibility)
transformers==4.53.3
datasets==4.0.0
accelerate==1.9.0
evaluate==0.4.5
bitsandbytes==0.46.1
trl==0.19.1
peft==0.16.0

# Other dependencies
tensorboard==2.20.0
protobuf==6.31.1
sentencepiece==0.2.0
    הערה: במדריך הזה לא משתמשים בגרסאות העדכניות ביותר של יחסי התלות של NVIDIA ו-PyTorch. אם אתם צריכים יחסי תלות חדשים יותר, עיינו במסמכי התיעוד של NVIDIA ובמסמכי התיעוד של PyTorch.

  5. כדי להוריד את מערך הנתונים, להמיר אותו לאסימונים ולעבד אותו מראש לפורמט שמוכן לאימון, יוצרים קובץ preprocess_data.py עם התוכן הבא:

    import argparse
from datasets import load_dataset
from transformers import AutoTokenizer
import os
from itertools import chain

def get_args():
   parser = argparse.ArgumentParser(description="Download and preprocess a dataset.")
   parser.add_argument("--dataset_name", type=str, required=True)
   parser.add_argument("--dataset_config", type=str, required=True)
   parser.add_argument("--tokenizer_id", type=str, required=True)
   parser.add_argument("--max_seq_length", type=int, required=True)
   parser.add_argument("--output_path", type=str, required=True, help="Path to save the processed dataset.")
   return parser.parse_args()

def main():
   args = get_args()

   if os.path.exists(args.output_path) and os.listdir(args.output_path):
       print(f"Processed dataset already exists at {args.output_path}. Skipping.")
       return

   # 1. Load tokenizer
   tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id)

   # 2. Load raw dataset
   print(f"Loading raw dataset {args.dataset_name}...")
   raw_dataset = load_dataset(args.dataset_name, name=args.dataset_config, split="train")

   # 3. Tokenize
   def tokenize_function(examples):
       return tokenizer(examples["text"])

   num_proc = os.cpu_count()
   print(f"Tokenizing dataset using {num_proc} processes...")
   print("Tokenizing dataset...")
   tokenized_dataset = raw_dataset.map(
       tokenize_function,
       batched=True,
       remove_columns=raw_dataset.column_names,
       desc="Running tokenizer on dataset",
       num_proc=num_proc,
   )

   # 4. Group texts
   def group_texts(examples):
       concatenated_examples = {k: list(chain.from_iterable(examples[k])) for k in examples.keys()}
       total_length = len(concatenated_examples[list(examples.keys())[0]])
       if total_length >= args.max_seq_length:
           total_length = (total_length // args.max_seq_length) * args.max_seq_length
       result = {
           k: [t[i : i + args.max_seq_length] for i in range(0, total_length, args.max_seq_length)]
           for k, t in concatenated_examples.items()
       }
       result["labels"] = result["input_ids"].copy()
       return result

   print("Grouping texts...")
   lm_dataset = tokenized_dataset.map(
       group_texts,
       batched=True,
       desc=f"Grouping texts in chunks of {args.max_seq_length}",
       num_proc=num_proc,
   )

   # 5. Save to disk
   print(f"Saving processed dataset to {args.output_path}...")
   lm_dataset.save_to_disk(args.output_path)
   print("Preprocessing complete.")

if __name__ == "__main__":
   main()

  6. כדי לציין את ההוראות למשימה, יוצרים קובץ train.py עם התוכן הבא:

    import torch
import argparse
from datasets import load_dataset, load_from_disk
import os
from transformers import (
    AutoConfig,
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)
from huggingface_hub import login

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_config_id", type=str, default="Qwen/Qwen2-1.5B", help="Hugging Face model config to use for architecture.")
    # Data arguments - used if preprocessed data is not available
    parser.add_argument("--dataset_name", type=str, default="HuggingFaceFW/fineweb-edu", help="Hugging Face dataset for pre-training.")
    parser.add_argument("--dataset_config", type=str, default="CC-MAIN-2024-10", help="Config for the fineweb-edu dataset, e.g., 'CC-MAIN-2024-10'.")
    parser.add_argument("--preprocessed_data_path", type=str, default=None, help="Path to a preprocessed dataset on disk. If provided, skips download and processing.")
    # General arguments
    parser.add_argument("--hf_token", type=str, default=None, help="Hugging Face token for private models/tokenizers")
    parser.add_argument("--output_dir", type=str, default="qwen2-from-scratch-on-smollm-fineweb", help="Directory to save model checkpoints")

    # TrainingArguments
    parser.add_argument("--max_seq_length", type=int, default=1024, help="Maximum sequence length")
    parser.add_argument("--num_train_epochs", type=int, default=1, help="Number of training epochs")
    parser.add_argument("--max_steps", type=int, default=-1, help="If set to a positive number, it overrides num_train_epochs.")
    parser.add_argument("--per_device_train_batch_size", type=int, default=4, help="Batch size per device during training")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=4, help="Gradient accumulation steps")
    parser.add_argument("--learning_rate", type=float, default=5e-5, help="Learning rate")
    parser.add_argument("--logging_steps", type=int, default=10, help="Log every X steps")
    parser.add_argument("--save_strategy", type=str, default="steps", help="Checkpoint save strategy")
    parser.add_argument("--save_steps", type=int, default=500, help="Save checkpoint every X steps")

    return parser.parse_args()

def main():
    args = get_args()

    # --- 1. Setup and Login ---
    if args.hf_token:
        login(args.hf_token)

    # --- 2. Load Tokenizer ---
    # We load the tokenizer from the specified config ID to ensure compatibility
    # with the model architecture (e.g., special tokens).
    tokenizer = AutoTokenizer.from_pretrained(args.model_config_id)

    # --- 3. Initialize Model from Scratch ---
    print(f"Initializing a new model from {args.model_config_id} configuration...")
    config = AutoConfig.from_pretrained(args.model_config_id)
    model = AutoModelForCausalLM.from_config(config)

    print(f"Model has {model.num_parameters():,} parameters.")

    # --- 4. Load or Create and prepare the training dataset ---
    if args.preprocessed_data_path and os.path.exists(args.preprocessed_data_path):
        print(f"Loading preprocessed dataset from {args.preprocessed_data_path}...")

        # Synchronization of distributed processes
        local_rank = int(os.environ.get("LOCAL_RANK", -1))
        if local_rank != -1:
            # Introducing a minimal time offset per GPU to avoid I/O collisions.
            import time
            time.sleep(local_rank * 0.2)

        lm_dataset = load_from_disk(args.preprocessed_data_path, keep_in_memory=False)

    else:
        print("No preprocessed dataset found, starting from raw data...")
        raw_dataset = load_dataset(args.dataset_name, name=args.dataset_config, split="train")

        # Tokenization function
        def tokenize_function(examples):
            return tokenizer(examples["text"])

        tokenized_dataset = raw_dataset.map(
            tokenize_function,
            batched=True,
            remove_columns=raw_dataset.column_names,
            desc="Running tokenizer on dataset",
        )

        # Main data processing function that will concatenate all texts from our dataset
        # and generate chunks of max_seq_length.
        def group_texts(examples):
            # Concatenate all texts.
            concatenated_examples = {k: [item for sublist in examples[k] for item in sublist] for k in examples.keys()}
            total_length = len(concatenated_examples[list(examples.keys())[0]])
            # We drop the small remainder.
            if total_length >= args.max_seq_length:
                total_length = (total_length // args.max_seq_length) * args.max_seq_length
            # Split by chunks of max_len.
            result = {
                k: [t[i : i + args.max_seq_length] for i in range(0, total_length, args.max_seq_length)]
                for k, t in concatenated_examples.items()
            }
            result["labels"] = result["input_ids"].copy()
            return result

        lm_dataset = tokenized_dataset.map(
            group_texts,
            batched=True,
            desc=f"Grouping texts in chunks of {args.max_seq_length}",
        )


    # --- 5. Configure Training Arguments ---
    # Check for bfloat16 support
    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()

    training_args = TrainingArguments(
        output_dir=args.output_dir,
        num_train_epochs=args.num_train_epochs,
        max_steps=args.max_steps,
        per_device_train_batch_size=args.per_device_train_batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate,
        logging_steps=args.logging_steps,
        save_strategy=args.save_strategy,
        save_steps=args.save_steps,
        save_total_limit=2, # Optional: Limit the number of checkpoints
        bf16=use_bf16,
        fp16=not use_bf16,
        optim="adamw_torch",
        lr_scheduler_type="cosine",
        warmup_ratio=0.03,
        report_to="tensorboard",
        gradient_checkpointing=True,
        # Required for gradient checkpointing with some parallelization strategies
        gradient_checkpointing_kwargs={"use_reentrant": False},
    )

    # --- 6. Create Trainer and Start Training ---
    # Data collator will take care of creating batches for causal language modeling
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=lm_dataset,
        # eval_dataset=... # Optional: if you have a validation set
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    print("Starting training from scratch...")
    trainer.train()
    print("Training finished.")

    # --- 7. Save the final model ---
    print(f"Saving final model to {args.output_dir}")
    trainer.save_model()

if __name__ == "__main__":
    main()

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

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

  1. מגדירים את המשתנה LOGIN_NODE על ידי אחזור השם של צומת ההתחברות של האשכול:

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

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

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

    # Identify the VPC network used by the Slurm cluster's login node
VPC_NETWORK_URI=$(gcloud compute instances describe "${LOGIN_NODE}" \
  --project="${PROJECT_ID}" \
  --zone="${ZONE}" \
  --format="value(networkInterfaces[0].network)")
CLUSTER_NETWORK=$(basename "${VPC_NETWORK_URI}")

# Create the IAP firewall rule to allow SSH access to the login node
gcloud compute firewall-rules create "allow-ssh-from-iap-${CLUSTER_NETWORK}" \
  --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.txt \
  ./submit.slurm \
  ./accelerate_config.yaml \
  ./train.py \
  ./preprocess_data.py \
  "${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 וכלים

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

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

    chmod +x install_environment.sh
./install_environment.sh

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

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

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

    sbatch submit.slurm

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

    tail -f ~/logs/slurm-1.err # (or .out, depending on where the script is currently sending logs)

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

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

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

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

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

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

  • משך העבודה: העבודה אמורה להימשך כשעה.

הורדת המודל

אחרי שהמשימה תרוץ בהצלחה, המודל המאומן יישמר בספרייה ~/qwen2-from-scratch-on-smollm-fineweb/ בצומת ההתחברות. מכיוון שהספרייה המשותפת הקבועה הזו מותקנת בכל הצמתים באשכול, נקודות הבדיקה של המודל נשארות זמינות גם אחרי שהמשימה מסתיימת או שההקצאה של צומתי המחשוב מבוטלת.

כדי להוריד את המודל השמור מצומת ההתחברות למכונה המקומית, משתמשים בפקודה gcloud compute scp, כמו בדוגמה הבאה:

# From your local machine
gcloud compute scp --project="${PROJECT_ID}" --zone="${ZONE}" --tunnel-through-iap \
  "${LOGIN_NODE}":~/qwen2-from-scratch-on-smollm-fineweb/ ./qwen2-trained-model/ --recurse

אחרי שמורידים את המודל, אפשר:

  • טעינת המודל להסקת מסקנות: משתמשים ב-framework של Hugging Face Transformers כדי לטעון את הספרייה qwen2-trained-model/ ולהסיק מסקנות באמצעות המודל החדש של Qwen2 שאומן.
  • כוונון עדין נוסף: משתמשים בנקודת הבדיקה השמורה כנקודת התחלה לכוונון עדין נוסף במערך נתונים ספציפי יותר.
  • העלאת המודל ל-Hugging Face Hub: משתפים את המודל המאומן על ידי העלאה שלו ל-Hugging Face Hub.

הסרת המשאבים

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

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

  1. כדי למחוק את אשכול Slurm, מריצים את הפקודה gcluster destroy:

    ./gcluster destroy "${CLUSTER_NAME}" --auto-approve

  2. כדי למחוק את הקטגוריה של Cloud Storage:

    gcloud storage buckets delete "gs://${GCS_BUCKET}" --quiet || true

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

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

  4. כדי למחוק את כל רשתות ה-VPC, כללי חומת האש, הנתבים, כתובות ה-IP ורשתות המשנה שמשויכים לפרויקט:

    echo "========================================================================="
echo " STARTING AUTOMATED NETWORK CLEANUP FOR CLUSTER: ${CLUSTER_NAME}"
echo "========================================================================="

echo "Discovering all VPC networks linked to the cluster..."
NETWORKS=$(gcloud compute networks list --project="${PROJECT_ID}" --format="value(name)" | grep "^${CLUSTER_NAME}" || true)

if [ -z "${NETWORKS}" ]; then
    echo "No VPC networks found starting with ${CLUSTER_NAME}. Everything is already clean!"
    exit 0
fi

echo "Found the following networks to process:"
echo "${NETWORKS}"
echo "-------------------------------------------------------------------------"

echo "=== 1. Wiping Global Firewall Rules ==="
FIREWALL_RULES=$(gcloud compute firewall-rules list \
    --project="${PROJECT_ID}" \
    --filter="network ~ ^${CLUSTER_NAME} OR name ~ ^${CLUSTER_NAME}" \
    --format="value(name)" || echo "")

if [ -n "${FIREWALL_RULES}" ]; then
    echo "Deleting matching firewall rules:"
    echo "${FIREWALL_RULES}"
    echo "${FIREWALL_RULES}" | xargs -r gcloud compute firewall-rules delete --project="${PROJECT_ID}" --quiet
else
    echo "No matching firewall rules found."
fi

echo "=== 2. Tearing Down Network-Specific Infrastructure ==="
echo "${NETWORKS}" | while read -r net_name; do
    [ -z "${net_name}" ] && continue
    echo "Processing resources for network: ${net_name}"

    ROUTERS=$(gcloud compute routers list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="network=${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${ROUTERS}" ]; then
        echo "  -> Deleting routers: ${ROUTERS}"
        echo "${ROUTERS}" | xargs -r gcloud compute routers delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi

    IPS=$(gcloud compute addresses list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="name ~ ^${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${IPS}" ]; then
        echo "  -> Deleting IP reservations: ${IPS}"
        echo "${IPS}" | xargs -r gcloud compute addresses delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi

    SUBNETS=$(gcloud compute networks subnets list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="network=${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${SUBNETS}" ]; then
        echo "  -> Deleting subnetworks:"
        echo "${SUBNETS}"
        echo "${SUBNETS}" | xargs -r gcloud compute networks subnets delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi
done

echo "-------------------------------------------------------------------------"
echo "Waiting 15 seconds for Google Cloud API dependencies to unlock..."
sleep 15

echo "=== 3. Final VPC Networks Destruction ==="
echo "${NETWORKS}" | while read -r net_name; do
    [ -z "${net_name}" ] && continue
    echo "Deleting core VPC network: ${net_name}..."
    gcloud compute networks delete "${net_name}" --project="${PROJECT_ID}" --quiet || \
    echo "Warning: Could not delete ${net_name} yet. If a lock occurred, please rerun in 1 minute."
done

echo "========================================================================="
echo " SUCCESS: All network resources for cluster ${CLUSTER_NAME} have been wiped!"
echo "========================================================================="

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

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

gcloud projects delete PROJECT_ID

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