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

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

מטרות

  • מגדירים מופע של 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="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 וירטואלית

יוצרים מכונת 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}"

אחרי שיוצרים את מופע ה-VM, מתחברים אליו באמצעות 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

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

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

הסרת המשאבים

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

מחיקת מופע של TPU VM

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

# 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.