‫Google משתמשת בטכנולוגיית AI כדי לתרגם תוכן לשפה המועדפת עליך. בתרגומים כאלו עשויות להיות שגיאות.

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

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

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

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

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

מטרות

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

עלויות

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

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

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

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

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

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

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

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

gcloud init

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

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

Select a project: כדי לבחור פרויקט לא צריך תפקיד IAM ספציפי – אפשר לבחור כל פרויקט שקיבלתם בו תפקיד.
יצירת פרויקט: כדי ליצור פרויקט, צריך את התפקיד Project Creator (יצירת פרויקטים) (roles/resourcemanager.projectCreator), שכולל את ההרשאה resourcemanager.projects.create. איך מקצים תפקידים

יוצרים Google Cloud פרויקט:
```
gcloud projects create PROJECT_ID
```
מחליפים את PROJECT_ID בשם של פרויקט Google Cloud שיוצרים.
בוחרים את הפרויקט שיצרתם: Google Cloud
```
gcloud config set project PROJECT_ID
```
מחליפים את PROJECT_ID בשם הפרויקט ב- Google Cloud .

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

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

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

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

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

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

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

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

gcloud init

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

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

Select a project: כדי לבחור פרויקט לא צריך תפקיד IAM ספציפי – אפשר לבחור כל פרויקט שקיבלתם בו תפקיד.
יצירת פרויקט: כדי ליצור פרויקט, צריך את התפקיד Project Creator (יצירת פרויקטים) (roles/resourcemanager.projectCreator), שכולל את ההרשאה resourcemanager.projects.create. איך מקצים תפקידים

יוצרים Google Cloud פרויקט:
```
gcloud projects create PROJECT_ID
```
מחליפים את PROJECT_ID בשם של פרויקט Google Cloud שיוצרים.
בוחרים את הפרויקט שיצרתם: Google Cloud
```
gcloud config set project PROJECT_ID
```
מחליפים את PROJECT_ID בשם הפרויקט ב- Google Cloud .

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

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

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

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

מעניקים תפקידים לחשבון המשתמש. מריצים את הפקודה הבאה לכל אחד מהתפקידים הבאים ב-IAM: roles/compute.admin, roles/iam.serviceAccountUser, roles/file.editor, roles/storage.admin, roles/serviceusage.serviceUsageAdmin
```
gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫PROJECT_ID: מזהה הפרויקט.
- ‫USER_IDENTIFIER: המזהה של חשבון המשתמש . לדוגמה, myemail@example.com.
- ‫ROLE: תפקיד ה-IAM שאתם מקצים לחשבון המשתמש.
מפעילים את חשבון השירות שמוגדר כברירת מחדל עבור Google Cloud הפרויקט:
```
gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --project=PROJECT_ID
```
מחליפים את PROJECT_NUMBER במספר הפרויקט. כדי לבדוק את מספר הפרויקט, אפשר לעיין במאמר קבלת פרויקט קיים.

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

gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \
    --role=roles/editor

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

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

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

נכנסים לחשבון Hugging Face או יוצרים חשבון חדש.

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

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

הכנת הסביבה

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

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

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

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

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

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

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

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

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

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

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

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

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

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

#!/bin/bash
# This script 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. ---"

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

# Default configuration for a 2-node, 8-GPU-per-node (16 total GPUs) FSDP training job.

compute_environment: "LOCAL_MACHINE"
distributed_type: "FSDP"
downcast_bf16: "no"
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

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

#!/bin/bash
#SBATCH --job-name=qwen2-pretrain-smollm-fineweb
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8 # 8 tasks per node
#SBATCH --gpus-per-task=1   # 1 GPU per task
#SBATCH --partition=a4high
#SBATCH --output=logs/slurm-%j.out
#SBATCH --error=logs/slurm-%j.err

set -e
echo "--- Slurm Job Started ---"

# --- STAGE 1: Setup environment and pre-process data on each node's local SSD ---
# This command runs once per node.
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 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}"
  rsync -a --info=progress2 ~/./.venv/ ${LOCAL_VENV}/
  mkdir -p ${LOCAL_CACHE} ${PROCESSED_DATA_DIR}

  echo "Pre-processing data on $(hostname)..."
  source ${LOCAL_VENV}/bin/activate
  export HF_HOME=${LOCAL_CACHE}
  export HF_DATASETS_CACHE=${LOCAL_CACHE}

  # This runs the new preprocessing script. It ensures only ONE process per node
  # downloads and processes the data, avoiding rate limiting and redundant work.
  python 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}"
mkdir -p ${LOCAL_OUTPUT_DIR}

# This is the main training command. It launches one Python process per GPU.
srun --ntasks=$((SLURM_NNODES * 8)) --gpus-per-task=1 bash -c "
  source ${LOCAL_VENV}/bin/activate

  # The training script now loads the pre-processed data from the local SSD.
  python 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
"

# --- STAGE 3: Copy Final Model from Local SSD to Home Directory ---
echo "--- Copying final model from local SSD to /home ---"
# This command runs only on the first node of the job allocation
# and copies the final model back to the persistent shared directory.
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 ---"

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

כדי להוריד, לבצע טוקניזציה ולעבד מראש את מערך הנתונים לפורמט שמוכן לאימון, יוצרים קובץ 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()

כדי לציין את ההוראות למשימה, יוצרים קובץ 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 smollm-corpus dataset, e.g., 'fineweb-edu-dedup'.")
    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-olmo", 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)

    # --- 4. 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.")

    # --- 3. 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}...")
        lm_dataset = load_from_disk(args.preprocessed_data_path)
    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: sum(examples[k], []) 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, פועלים לפי השלבים הבאים:

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

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

gcloud compute scp \
  --project=PROJECT_ID \
  --zone=ZONE \
  --tunnel-through-iap \
  ./train.py \
  ./requirements.txt \
  ./submit.slurm \
  ./install_environment.sh \
  ./accelerate_config.yaml \
  "LOGIN_NODE_NAME":~/

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

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

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

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

התקנת frameworks וכלים

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

יוצרים משתנה סביבה לאסימון הגישה שלכם ל-Hugging Face:
```
export HUGGING_FACE_TOKEN="HUGGING_FACE_TOKEN"
```
מגדירים סביבה וירטואלית של Python עם כל הרכיבים התלויים הנדרשים:
```
chmod +x install_environment.sh
./install_environment.sh
```

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

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

שולחים את העבודה למתזמן של Slurm:
```
sbatch submit.slurm
```
בצומת הכניסה באשכול Slurm, אפשר לעקוב אחרי התקדמות העבודה על ידי בדיקת קובצי הפלט שנוצרו בספרייה home:
```
tail -f logs/slurm-qwen2-pretrain-smollm-fineweb.err
```
אם העבודה מתחילה בהצלחה, בקובץ .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

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

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

הורדת המודל

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

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

# From your local machine
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 \
  "${LOGIN_NODE_NAME}":~/qwen2-from-scratch-on-smollm-fineweb/ ./qwen2-trained-model/ --recurse

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

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

הסרת המשאבים

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

מחיקת פרויקט

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

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

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

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

gcloud projects delete PROJECT_ID

מחיקת אשכול Slurm

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

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