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

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

מטרות

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

עלויות

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

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

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

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

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

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

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

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

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

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

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

הגדרת הסביבה

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

export PROJECT="YOUR_PROJECT_ID"
export ZONE="ZONE_NAME"
export RESERVATION="RESERVATION_NAME"
export NAME="TPU_MACHINE_NAME"

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

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

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

gcloud auth login

יצירת מכונת Cloud TPU וירטואלית

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

gcloud alpha compute tpus tpu-vm create $NAME \
    --zone=$ZONE \
    --project=$PROJECT \
    --accelerator-type=v6e-8 \
    --version=v2-alpha-tpuv6e \
    --provisioning-model=reservation-bound \
    --reservation=$RESERVATION

אחרי שיוצרים את מופע ה-VM, מתחברים אליו באמצעות SSH.

gcloud compute tpus tpu-vm 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 python3.12 python3.12-venv

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

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

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

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

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

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

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

#install_maxtext_tpu_post_train_extra_deps
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/
export RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
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

  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 יהיה מוכן לשימוש. כי כשמפעילים מחדש את ה-VM, מאבדים את הגישה לתוכן של התיקייה /dev/shm. לכן, כדאי להעביר את המודל המכוונן לאחסון מתמיד או להעלות אותו ל-Hugging Face Hub.

הסרת המשאבים

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

מחיקת TPU VM

יוצאים ממופע Cloud TPU VM ואז מוחקים אותו.

gcloud alpha compute tpus tpu-vm delete $NAME --zone=$ZONE --project=$PROJECT --quiet

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

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