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

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

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

למידע על הפעלת הסקה ב-Cloud TPU v6e, אפשר לעיין במדריכים הבאים:

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

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

  • יצירת חשבון ופרויקט ב- Google Cloud עם חיוב מופעל
  • התקנה של רכיבי אלפא של Google Cloud CLI
  • הפעלת Cloud TPU API
  • יצירת סוכן שירות של Cloud TPU
  • יצירה של חשבון שירות של Cloud TPU והענקת הרשאות

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

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

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

אם אתם משתמשים ב-Google Kubernetes Engine‏ (GKE) עם XPK (חבילת עיבוד מואץ), תצטרכו הרשאות נוספות במסוף Google Cloud . מידע נוסף זמין במאמר ההרשאות שנדרשות במסוף Google Cloud .

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

אפשר להקצות ולנהל TPU v6e באמצעות השיטות הבאות:

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

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

הגדרת JAX או PyTorch

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

כדי להגדיר ולהפעיל XPK עם MaxText, אפשר לעיין במאמר הפעלת MaxText בהיקף גדול באמצעות XPK .

שיפור הגדרות TCP

אם הקצאתם את ה-TPU מדגם v6e באמצעות משאבים בתור, אתם יכולים להריץ את הפקודה הבאה כדי לשפר את ביצועי הרשת על ידי הגדלת המגבלות של מאגר הנתונים הזמני לקבלת TCP.

gcloud alpha compute tpus queued-resources ssh "${QUEUED_RESOURCE_ID}" \
    --project "${PROJECT_ID}" \
    --zone "${ZONE}" \
    --node=all \
    --worker=all \
    --command='
    sudo sh -c "echo \"4096 41943040 314572800\" > /proc/sys/net/ipv4/tcp_rmem"'

שימוש ב-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 גרסה 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

    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 docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.4.35

    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 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.train 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.train 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.train 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 ומבקשים גישה.

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