הפעלת כוונון מפוקח (SFT) במכונה וירטואלית של TPU באמצעות MaxText

במדריך הזה מוסבר איך להריץ כוונון עדין מפוקח (SFT) במכונה וירטואלית (VM) אחת של v6e-8 יחידת עיבוד טנסור (TPU) ב- Google Cloud באמצעות MaxText, מחסנית אימון מבוססת JAX עם ביצועים גבוהים למודלים גדולים של שפה (LLM).

סיכום עומס העבודה

הרצת עומס העבודה של SFT במכונה וירטואלית של Cloud TPU כוללת את הרכיבים הבאים:

• מודל: Gemma 3 4B (gemma3-4b), מודל שפה קל משקל עם משקלים פתוחים. במדריך הזה, משקלי המודל מומרים מהפורמט של Hugging Face לפורמט שתואם ל-MaxText‏ (mt-format) ומוכנים להפעלה במכונה וירטואלית אחת של Cloud TPU v6e-8 באמצעות JAX.
• שיטת האימון: SFT באמצעות MaxText. עומס העבודה מריץ צינור (pipeline) של כוונון עדין במופע יחיד של Cloud TPU v6e-8 באמצעות JAX. הוא מבצע אופטימיזציה של פרמטרים של מודלים ביותר מ-1,000 שלבים באמצעות כלי SFT מובנה של MaxText‏ (maxtext.trainers.post_train.sft.train_sft) ומערכי נתונים של כוונון הוראות שמתמקדים בדיאלוג.
• קבוצת נתונים ותבנית: ‏UltraChat 200k‏ (HuggingFaceH4/ultrachat_200k), קבוצת נתונים גדולה ואיכותית של שיחות רב-שלביות (פיצול train_sft עם עמודה messages). הוא משמש לכוונון המודל ולשיפור היכולות שלו לנהל שיחות ולמלא הוראות.
• תשתית ומסגרת: מכונת VM אחת של Cloud TPU‏ v6e-8 (ct6e-standard-8t). המסגרת היא MaxText‏ (maxtext[tpu-post-train]) שמבוססת על JAX.

מטרות

• מגדירים מכונה וירטואלית של Cloud TPU.
• מתקינים את MaxText ואת יחסי התלות שלו.
• המרת מודל של Hugging Face לפורמט MaxText.
• הרצת עומס עבודה של אימון SFT ב-TPU.
• ממירים את המודל המכוונן בחזרה לפורמט Hugging Face לצורך הצגה.

עלויות

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

• Cloud TPU

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

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

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

חשוב: כדי להשלים את המדריך הזה, צריכה להיות לכם גישה ל-v6e-8 TPUs. במדריך הזה אנחנו מניחים שכבר יש לכם הזמנה של יחידות TPU בפרויקטGoogle Cloud כדי לוודא שהקיבולת תהיה זמינה כשאתם צריכים אותה. במאמר מידע על הקצאות מקום ל-Cloud TPU תוכלו לקרוא איך מקצים מקום ל-Cloud TPU לתאריך ולשעה בעתיד.

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

• כדי להשתמש במדריך הזה, צריך אסימון גישה של Hugging Face. אתם יכולים להירשם לחשבון בחינם ב-Hugging Face. אחרי שיוצרים חשבון, יוצרים אסימון גישה:

  1. בדף Welcome to Hugging Face, לוחצים על הדמות של החשבון ובוחרים באפשרות טוקן גישה.
  2. בדף Access tokens (אסימוני גישה), לוחצים על Create new token (יצירת אסימון חדש).
  3. בוחרים את סוג הטוקן Read ומזינים שם לטוקן.
  4. טוקן הגישה שלכם יוצג. שומרים את הטוקן במקום בטוח.

• באתר Hugging Face, מאשרים את הסכם הרישיון של המודל שמתכננים לאמן. במדריך הזה משתמשים במודל gemma3-4b.

כדי לקבל את ההרשאות שדרושות להשלמת המדריך הזה, צריך לבקש מהאדמין להקצות לכם בפרויקט את תפקידי ה-IAM הבאים:

• אדמין TPU (roles/tpu.admin)
• משתמש בחשבון שירות (roles/iam.serviceAccountUser)
• עורך מחשוב (roles/compute.editor)

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

יכול להיות שאפשר לקבל את ההרשאות הנדרשות גם באמצעות תפקידים בהתאמה אישית או תפקידים מוגדרים מראש.

הגדרת הסביבה

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

export PROJECT="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export RESERVATION="YOUR_RESERVATION_NAME"
export NAME="YOUR_TPU_NAME"
export NETWORK="default"

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

• ‫YOUR_PROJECT_ID: מזהה הפרויקט ב- Google Cloud
• ‫YOUR_ZONE: האזור שבו רוצים להשתמש
• ‫YOUR_RESERVATION_NAME: הקצאת הקיבולת שלכם
• YOUR_TPU_NAME: השם של מכונת ה-TPU VM של Cloud TPU

מבצעים אימות באמצעות Google Cloud על ידי הפעלת הפקודה הבאה:

gcloud auth login

יצירת Cloud TPU VM

יוצרים מכונת Cloud TPU VM עם 8 שבבי v6e TPU, שמקושרת להזמנת הקיבולת.

gcloud compute instances create "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --network="${NETWORK}" \
    --tags="${NAME}" \
    --machine-type=ct6e-standard-8t \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --boot-disk-size=200GB \
    --maintenance-policy=TERMINATE \
    --instance-termination-action=DELETE \
    --provisioning-model=RESERVATION_BOUND \
    --reservation-affinity=specific \
    --reservation="${RESERVATION}"

אחרי שיוצרים את המכונה הווירטואלית, מתחברים אליה באמצעות SSH.

gcloud compute ssh "${NAME}" --zone "${ZONE}" --project "${PROJECT}"

מבצעים את השלבים הבאים ב-TPU VM.

התקנת MaxText

מעדכנים את חבילות המערכת במכונה הווירטואלית של TPU.

sudo apt update && sudo apt upgrade -y --fix-missing

מתקינים את Python 3.12, שנדרש ל-MaxText, ואת חבילת הסביבה הווירטואלית שלו.

sudo apt install -y build-essential cmake ninja-build

אפשר להשתמש ב-uv כדי להאיץ את ההתקנה של חבילת Python.

curl -LsSf https://astral.sh/uv/install.sh | sh
source "${HOME}/.local/bin/env"

יוצרים סביבה וירטואלית בשם maxtext_venv ומפעילים אותה.

uv python install 3.12
uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

מתקינים את MaxText ואת יחסי התלות שנדרשים למשימות שאחרי האימון.

UV_TORCH_BACKEND=cpu uv pip install "maxtext[tpu-post-train]==0.2.2" --resolution=lowest

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

install_tpu_post_train_extra_deps

המרת המודל לפורמט MaxText

כדי לאמן את המודל בפורמט MaxText, צריך להמיר אותו מפורמט Hugging Face לפורמט MaxText.

מציינים את משתני הסביבה, כמו אסימון הגישה שלכם ל-Hugging Face, השם של המודל שבו רוצים להשתמש והספרייה שבה רוצים לשמור את המודל בפורמט MaxText.

export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='gemma3-4b'
export MODEL_CHECKPOINT_DIRECTORY=/dev/shm/$MODEL_NAME/mt-format/
export USE_PATHWAYS=0 # Set to 1 for Pathways, 0 for McJAX
export LAZY_LOAD_TENSORS=False # True to use lazy load, False to use eager load.

מחליפים את YOUR_HF_TOKEN באסימון הגישה של Hugging Face שיצרתם קודם.

כדי להמיר את המודל מפורמט Hugging Face לפורמט MaxText, מריצים את הסקריפט הבא. ההמרה הזו נמשכת כחמש דקות.

python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name="${MODEL_NAME?}" \
    hf_access_token="${HF_TOKEN?}" \
    base_output_directory="${MODEL_CHECKPOINT_DIRECTORY?}" \
    scan_layers=True \
    use_multimodal=False \
    hardware=cpu \
    skip_jax_distributed_system=true \
    checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) \
    checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) \
    --lazy_load_tensors="${LAZY_LOAD_TENSORS?}"

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

אחרי שתהליך ההמרה יסתיים, תוכלו להתחיל את עומס העבודה של SFT.

1. מגדירים את הפרמטרים של אימון עומס העבודה של SFT.

   # -- MaxText configuration --
   export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
   RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
   export RUN_NAME
   export STEPS=1000
   export PER_DEVICE_BATCH_SIZE=1
   
   # -- Dataset configuration --
   export DATASET_NAME="HuggingFaceH4/ultrachat_200k"
   export TRAIN_SPLIT="train_sft"
   export TRAIN_DATA_COLUMNS="['messages']"
   
   export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items
   export TPU_ACCELERATOR_TYPE=v6e-8
   export TPU_WORKER_ID=0
   TPU_NAME=$(hostname)
   export TPU_NAME
   export TPU_SKIP_MDS_QUERY=1
   export TPU_WORKER_HOSTNAMES=localhost
   export TPU_TOPOLOGY=2x4
   export TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
   export TPU_HOST_BOUNDS=1,1,1

2. מתחילים את משימת האימון. התהליך נמשך כ-10 דקות במכונה וירטואלית של v6e-8.

   python3 -m maxtext.trainers.post_train.sft.train_sft \
       run_name="${RUN_NAME?}" \
       base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
       model_name="${MODEL_NAME?}" \
       load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
       per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \
       steps="${STEPS?}" \
       hf_path="${DATASET_NAME?}" \
       train_split="${TRAIN_SPLIT?}" \
       train_data_columns="${TRAIN_DATA_COLUMNS?}" \
       profiler=xplane

המרת המודל המאומן בחזרה לפורמט Hugging Face

אחרי שסיימתם את עומס העבודה של האימון, המירו את המודל בחזרה לפורמט Hugging Face.

1. מגדירים את הנתיבים לייצוא ואת הפרמטרים שאומנו.

   export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
   export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/$STEPS/model_params

2. ממירים את המודל בחזרה לפורמט של Hugging Face.

   python3 -m maxtext.checkpoint_conversion.to_huggingface \
       model_name="${MODEL_NAME}" \
       load_parameters_path="${POST_TRAIN_PATH}" \
       base_output_directory="${HF_EXPORT}" \
       scan_layers=True \
       use_multimodal=False \
       weight_dtype=bfloat16

אחרי שההמרה תושלם, המודל המותאם שמאוחסן ב-/dev/shm/gemma3-4b/hf-trained יהיה מוכן לשימוש. כי כשמפעילים מחדש את המכונה הווירטואלית, מאבדים את הגישה לתוכן של התיקייה /dev/shm. לכן, כדאי להעביר את המודל המכוונן לאחסון מתמיד או להעלות אותו ל-Hugging Face Hub.

הסרת המשאבים

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

מחיקת מופע TPU VM

יוצאים מהמכונה הווירטואלית של Cloud TPU ומוחקים אותה.

# 1. Delete TPU instance
echo "Deleting TPU instance: ${NAME}..."
gcloud compute instances delete "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --quiet || true

# 2. Delete IAP firewall rule
echo "Deleting firewall rule: ${FIREWALL_RULE_NAME}..."
gcloud compute firewall-rules delete "${FIREWALL_RULE_NAME}" \
    --project="${PROJECT}" \
    --quiet || true

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

• למידע נוסף על Cloud TPU, קראו את המאמר מבוא ל-Cloud TPU.
• פרטים על הארכיטקטורה וההגדרה של v6e-8 TPU זמינים במאמר TPU v6e.