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

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

‫TPU v6e, שנקרא גם Trillium, הוא הדור השישי של TPUs מבית 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 (ערכת עיבוד מואץ), אתם צריכים הרשאות נוספות במסוף 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 עם תמיכה בכרטיס רשת יחיד או עם תמיכה בכמה כרטיסי רשת באמצעות XPK. ברוב עומסי העבודה של פרוסות יחידות, כרטיס רשת יחיד מספק ביצועים מספיקים עם פחות הגדרות. לעומסי עבודה של Multislice ולעומסי עבודה שדורשים מהירויות גבוהות של הכנסת נתונים, משתמשים ב-multi-NIC.

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

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

בקטעים הבאים מוסבר איך ליצור אשכול GKE עם תמיכה בכרטיס רשת יחיד באמצעות 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, מציינים את מספר ה-slices על סמך דרישות המדרגיות של עומס העבודה. המספר הכולל של הצ'יפים באשכול הוא מספר הצ'יפים ב-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.

יצירת אשכול עם תמיכה בכרטיסי רשת מרובים באמצעות XPK

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

בקטעים הבאים מוסבר איך ליצור אשכול GKE עם תמיכה בכמה כרטיסי 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, מציינים את מספר ה-slices על סמך דרישות המדרגיות של עומס העבודה. המספר הכולל של הצ'יפים באשכול הוא מספר הצ'יפים ב-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}"

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

  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.

אימות ההגדרה של כרטיסי רשת מרובים

אחרי שיוצרים אשכול עם תמיכה בכמה כרטיסי רשת, אפשר לוודא ששני כרטיסי הרשת נמצאים בשימוש על ידי יצירת עומס עבודה של 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. בודקים את הפלט של עומס העבודה XPK ביומני המסוף Google Cloud כדי לוודא שערך ה-MTU של eth0 ו-eth1 מוגדר ל-8,896.

הגדרת 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
    • ‫LibTPU יומי 0.1.dev20240912+nightly
  • גרסאות יציבות של התוכנה:
    • ‫JAX + JAX Lib גרסה v0.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 טהורים ומיועד ל-TPU ול-GPU לאימון ולהיקש. Google Cloud 

    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) דליל.

    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

    Llama3-8b

    ‫Llama היא משפחה של מודלים גדולים של שפה (LLM) בקוד פתוח שפותחו על ידי Meta.

    דוגמה להרצת Llama3 ב-PyTorch מופיעה במודלים של torch_xla במאגר torchprime.

    MaxDiffusion

    ‫MaxDiffusion הוא אוסף של יישומי הפניה של מודלים שונים של דיפוזיה סמויה שנכתבו ב-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 ‫Two v6e-16
שלבי ההדרכה 0.069 0.073 0.13
גודל אצווה גלובלי 8 32 64
תפוקה (דוגמאות לשנייה) 115.9 438.4 492.3

אימון מודלים של Llama באמצעות PyTorch/XLA ב-Cloud TPU v6e

בקטע הזה מוסבר איך לאמן מודלים של Llama באמצעות PyTorch/XLA ב-Cloud TPU v6e באמצעות מערך הנתונים WikiText.

גישה ל-Hugging Face ולמודל Llama 3

כדי להשתמש בדוגמה הזו, צריך טוקן גישה למשתמש ב-Hugging Face. מידע על יצירת אסימוני גישה למשתמשים זמין במסמכי Hugging Face בנושא אסימוני גישה למשתמשים.

צריך גם הרשאה לגשת למודל Llama-3-8B ב-Hugging Face. כדי לקבל גישה, עוברים אל המודל Meta-Llama-3-8B ב-HuggingFace ומבקשים גישה.

יצירת מכונת VM של Cloud TPU

בדוגמה הזו, יוצרים Cloud TPU v6e עם 8 שבבים.

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

    export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-east1-d
export ACCELERATOR_TYPE=v6e-8
export RUNTIME_VERSION=v2-alpha-tpuv6e

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

    • PROJECT_ID: מזהה הפרויקט ב- Google Cloud . משתמשים בפרויקט קיים או יוצרים פרויקט חדש.
    • ‫TPU_NAME: השם של ה-TPU.
    • ‫ZONE: האזור שבו יוצרים את מכונת ה-TPU. מידע נוסף על אזורים נתמכים זמין במאמר אזורים ותחומים של TPU.
    • ‫ACCELERATOR_TYPE: סוג המאיץ מציין את הגרסה והגודל של Cloud TPU שרוצים ליצור. מידע נוסף על סוגי המאיצים הנתמכים בכל גרסת TPU זמין במאמר בנושא גרסאות TPU.
    • ‫RUNTIME_VERSION: גרסת התוכנה של Cloud TPU.

  2. יוצרים מכונת VM של Cloud TPU:

    gcloud alpha compute tpus tpu-vm create ${TPU_NAME} --version=${RUNTIME_VERSION} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --zone=${ZONE} \
   --project=${PROJECT_ID}

התקנה

מתקינים את pytorch-tpu/transformers fork של Hugging Face transformers ואת יחסי התלות. הדוגמה הזו נבדקה עם גרסאות התלות הבאות:

  • ‫torch: תואם לגרסה 2.5.0
  • ‫torch_xla[tpu]: תואם לגרסה 2.5.0
  • jax: 0.4.33
  • jaxlib: 0.4.33
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='git clone -b flash_attention https://github.com/pytorch-tpu/transformers.git
   cd transformers
   sudo pip3 install -e .
   pip3 install datasets
   pip3 install evaluate
   pip3 install scikit-learn
   pip3 install accelerate
   pip install torch~=2.6.0 torch_xla[tpu]~=2.6.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html
   pip install jax==0.4.38 jaxlib==0.4.38 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/'

הגדרת קובצי תצורה של מודלים

פקודת האימון בקטע הבא, הפעלת המודל, משתמשת בשני קובצי הגדרה בפורמט JSON כדי להגדיר פרמטרים של המודל והגדרה של Fully Sharded Data Parallel (FSDP). חלוקת FSDP מאפשרת להשתמש בגודל אצווה גדול יותר במהלך האימון על ידי חלוקת משקלי המודל בין כמה יחידות TPU. כשמאמנים מודלים קטנים יותר, יכול להיות שמספיק להשתמש במקביליות נתונים ולשכפל את המשקלים בכל מכשיר. מידע נוסף על חלוקת טנסורים למקטעים במכשירים ב-PyTorch/XLA זמין במדריך למשתמש של PyTorch/XLA SPMD.

  1. יוצרים את קובץ התצורה של פרמטר המודל. בהמשך מוצגת הגדרת הפרמטרים של מודל Llama-3-8B. למודלים אחרים, אפשר למצוא את קובץ ההגדרות ב-Hugging Face. לדוגמה, אפשר לעיין בהגדרות של Llama-2-7B.

    cat > llama-config.json << EOF
{
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 128000,
  "eos_token_id": 128001,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 14336,
  "max_position_embeddings": 8192,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 8,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 500000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.40.0.dev0",
  "use_cache": false,
  "vocab_size": 128256
}
EOF

  2. יוצרים את קובץ ההגדרות של FSDP:

    cat > fsdp-config.json << EOF
{
  "fsdp_transformer_layer_cls_to_wrap": [
    "LlamaDecoderLayer"
  ],
  "xla": true,
  "xla_fsdp_v2": true,
  "xla_fsdp_grad_ckpt": true
}
EOF

    מידע נוסף על FSDP זמין במאמר Fully Sharded Data Parallel using SPMD .

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

    gcloud alpha compute tpus tpu-vm scp llama-config.json fsdp-config.json ${TPU_NAME}:. \
   --worker=all \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

הפעלת המודל

בעזרת קובצי ההגדרות שיצרתם בקטע הקודם, מריצים את הסקריפט run_clm.py כדי לאמן את מודל Llama-3-8B על מערך הנתונים WikiText. סקריפט האימון נמשך כ-10 דקות בהרצה ב-Cloud TPU v6e-8.

  1. נכנסים אל Hugging Face ב-Cloud TPU באמצעות הפקודה הבאה:

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   pip3 install "huggingface_hub[cli]"
   huggingface-cli login --token HUGGING_FACE_TOKEN'

  2. מריצים את אימון המודל:

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   export PJRT_DEVICE=TPU
   export XLA_USE_SPMD=1
   export ENABLE_PJRT_COMPATIBILITY=true
   # Optional variables for debugging:
   export XLA_IR_DEBUG=1
   export XLA_HLO_DEBUG=1
   export PROFILE_EPOCH=0
   export PROFILE_STEP=3
   export PROFILE_DURATION_MS=100000
   # Set PROFILE_LOGDIR to a local VM path or gs://my-bucket/profile_path
   export PROFILE_LOGDIR=PROFILE_PATH
   python3 transformers/examples/pytorch/language-modeling/run_clm.py \
     --dataset_name wikitext \
     --dataset_config_name wikitext-2-raw-v1 \
     --per_device_train_batch_size 16 \
     --do_train \
     --output_dir /home/$USER/tmp/test-clm \
     --overwrite_output_dir \
     --config_name /home/$USER/llama-config.json \
     --cache_dir /home/$USER/cache \
     --tokenizer_name meta-llama/Meta-Llama-3-8B \
     --block_size 8192 \
     --optim adafactor \
     --save_strategy no \
     --logging_strategy no \
     --fsdp "full_shard" \
     --fsdp_config /home/$USER/fsdp-config.json \
     --torch_dtype bfloat16 \
     --dataloader_drop_last yes \
     --flash_attention \
     --max_steps 20'

פתרון בעיות ב-PyTorch/XLA

אם הגדרתם את המשתנים האופציונליים לניפוי באגים בקטע הקודם, הפרופיל של המודל יאוחסן במיקום שצוין על ידי המשתנה PROFILE_LOGDIR. אפשר לחלץ את הקובץ xplane.pb שמאוחסן במיקום הזה ולהשתמש ב-tensorboard כדי להציג את הפרופילים בדפדפן באמצעות ההוראות של TensorBoard.

אם הביצועים של PyTorch/XLA לא עומדים בציפיות, אפשר לעיין במדריך לפתרון בעיות, שכולל הצעות לניפוי באגים, ליצירת פרופילים ולאופטימיזציה של המודל.