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

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

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

הפעלת עומס העבודה הזה של RL אחרי אימון ב-Cloud TPU כוללת את הרכיבים הבאים:

• מודל: Llama 3.1 8B Instruct (llama3.1-8b-Instruct), מודל LLM עם משקלים פתוחים שמותאם לדיאלוגים ולמשימות של מילוי הוראות. במדריך הזה, משקלי המודל מומרים מהפורמט של Hugging Face לפורמט שתואם ל-MaxText‏ (mt-format) ומוכנים להפעלה מבוזרת ב-TPU של v6e-8 באמצעות JAX.

• אלגוריתם של למידת חיזוק (RL): אופטימיזציה של מדיניות יחסית לקבוצה (GRPO), שבה עומס העבודה מפעיל צינור עיבוד נתונים (pipeline) ניתן להרחבה אחרי האימון באמצעות כלי האימון המובנה של MaxText ללמידת חיזוק. בשיטת GRPO אין צורך במודל נפרד של ערך או של מבקר. במקום זאת, המערכת מעריכה קבוצה של פלטים שנוצרו לאותה הנחיה, ומשתמשת בתגמולים יחסיים כדי לבצע אופטימיזציה יעילה של מדיניות השחקן.

• מערך נתונים ותבנית: Grade School Math 8K (GSM8K), מערך נתונים איכותי שמורכב מבעיות מילוליות במתמטיקה לבית הספר היסודי, עם מגוון לשוני. עומס העבודה הזה משתמש בו כדי ליצור השקות, לחשב תגמולים ולשפר את יכולות החשיבה הרציונלית הרב-שלבית במהלך שלב ההתאמה של RL.

• תשתית ומסגרת: צומת Cloud TPU v6e-8 יחיד (ct6e-standard-8t) באמצעות MaxText‏ (maxtext[tpu-post-train]), מחסנית אימון עם ביצועים גבוהים שמופעלת על ידי JAX.

מטרות

• מגדירים מכונה וירטואלית של Cloud TPU.
• מתקינים את MaxText ואת יחסי התלות שלו.
• המרת מודל של Hugging Face לפורמט MaxText.
• מריצים עומס עבודה של אופטימיזציה יחסית של מדיניות קבוצת RL ‏ (GRPO) ב-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, מאשרים את הסכם הרישיון של המודל שמתכננים לאמן. במדריך הזה נעשה שימוש במודל llama3.1-8b-Instruct.

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

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

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

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

הגדרת הסביבה

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

export PROJECT="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export RESERVATION="YOUR_RESERVATION_NAME"
export TPU_NAME="YOUR_TPU_NAME"

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

• ‫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

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

gcloud compute instances create "${TPU_NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --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 "${TPU_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.3" --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/
   RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
   export RUN_NAME
   export CHIPS_PER_VM=8
   export NUM_BATCHES=50
   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.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 מוכן לשימוש. כי כשמפעילים מחדש את המכונה הווירטואלית, מאבדים את הגישה לתוכן של התיקייה /dev/shm. לכן, כדאי להעביר את המודל המכוונן לאחסון מתמיד או להעלות אותו ל-Hugging Face Hub.

הסרת המשאבים

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

מחיקת המכונה הווירטואלית של ה-TPU

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

gcloud compute instances delete "${TPU_NAME}" --zone="${ZONE}" --project="${PROJECT}" --quiet

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

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