אימון מודל באמצעות TPU v6e

במאמר הזה נסביר איך לאמן מודלים ב-Cloud TPU v6e (שנקרא גם Trillium), כולל הגדרת סביבה, אופטימיזציה של הביצועים ודוגמאות מעשיות לאימון באמצעות JAX ו-PyTorch.

‫TPU v6e, שנקרא גם Trillium, הוא הדור השישי של יחידות ה-TPU של Google. בכל הממשקים הטכניים, כמו ה-API והיומנים, ובכל המסמך הזה, Trillium נקרא v6e. עם 256 צ'יפים לכל Pod, הארכיטקטורה של TPU v6e דומה מאוד ל-v5e. ‫TPU v6e מותאם לאימון, לכוונון ולהצגה של מודלים מסוג Transformer, יצירת תמונות לפי טקסט ורשתות נוירונים מלאכותיות (CNN). מידע נוסף על הארכיטקטורה וההגדרות של מערכת TPU v6e זמין במאמר TPU v6e.

במדריכים הבאים מוסבר איך להריץ הסקה ב-Cloud TPU v6e:

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

לפני שמתחילים, צריך:

  • הגדרת פרויקט ב- Google Cloud ל-TPU
  • התקנת Google Cloud CLI
  • הפעלת Compute Engine API
  • יצירת חשבון שירות ב-Compute Engine והענקת הרשאות

מידע נוסף זמין במאמר הגדרת סביבת Cloud TPU.

אימות המכסה וההרשאות

מוודאים שלפרויקט יש את המכסות הבאות:

אם אתם משתמשים ב-Google Kubernetes Engine ‏ (GKE) עם XPK ‏ (Accelerated Processing Kit), אתם צריכים הרשאות נוספות במסוף Google Cloud . למידע נוסף, קראו את המאמר ההרשאות שנדרשות במסוף Google Cloud .

אפשרויות הקצאת הרשאות ידנית

אתם יכולים להקצות ולנהל את TPU v6e בדרכים הבאות:

  • ‫Compute Engine (Google Cloud CLI): אתם יכולים להשתמש בכלי Google Cloud CLI כדי לנהל את מכונות ה-TPU הווירטואליות ישירות כמשאבים של Compute Engine. אתם יכולים להקצות מכונות וירטואליות נפרדות של TPU או להשתמש בקבוצות של מופעי מכונה מנוהלים (MIG) כדי לנהל קבוצות של מכונות וירטואליות של TPU. למידע נוסף, קראו את הסקירה הכללית על יצירת TPU.
  • ‫GKE: אתם יכולים להשתמש ב-GKE כדי להקצות ולנהל יחידות TPU כמאגר של מאיצים לעומסי העבודה של למידת המכונה בקונטיינרים. מידע נוסף זמין במאמר מידע על יחידות TPU ב-GKE.
  • ‫GKE ו-XPK: ‏XPK הוא כלי שורת פקודה שמפשט את יצירת האשכולות ואת הרצת עומסי העבודה ב-GKE. הוא מיועד למומחי ML שרוצים להקצות יחידות TPU ולהריץ משימות אימון בלי צורך בידע מעמיק ב-Kubernetes. מידע נוסף זמין במאגר של XPK ב-GitHub.

הקצאת מעבדי TPU מדור v6e באמצעות GKE ו-XPK

אם אתם משתמשים ב-GKE עם יחידות TPU מסוג v6e, אתם יכולים להשתמש בפקודות Kubernetes או ב-XPK כדי להקצות יחידות TPU ולאמן או להציג מודלים. מידע נוסף על השימוש ב-GKE עם יחידות TPU זמין במאמר מידע על יחידות TPU ב-GKE.

ב-Cloud TPU v6e יש תמיכה בהגדרות של כרטיס רשת (NIC) שמאפשרות לשנות את קצב הנתונים בכמה רשתות. בקטעים הבאים מפורטות פקודות ליצירת אשכול GKE עם תמיכה ב-NIC יחיד או ב-multi-NIC באמצעות XPK. ברוב עומסי העבודה בפרוסה אחת, כרטיס רשת אחד מספק ביצועים מספיקים עם פחות הגדרות אישיות. לעומסי עבודה של ריבוי-פרוסות (Multislice) ולעומסי עבודה שדורשים מהירויות גבוהות של הטמעת נתונים, השתמשו ב-multi-NIC.

יצירת אשכול עם תמיכה ב-NIC יחיד באמצעות XPK

ברוב עומסי העבודה בפרוסה אחת, כרטיס רשת אחד מספק ביצועים מספיקים עם פחות הגדרות. לעומסי עבודה של ריבוי-פרוסות (Multislice) ולעומסי עבודה שדורשים מהירויות גבוהות של הטמעת נתונים, השתמשו ב-multi-NIC.

בקטעים הבאים מוסבר איך ליצור אשכול GKE עם תמיכה ב-NIC יחיד באמצעות XPK.

התקנת XPK והגדרת משתני סביבה

  1. התקנת XPK. פועלים לפי ההוראות במאגר של XPK ב-GitHub.

  2. מגדירים משתני סביבה לאשכול:

    export CLUSTER_NAME=XPK_CLUSTER_NAME
export ZONE=us-east1-d
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=1

    מגדירים את משתני הסביבה הבאים:

    • ‫CLUSTER_NAME: שם לאשכול.
    • ‫ZONE: האזור שבו ייווצר אשכול ה-TPU. מידע נוסף על אזורים נתמכים זמין במאמר אזורים ותת-אזורים.
    • ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
    • ‫ACCELERATOR_TYPE: סוג ה-TPU, שנקרא גם סוג המאיץ, מציין את הגרסה והגודל של ה-Cloud TPU שרוצים ליצור. לדוגמה, v6e-256. למידע נוסף על סוגי המאיצים שנתמכים בכל גרסת TPU, קראו את המאמר גרסאות TPU.
    • ‫NUM_SLICES: מספר פרוסות ה-TPU של האשכול. בכל פרוסה יש את מספר הצ'יפים שצוין ב-ACCELERATOR_TYPE. באשכול עם פרוסה אחת, מגדירים את NUM_SLICES ל-1. באשכול Multislice, מציינים את מספר הפרוסות בהתאם לדרישות של עומס העבודה בנוגע ליכולת ההתאמה לעומס. מספר הצ'יפים הכולל באשכול הוא מספר הצ'יפים ב-ACCELERATOR_TYPE כפול NUM_SLICES.

יצירת האשכול

בוחרים אחת מהאפשרויות הבאות כדי ליצור את האשכול. כדי להגיע לביצועים אופטימליים, מומלץ להשתמש ברשת בהתאמה אישית עם MTU של 8,896. מידע נוסף זמין במאמר הגדרת MTU.

רשת בהתאמה אישית

כדי ליצור רשת בהתאמה אישית עם MTU של 8,896 ולהשתמש בה באשכול:

  1. מגדירים משתני סביבה לשמות של הרשת וחומת האש:

    export NETWORK_NAME=NETWORK_NAME
export NETWORK_FW_NAME=FIREWALL_NAME

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

    • ‫NETWORK_NAME: שם לרשת.
    • ‫FIREWALL_NAME: שם לכלל של חומת האש ברשת.

  2. יוצרים רשת בהתאמה אישית עם MTU של 8,896:

    gcloud compute networks create ${NETWORK_NAME} \
    --mtu=8896 \
    --project=${PROJECT_ID} \
    --subnet-mode=auto \
    --bgp-routing-mode=regional

  3. יוצרים כלל של חומת אש שמאפשר תעבורת נתונים מסוג TCP,‏ ICMP ו-UDP ברשת:

    gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
    --network=${NETWORK_NAME} \
    --allow tcp,icmp,udp \
    --project=${PROJECT_ID}

  4. מגדירים משתנה סביבה לארגומנטים של אשכול XPK כדי להשתמש ברשת שיצרתם:

    export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"

  5. יוצרים את אשכול ה-XPK. הפקודה הבאה מקצה קיבולת על פי דרישה:

    xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=e2-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"

    כדי להשתמש בקיבולת שמורה, מחליפים את --on-demand ב---reservation=RESERVATION_NAME. כדי להשתמש במכונות וירטואליות מסוג TPU Spot, מחליפים את --on-demand ב---spot.

רשת ברירת המחדל

אם אתם לא צריכים רשת עם MTU גבוה, אתם יכולים ליצור אשכול שמשתמש ברשת ברירת המחדל של VPC. הפקודה הבאה מקצה קיבולת על פי דרישה:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=e2-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand

כדי להשתמש בקיבולת שמורה, מחליפים את --on-demand ב---reservation=RESERVATION_NAME. כדי להשתמש במכונות וירטואליות מסוג TPU Spot, מחליפים את --on-demand ב---spot.

יצירת אשכול עם תמיכה ב-multi-NIC באמצעות XPK

בעומסי עבודה של ריבוי-פרוסות (Multislice) או בעומסי עבודה אחרים שדורשים רוחב פס גבוה ברשת, כמו הטמעת נתונים, אתם יכולים להשתמש ב-multi-NIC כדי לשפר את הביצועים. כשמשתמשים ב-multi-NIC, לכל TPU VM מוקצים ממשקי רשת נוספים, שכל אחד מהם מחובר לרשת VPC ייחודית, וכך משפרים את קצב העברת הנתונים הכולל ברשת. ברוב עומסי העבודה בפרוסה אחת, כרטיס רשת אחד מספק ביצועים מספיקים עם פחות הגדרות אישיות.

בקטעים הבאים מוסבר איך ליצור אשכול GKE עם תמיכה ב-multi-NIC באמצעות XPK.

התקנת XPK והגדרת משתני סביבה

  1. התקנת XPK. פועלים לפי ההוראות במאגר של XPK ב-GitHub.

  2. מגדירים משתני סביבה לאשכול ולרשת הראשית:

    export CLUSTER_NAME=XPK_CLUSTER_NAME
export REGION=REGION
export ZONE=us-east1-d
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=2

export NETWORK_NAME_1=${CLUSTER_NAME}-mtu9k-1-${ZONE}
export SUBNET_NAME_1=${CLUSTER_NAME}-privatesubnet-1-${ZONE}
export FIREWALL_RULE_NAME_1=${CLUSTER_NAME}-privatefirewall-1-${ZONE}
export ROUTER_NAME_1=${CLUSTER_NAME}-network-1-${ZONE}
export NAT_CONFIG_1=${CLUSTER_NAME}-natconfig-1-${ZONE}

export NETWORK_NAME_2=${CLUSTER_NAME}-mtu9k-2-${ZONE}
export SUBNET_NAME_2=${CLUSTER_NAME}-privatesubnet-2-${ZONE}
export FIREWALL_RULE_NAME_2=${CLUSTER_NAME}-privatefirewall-2-${ZONE}
export ROUTER_NAME_2=${CLUSTER_NAME}-network-2-${ZONE}
export NAT_CONFIG_2=${CLUSTER_NAME}-natconfig-2-${ZONE}

    מגדירים את משתני הסביבה הבאים:

    • ‫CLUSTER_NAME: שם לאשכול.
    • ‫REGION: האזור שבו ייווצר אשכול ה-TPU.
    • ‫ZONE: האזור שבו ייווצר אשכול ה-TPU. מידע נוסף על אזורים נתמכים זמין במאמר אזורים ותת-אזורים.
    • ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
    • ‫ACCELERATOR_TYPE: סוג המאיץ מציין את הגרסה והגודל של ה-Cloud TPU שרוצים ליצור. לדוגמה, v6e-256. למידע נוסף על סוגי המאיצים שנתמכים בכל גרסת TPU, קראו את המאמר גרסאות TPU.
    • ‫NUM_SLICES: מספר פרוסות ה-TPU של האשכול. בכל פרוסה יש את מספר הצ'יפים שצוין ב-ACCELERATOR_TYPE. באשכול עם פרוסה אחת, מגדירים את NUM_SLICES ל-1. באשכול Multislice, מציינים את מספר הפרוסות בהתאם לדרישות של עומס העבודה בנוגע ליכולת ההתאמה לעומס. מספר הצ'יפים הכולל באשכול הוא מספר הצ'יפים ב-ACCELERATOR_TYPE כפול NUM_SLICES.

יצירת משאבי הרשת הראשיים

  1. יוצרים את הרשת הראשית עם יחידת העברה מקסימלית (MTU) של 8,896:

    gcloud compute networks create ${NETWORK_NAME_1} \
    --mtu=8896 \
    --bgp-routing-mode=regional \
    --subnet-mode=custom \
    --project=${PROJECT_ID}

    שימוש ברשת בהתאמה אישית עם MTU של 8,896 מספק ביצועים טובים יותר. מידע נוסף זמין במאמר הגדרת MTU.

  2. יוצרים את רשת המשנה הראשית:

    gcloud compute networks subnets create ${SUBNET_NAME_1} \
    --network=${NETWORK_NAME_1} \
    --range=10.11.0.0/18 \
    --region=${REGION} \
    --project=${PROJECT_ID}

  3. יוצרים כלל של חומת אש לרשת הראשית, שמאפשר תעבורת נתונים מסוג tcp,‏ icmp ו-udp ברשת הראשית:

    gcloud compute firewall-rules create ${FIREWALL_RULE_NAME_1} \
    --network=${NETWORK_NAME_1} \
    --allow tcp,icmp,udp \
    --project=${PROJECT_ID}

  4. יוצרים Cloud Router לרשת הראשית:

    gcloud compute routers create ${ROUTER_NAME_1} \
    --project=${PROJECT_ID} \
    --network=${NETWORK_NAME_1} \
    --region=${REGION}

  5. מגדירים NAT לרשת הראשית. הפקודה הבאה מאפשרת לתעבורת הנתונים מהאשכול להגיע לאינטרנט:

    gcloud compute routers nats create ${NAT_CONFIG_1} \
    --router=${ROUTER_NAME_1} \
    --region=${REGION} \
    --auto-allocate-nat-external-ips \
    --nat-all-subnet-ip-ranges \
    --project=${PROJECT_ID} \
    --enable-logging

יצירת משאבי הרשת המשנית

  1. יוצרים את הרשת המשנית:

    gcloud compute networks create ${NETWORK_NAME_2} --mtu=8896 \
    --bgp-routing-mode=regional \
    --subnet-mode=custom \
    --project=${PROJECT_ID}

  2. יוצרים רשת משנה לרשת המשנית:

    gcloud compute networks subnets create ${SUBNET_NAME_2} \
    --network=${NETWORK_NAME_2} \
    --range=10.10.0.0/18 \
    --region=${REGION} \
    --project=${PROJECT_ID}

  3. יוצרים כלל של חומת אש שמאפשר תעבורת נתונים ברשת החדשה:

    gcloud compute firewall-rules create ${FIREWALL_RULE_NAME_2} \
    --network=${NETWORK_NAME_2} \
    --allow tcp,icmp,udp \
    --source-ranges 10.10.0.0/18 \
    --project=${PROJECT_ID}

  4. יוצרים Cloud Router לרשת המשנית:

    gcloud compute routers create ${ROUTER_NAME_2} \
    --project=${PROJECT_ID} \
    --network=${NETWORK_NAME_2} \
    --region=${REGION}

  5. יוצרים הגדרת NAT ל-Cloud Router:

    gcloud compute routers nats create ${NAT_CONFIG_2} \
    --router=${ROUTER_NAME_2} \
    --region=${REGION} \
    --auto-allocate-nat-external-ips \
    --nat-all-subnet-ip-ranges \
    --project=${PROJECT_ID} \
    --enable-logging

יצירת האשכול

  1. מגדירים משתנה סביבה לארגומנטים של האשכול ומאגר הצמתים, כדי להשתמש ברשתות וברשתות המשנה שיצרתם:

    export CLUSTER_ARGUMENTS="--enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}"
export NODE_POOL_ARGUMENTS="--additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}"

    הארגומנטים האלה מגדירים את האשכול כך שישתמש בשתי הרשתות שיצרתם לתמיכה ב-multi-NIC.

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

    xpk cluster create \
    --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=e2-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE}  \
    --project=${PROJECT_ID} \
    --on-demand \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
    --custom-nodepool-arguments="${NODE_POOL_ARGUMENTS}" \
    --create-vertex-tensorboard

    כדי להשתמש בקיבולת שמורה, מחליפים את --on-demand ב---reservation=RESERVATION_NAME. כדי להשתמש במכונות וירטואליות מסוג TPU Spot, מחליפים את --on-demand ב---spot.

אימות ההגדרה של multi-NIC

אחרי שיוצרים אשכול עם תמיכה ב-multi-NIC, אפשר לוודא ששני כרטיסי ה-NIC נמצאים בשימוש על ידי יצירת עומס עבודה של XPK והוספת הדגל --command ifconfig.

  1. כדי להציג את הפלט של הפקודה ifconfig ביומני המסוף שלGoogle Cloud , משתמשים בפקודה הבאה: צריך לציין את הדגל --base-docker-image maxtext_base_image כדי להשתמש בקובץ האימג' הבסיסי של MaxText, כמו בדוגמה הבאה, או לציין את הדגל --docker-image ואת קובץ האימג' שבו רוצים להשתמש.

    xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --num-slices=${NUM_SLICES}  \
    --on-demand \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --command "ifconfig"

    אם רוצים להפעיל יומני ניפוי באגים או להשתמש ב- Vertex AI TensorBoard, מוסיפים לפקודה את הארגומנטים האופציונליים הבאים:

    --enable-debug-logs \
--use-vertex-tensorboard

  2. מוודאים שגם ב-eth0 וגם ב-eth1 ה-MTU מוגדר ל-8,896. כדי לעשות את זה, בודקים את הפלט של עומס העבודה של XPK ביומני המסוף של Google Cloud .

הגדרת JAX או PyTorch

במקורות המידע הבאים מוסבר איך להגדיר את JAX או את PyTorch ב-TPU, בהתאם לשיטת ההקצאה והניהול שבה אתם משתמשים:

כדי להגדיר ולהריץ את XPK עם MaxText, קראו את המאמר הרצת MaxText בקנה מידה גדול באמצעות XPK .

שימוש ב-SkyPilot

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

דוגמאות לאימון

בקטעים הבאים מופיעות דוגמאות לאימון מודלים של MaxText,‏ MaxDiffusion ו-PyTorch ב-Cloud TPU v6e.

הדוגמאות האלה נבדקו עם גרסאות התוכנה הבאות:

  • ‫Python 3.10 ואילך
  • גרסאות תוכנה ליליות:
    • ‫JAX 0.4.32.dev20240912 לילי
    • ‫Nightly LibTPU 0.1.dev20240912+nightly
  • גרסאות תוכנה יציבות:
    • ‫JAX + JAX Lib בגרסה 0.4.37

אימון של MaxText ו-MaxDiffusion ב-Cloud TPU v6e

בקטעים הבאים מוסבר על מחזור החיים של אימון המודלים MaxText ו-MaxDiffusion.

באופן כללי, השלבים העיקריים הם:

  1. יוצרים את קובץ האימג' הבסיסי של עומס העבודה.
  2. מריצים את עומס העבודה באמצעות XPK.
    1. יוצרים את פקודת האימון של עומס העבודה.
    2. פורסים את עומס העבודה.
  3. עוקבים אחרי עומס העבודה וצופים במדדים.
  4. אם אין צורך בעומס העבודה של XPK, מוחקים אותו.
  5. מוחקים את האשכול כשלא צריכים אותו יותר.

יצירת קובץ אימג' בסיסי

מתקינים את MaxText או MaxDiffusion ויוצרים את קובץ האימג' של Docker:

  1. משכפלים את המאגר שבו רוצים להשתמש ועוברים לספרייה של המאגר:

    MaxText:

    git clone https://github.com/google/maxtext.git
cd maxtext
git checkout maxtext-v0.2.0

    ‫MaxDiffusion:

    git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion && git checkout 4a8155ec0129512812b31930f0a91c6d5a141103

  2. מגדירים את Docker כך שישתמש ב-Google Cloud CLI:

    gcloud auth configure-docker

  3. יוצרים את קובץ האימג' של Docker באמצעות הפקודה הבאה או באמצעות קובץ אימג' של JAX AI. מידע נוסף על תמונות מ-AI של JAX זמין במאמר תמונות מ-AI של JAX.

    MaxText:

    bash dependencies/scripts/docker_build_dependency_image.sh MODE=stable

    ‫MaxDiffusion:

    bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_stack MODE=jax_ai_image PROJECT=${PROJECT_ID} LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest

  4. מגדירים את מזהה הפרויקט בהגדרות האישיות הפעילות של ה-CLI של gcloud:

    gcloud config set project ${PROJECT_ID}

  5. אם אתם מפעילים את עומס העבודה ממכונה שאין בה את קובץ האימג' שנוצר באופן מקומי, העלו את קובץ האימג'.

    1. מגדירים את משתנה הסביבה CLOUD_IMAGE_NAME:

      export CLOUD_IMAGE_NAME=${USER}_runner

    2. מעלים את התמונה:

      bash dependencies/scripts/docker_upload_runner.sh ${CLOUD_IMAGE_NAME}

הרצת עומס העבודה באמצעות XPK

  1. אם אתם לא משתמשים בערכי ברירת המחדל שמוגדרים על ידי MaxText או MaxDiffusion, הגדירו את משתני הסביבה הבאים:

    export BASE_OUTPUT_DIR=gs://YOUR_BUCKET
export PER_DEVICE_BATCH_SIZE=2
export NUM_STEPS=30
export MAX_TARGET_LENGTH=8192

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

    אל תפעילו עדיין את הסקריפט של המודל.

    MaxText

    ‫MaxText הוא מודל LLM בקוד פתוח, בעל ביצועים גבוהים ורמת סקיילביליות גבוהה, שכתוב ב-Python וב-JAX טהורים ומטרגט יחידות Google Cloud TPU ו-GPU לאימון ולהסקה.

    JAX_PLATFORMS=tpu,cpu \
ENABLE_PJRT_COMPATIBILITY=true \
TPU_SLICE_BUILDER_DUMP_CHIP_FORCE=true \
TPU_SLICE_BUILDER_DUMP_ICI=true && \
python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
     base_output_directory=${BASE_OUTPUT_DIR} \
     dataset_type=synthetic \
     per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
     enable_checkpointing=false \
     gcs_metrics=true \
     profiler=xplane \
     skip_first_n_steps_for_profiler=5 \
     steps=${NUM_STEPS}  # attention='dot_product'"

    Gemma2

    ‫Gemma היא משפחה של מודלי שפה גדולים (LLM) עם משקלים פתוחים שפותחו על ידי Google DeepMind, על בסיס המחקר והטכנולוגיה של Gemini.

    python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
    model_name=gemma2-27b \
    run_name=gemma2-27b-run \
    base_output_directory=${BASE_OUTPUT_DIR} \
    max_target_length=${MAX_TARGET_LENGTH} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    steps=${NUM_STEPS} \
    enable_checkpointing=false \
    use_iota_embed=true \
    gcs_metrics=true \
    dataset_type=synthetic \
    profiler=xplane \
    attention=flash

    Mixtral 8x7b

    ‫Mixtral הוא מודל AI מתקדם שפותח על ידי Mistral AI, ומשתמש בארכיטקטורה של שילוב מומחים (MoE) חסכוני (sparse).

    python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
    base_output_directory=${BASE_OUTPUT_DIR} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    model_name=mixtral-8x7b \
    steps=${NUM_STEPS} \
    max_target_length=${MAX_TARGET_LENGTH} \
    tokenizer_path=assets/tokenizer.mistral-v1 \
    attention=flash \
    dtype=bfloat16 \
    dataset_type=synthetic \
    profiler=xplane

    MaxDiffusion

    ‫MaxDiffusion הוא אוסף של הטמעות לדוגמה של מודלים שונים של דיפוזיה סמויה (latent diffusion) שכתובים ב-Python וב-JAX, ופועלים במכשירי XLA, כולל מעבדי Cloud TPU ו-GPU. ‫Stable Diffusion הוא מודל סמוי ליצירת תמונות לפי טקסט, שיוצר תמונות פוטו-ריאליסטיות מכל קלט טקסט.

    כדי להריץ את MaxDiffusion, צריך להתקין הסתעפות ספציפית של Git, כמו שמוצג בסקריפט האימון הבא.

    git clone https://github.com/google/maxdiffusion.git
&& cd maxdiffusion
&& git checkout 4a8155ec0129512812b31930f0a91c6d5a141103
&& pip install -r requirements.txt && pip install .
&& pip install huggingface_hub==0.30.2 && OUT_DIR=${BASE_OUTPUT_DIR}
&& python src/maxdiffusion/train_sdxl.py \
    src/maxdiffusion/configs/base_xl.yml \
    revision=refs/pr/95 \
    activations_dtype=bfloat16 \
    weights_dtype=bfloat16 \
    resolution=1024 \
    per_device_batch_size=1 \
    output_dir=${OUT_DIR} \
    jax_cache_dir=${OUT_DIR}/cache_dir/ \
    max_train_steps=200 \
    attention=flash \
    run_name=sdxl-ddp-v6e

  3. מייצאים את המשתנים הבאים:

    export CLUSTER_NAME=CLUSTER_NAME
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=NUM_SLICES
export YOUR_MODEL_SCRIPT=YOUR_MODEL_SCRIPT

    תיאורים של משתני סביבה

    • ‫CLUSTER_NAME: שם האשכול.
    • ‫ACCELERATOR_TYPE: סוג המאיץ מציין את הגרסה והגודל של ה-Cloud TPU שרוצים ליצור. למידע נוסף על סוגי המאיצים שנתמכים בכל גרסת TPU, ראו גרסאות TPU.
    • ‫NUM_SLICES: מספר פרוסות ה-TPU.
    • ‫YOUR_MODEL_SCRIPT: הסקריפט של המודל שצריך להריץ כפקודת אימון.

  4. מריצים את המודל באמצעות הסקריפט שיצרתם בשלב הקודם. צריך לציין את הדגל --base-docker-image כדי להשתמש בתמונת הבסיס של MaxText, או לציין את הדגל --docker-image ואת התמונה שבה רוצים להשתמש.

    אפשר להוסיף את הדגלים האופציונליים הבאים:

    xpk workload create \
  --cluster ${CLUSTER_NAME} \
  {--base-docker-image maxtext_base_image | --docker-image gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}:latest} \
  --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
  --tpu-type=${ACCELERATOR_TYPE} \
  --num-slices=${NUM_SLICES}  \
  --on-demand \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --command="${YOUR_MODEL_SCRIPT}"

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

ניפוי באגים ב-JAX ב-MaxText

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

מעקב אחרי JAX ב-MaxText באמצעות Vertex AI

כדי להשתמש ב-TensorBoard, צריך להיות לכם Google Cloud חשבון משתמש עם התפקיד aiplatform.user. מריצים את הפקודה הבאה כדי להקצות את התפקיד הזה:

gcloud projects add-iam-policy-binding your-project-id \
   --member='user:your-email' \
   --role='roles/aiplatform.user'

צפייה בנתוני סקלר ופרופיל דרך TensorBoard המנוהל של Vertex AI.

  1. מגדילים את הבקשות לניהול משאבים (CRUD) באזור שבו אתם משתמשים מ-600 ל-5,000. יכול להיות שזו לא תהיה בעיה בעומסי עבודה קטנים שמשתמשים בפחות מ-16 מכונות וירטואליות.

  2. מתקינים יחסי תלות כמו cloud-accelerator-diagnostics ל-Vertex AI:

    # xpk dependencies will install cloud-accelerator-diagnostics for Vertex AI
cd ~/xpk
pip install .

  3. יוצרים את האשכול באמצעות הדגל --create-vertex-tensorboard, כמו שמתואר במאמר יצירת Vertex AI TensorBoard. אפשר להריץ את הפקודה הזו גם על אשכולות קיימים.

  4. יוצרים את הניסוי ב-Vertex AI כשמריצים את עומס העבודה של XPK באמצעות הדגל --use-vertex-tensorboard והדגל האופציונלי --experiment-name. לרשימת השלבים המלאה, ראו יצירת ניסוי ב-Vertex AI להעלאת נתונים ל-Vertex AI TensorBoard.

היומנים כוללים קישור ל-Vertex AI TensorBoard, בדומה לקישור הבא:

View your TensorBoard at https://us-central1.tensorboard.googleusercontent.com/experiment/project_id+locations+us-central1+tensorboards+hash+experiments+name

אפשר גם למצוא את הקישור ל-Vertex AI TensorBoard במסוף Google Cloud . נכנסים לדף Vertex AI Experiments במסוף Google Cloud . בוחרים את האזור המתאים מהתפריט הנפתח.

ספריית TensorBoard נכתבת גם בקטגוריה של Cloud Storage שציינתם באמצעות ${BASE_OUTPUT_DIR}.

מחיקת עומס העבודה ב-XPK

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

מחיקת האשכול

משתמשים בפקודה xpk cluster delete כדי למחוק את האשכול:

xpk cluster delete --cluster ${CLUSTER_NAME} \
    --zone=${ZONE} --project=${PROJECT_ID}

תוצאות של בדיקת ביצועים ב-MaxDiffusion

הרצנו את סקריפט האימון של MaxDiffusion ב-v6e-4, ב-v6e-16 ובשני v6e-16. בטבלה הבאה מוצגים קצבי העברת הנתונים שנמדדו.

v6e-4 v6e-16 שני v6e-16
שלבי האימון 0.069 0.073 0.13
גודל אצווה גלובלי 8 32 64
תפוקה (דוגמאות לשנייה) 115.9 438.4 492.3

אימון מודל באמצעות PyTorch ב-Cloud TPU

אתם יכולים להשתמש ב-TorchTPU ישירות כדי לאמן את המודל או להשתמש שוב בספריית העזר TorchTitan. כדי להתחיל לעבוד עם עומסי עבודה ב-v6e, קראו את המאמר תחילת העבודה עם TorchTPU ב-TPU VM ואת המתכונים של TorchTitan Trillium. מידע נוסף על סביבת PyTorch ב-Cloud TPU זמין במאמר PyTorch ב-Cloud TPU.