איך מריצים אימון של למידה חיזוקית במכונה וירטואלית של TPU באמצעות MaxText

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

מטרות

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

עלויות

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

• Cloud TPU

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

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

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

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

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

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

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

• באתר Hugging Face, מאשרים את הסכם הרישיון של המודל שרוצים לאמן. במדריך הזה נעשה שימוש במודל llama3.1-8b-Instruct.

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

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

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

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

הגדרת הסביבה

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

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

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

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

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

gcloud auth login

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

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

gcloud alpha compute tpus tpu-vm create $TPU_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 $TPU_NAME --zone $ZONE --project $PROJECT

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

התקנת 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_tpu_post_train_extra_deps

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

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

מספקים את הערכים הבאים:

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

export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='llama3.1-8b-Instruct'
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?}

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

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

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

   # -- MaxText configuration --
   export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
   export RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
   export CHIPS_PER_VM=8
   export NUM_BATCHES=50
   export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items

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

   python3 -m maxtext.trainers.post_train.rl.train_rl \
       model_name=${MODEL_NAME?} \
       load_parameters_path=${MAXTEXT_CKPT_PATH?} \
       run_name=${RUN_NAME?} \
       base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
       chips_per_vm=${CHIPS_PER_VM?} \
       num_batches=${NUM_BATCHES?} \
       num_test_batches=10 \
       rollout_data_parallelism=1 \
       rollout_tensor_parallelism=-1

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

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

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

   export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
   export HF_MODEL_NAME=llama3.1-8b
   export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/actor/$NUM_BATCHES/model_params

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

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

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

הסרת המשאבים

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

מחיקת מופע של מכונת TPU וירטואלית

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

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

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

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