מדריך לשיפור הביצועים של Cloud TPU
השלב הראשון בפתרון בעיות שקשורות לביצועים של TPU הוא ליצור פרופיל של המודל. מידע נוסף על יצירת פרופיל ביצועים זמין במאמר יצירת פרופיל של המודל ב-Cloud TPU.
ביצועי מודל TPU
בקטע הזה מתוארות בעיות כלליות שיכולות להפחית את הביצועים של המודל, ומוסבר איך אפשר לטפל בהן.
מודלים שמוגבלים לקלט
יחידות TPU מבצעות חישובים מהר מאוד. כדי לוודא ש-TPU לא נמצא בלי פעילות, חשוב לוודא שמתבצע טעינה של נתונים ל-TPU באופן רציף. האופן שבו זה נעשה תלוי באופן הטעינה והעיבוד המקדים של מערך הנתונים.
לדוגמה, אפשר לקרוא קובצי נתונים במקביל באמצעות tf.data.TFRecordset() והפרמטר num_parallel_reads.
גודל אצווה קטן בגלל חלוקה
זמן הריצה של ה-TPU מפצל את האצווה בין כל 8 הליבות של מכשיר TPU (לדוגמה v2-8 או v3-8). אם מציינים גודל אצווה גלובלי של 128, כל ליבה מקבלת גודל אצווה של 16 (128 חלקי 8).
כדי להשתמש בזיכרון בצורה אופטימלית, צריך להשתמש בגודל האצווה הגדול ביותר שמתאים לזיכרון של TPU. כל ליבת TPU משתמשת ברישום וקטורי דו-ממדי בגודל 8X128 לעיבוד של מכפלות מטריצות. באופן כללי, גודל האצווה צריך להתחלק ב-8 או ב-128 ללא שארית.
שיפור ניהול הזיכרון
אתם יכולים להשתמש במשתני סביבה שקשורים לזיכרון כדי לכוונן התנהגויות של זמן ריצה ברמה נמוכה.
TPU_PREMAPPED_BUFFER_SIZE
TPU_PREMAPPED_BUFFER_SIZE מגדיר את הגודל של מאגר הזיכרון של המארח (בבייט) שממופה מראש ומוצמד לשימוש על ידי זמן הריצה של TPU להעברות נתונים (לדוגמה, DMA). ערך ברירת המחדל הוא 4,294,967,296 בייט. הערך חייב להיות כפולה של 2^12 (4KB = 4 * 1024 Bytes = 4096 = 2^12).
הדוגמאות הבאות הן ערכים תקינים של TPU_PRE_MAPPED_BUFFER_SIZE.
17179869184 = 2^34 = 2^22 * 2^12 (2^22 4KB pages will be premapped).
40000000000 = 5^10 * 2^12 = (5^10 4KB pages will be premapped).
הגדלת הגודל הזה עשויה לשפר את הביצועים של העברת הנתונים בין המארח לבין מכשיר ה-TPU, במיוחד בעומסי עבודה עם טנסורים גדולים או עם תקשורת תכופה בין המארח למכשיר. עם זאת, היא גם מגדילה את כמות הזיכרון המארח שמוצמד, ומקטינה את הזיכרון שזמין לתהליכים אחרים.
פתרון בעיות שקשורות לזיכרונות
אם אזור החיץ שמופה מראש לא גדול מספיק כדי להקצות זיכרון במהלך זמן הריצה של התוכנית, עומס העבודה ייכשל ותוחזר שגיאה RESOURCE_EXHAUSTED דומה ל:
"Allocating buffer from premmaped region failed with: RESOURCE_EXHAUSTED:
Attempting to allocate allocation_size. זה לא היה אפשרי. יש available_size בחינם".
אם המאגר גדול מדי, יכול להיות שייקח הרבה יותר זמן לאתחל את ה-TPU (יכול להיות יותר מ-15 שניות), ולכן נראה שה-TPU נתקע.
כדי לאבחן את הבעיה, בודקים את יומני זמן הריצה של TPU. ביומנים האלה מפורטות הפעולות שמבוצעות, כולל מיפוי מראש של מאגרי נתונים זמניים. אפשר למצוא את היומנים בנתיב /tmp/tpu_logs/tpu_driver.INFO או להדפיס אותם ישירות במסוף על ידי הגדרת משתנה הסביבה TPU_STDERR_LOG_LEVEL=0. ההגדרה הזו תיצור פלט שדומה לזה:
I0604 12:45:24.926233 62136 tpu_hal.cc:214] Starting premapped memory manager initialization...
I0604 12:45:29.411218 62136 system.cc:1059] tpu::System initialized, current host id: 0, logical device ids: 0
I0604 12:45:29.411244 61600 tfrt_tpu_system_state.cc:216] CreateTpuSystemState: TPU initialization is successful and it took 5.583190661s
I0604 12:45:29.411267 61600 tfrt_tpu_system_state.cc:220] CreateTpuSystemState: using TPU host premapped buffer of size: 4294967296
בפלט הזה יופיע כמה זמן לקח לאתחל את ה-TPU ומה הגודל של מאגר הנתונים הממופה מראש.
הגדרת גודל שטח אחסון זמני
אם המאגר הממופה מראש קטן מדי או גדול מדי, אפשר להגדיר את גודל שטח האחסון הזמני באופן ידני באמצעות משתני הסביבה הבאים.
TPU_PREMAPPED_BUFFER_SIZE: הגדרת הגודל הכולל (בבייטים) של אזור המאגר הממופה מראש.
TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES: הגדרת הגודל המקסימלי של מאגר יחיד שאפשר להקצות מהאזור שמופה מראש.
לדוגמה, אתם יכולים:
export TPU_PREMAPPED_BUFFER_SIZE=4294967296
כדי להגדיר את שטח האחסון הזמני:
export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES
כדי להפעיל אותה.
בייצוא הזה, הגודל מוגדר כברירת המחדל.
משנים את הערך של TPU_PREMAPPED_BUFFER_SIZE אם חושדים שהעברת נתונים ממכשיר מארח היא צוואר בקבוק. כדי למצוא את האיזון האופטימלי, כדאי לעקוב אחרי השימוש בזיכרון של המארח וביצועי המודל. בדרך כלל ערך ברירת המחדל מספיק לרוב תרחישי השימוש.
הגדרת tcmalloc
ספריית tcmalloc משמשת כברירת מחדל במכונות וירטואליות של Cloud TPU כדי לשפר את הביצועים של מודלים עם הקצאות זיכרון גדולות ותכופות. ההגדרה הזו מתבצעת באמצעות משתנה הסביבה LD_PRELOAD.
עם זאת, בעומסי עבודה מסוימים (לדוגמה, DLRM עם הקצאות גדולות מאוד של טבלת הטמעה), tcmalloc יכול לגרום להאטה. במקרים כאלה, אפשר לחזור לפונקציה הרגילה malloc על ידי ביטול ההגדרה של המשתנה LD_PRELOAD בסשן של מעטפת הפקודות לפני שמריצים את סקריפט האימון:
unset LD_PRELOAD
אופטימיזציות של ביצועי הרשת
בקטעים הבאים מוסבר איך לבצע אופטימיזציה של ביצועי הרשת על ידי הגדרת יחידת השידור המקסימלית (MTU) ושימוש בכרטיסי רשת מרובים (multi-NIC) בסביבות ריבוי-פרוסות (Multislice).
הגדרת MTU
כדי לקבל את ביצועי הרשת הטובים ביותר, מומלץ להשתמש ברשת עם MTU (יחידת שידור מקסימלית) של 8,896.
כברירת מחדל, ענן וירטואלי פרטי (VPC) מספק רק MTU של 1,460 בייט, מה שמוביל לביצועי רשת לא אופטימליים. אפשר להגדיר את ה-MTU של רשת VPC לכל ערך בין 1,300 בייט ל-8,896 בייט (כולל). גודלי MTU מותאמים אישית נפוצים הם 1,500 בייט (Ethernet רגיל) או 8,896 בייט (הגודל המקסימלי האפשרי). מידע נוסף זמין במאמר בנושא גדלי MTU תקינים ברשת VPC.
מידע נוסף על שינוי הגדרת ה-MTU ברשת קיימת או ברשת ברירת מחדל זמין במאמר שינוי הגדרת ה-MTU ברשת VPC.
שימוש באפשרות multi-NIC ל-Multislice
כשמאמנים מודלים גדולים בסביבות Multislice שמורכבות מאלפי שבבי TPU, תקשורת בין פרוסות ברשת מרכז הנתונים (DCN) עלולה להוות צוואר בקבוק. כדי לשפר את רוחב הפס ברשת לעומסי עבודה שמוגבלים על ידי הרשת, אפשר להשתמש בכמה כרטיסי NIC כדי להגדיל את מספר ממשקי הרשת במכונות הווירטואליות של TPU. כשמשתמשים בכמה ממשקי רשת, לכל מכונת TPU VM מוקצים ממשקי רשת נוספים, שכל אחד מהם מחובר לרשת VPC ייחודית, וכך משפרים את התפוקה הכוללת של הרשת. כרטיסי ה-NIC הנוספים צריכים להיות בטווחים של כתובות IP שאינם חופפים.
למידע נוסף על הפעלת ריבוי רשתות כשמשתמשים ב-Google Kubernetes Engine (GKE), אפשר לעיין במאמר שיפור ביצועי הרשת ללא hostNetwork ב-TPU Trillium או ב-Ironwood (TPU7x). דוגמה לשימוש בכרטיסי רשת מרובים עם XPK מופיעה במאמר יצירת אשכול עם תמיכה בכרטיסי רשת מרובים באמצעות XPK.
אופטימיזציות של קומפיילר XLA
XLA הוא קומפיילר ללמידת מכונה שיכול ליצור קבצים בינאריים עבור מעבדי TPU, מעבדי CPU, מעבדי GPU ופלטפורמות אחרות. XLA הוא חלק מבסיס הקוד הסטנדרטי של TensorFlow, אבל אפשר להשתמש בו גם במודלים של PyTorch ו-JAX. מודלים של Cloud TPU מתורגמים לגרף XLA, ש-XLA מהדר אותו לקובץ הפעלה של TPU. מידע נוסף על XLA זמין במאמר XLA: Optimizing Compiler for Machine Learning.
מרווח פנימי
כדי להשתמש בזיכרון TPU בצורה יעילה, צריך לבנות את הנתונים כך שאפשר יהיה לחלק אותם לחלקים בגודל 128x8. אם הנתונים לחישוב מטריצה לא ממלאים נתח שלם בגודל 128x8, מהדר XLA מוסיף ריפוד לטנסורים. יש שני חסרונות לשימוש במרווח פנימי:
- טנסורים עם ריפוד לא מנצלים את ליבת ה-TPU בצורה יעילה.
- הוספת Padding מגדילה את כמות הזיכרון שנדרשת לטנסור בשבב, ויכולה לגרום לשגיאה של חוסר זיכרון.
הוספת הריפוד מתבצעת אוטומטית על ידי מהדר XLA כשצריך, אבל אפשר לקבוע את כמות הריפוד באמצעות הכלי לבדיקת הזיכרון. כדי להימנע מריפוד, אפשר לבחור ממדי טנסור שמתאימים ל-TPU.
מידות של טנסור
כדי להשיג ביצועים מקסימליים של FLOPs, המימדים של הכפל מטריצות צריכים להיות גדולים יותר מגודל ה-MXU של גרסת ה-TPU שבה אתם משתמשים. הגודל של MXU הוא 256 x 256 בגרסה v6e ו-128 x 128 בגרסאות שקודמות ל-v6e. מידע נוסף זמין במאמר ארכיטקטורת מערכת Cloud TPU.
גודל אצווה
הקומפיילר של XLA מעגל כלפי מעלה את הגדלים של טנסורים שמאוחסנים בזיכרון HBM של TPU כדי לבצע חישובים בצורה יעילה יותר. הריפוד הזה מתבצע באופן שקוף ברמת החומרה, והוא לא משפיע על התוצאות. עם זאת, במקרים מסוימים, הוספת הריפוד עלולה להגדיל באופן משמעותי את השימוש בזיכרון ואת זמן הביצוע.
זמן הריצה של TPU מסדר טנסורים בזיכרון כדי למקסם את יעילות החישוב ולמזער את הריפוד. כדי לצמצם את התקורה של הזיכרון ולמקסם את יעילות החישוב, אחד מהתנאים הבאים צריך להתקיים:
הגודל הכולל של אצווה צריך להיות כפולה של 64 (8 לכל ליבת TPU), והגדלים של ממדי התכונות צריכים להיות כפולה של 128.
הגודל הכולל של אצווה צריך להיות כפולה של 1024 (128 לכל ליבת TPU), והגדלים של מאפייני המימד צריכים להיות כפולה של 8.
שימוש בגודל אצווה של 1,024 ובמאפיינים של מימדים שהם כפולה של 128 מניב את היעילות הכי טובה, אבל יכול להיות שזה לא אפשרי בכל המודלים.
תצוגה משולבת
Fusion היא טכניקה כללית שמהדר XLA משתמש בה כדי לבצע אופטימיזציה של תוכניות. פעולה משולבת היא שילוב של כמה פעולות שמרכיבות אותה, שצריך לבצע אותן יחד.
לדוגמה, נניח את סדר הפעולות הבא:
tmp = tf.add(x, y)
result = tf.multiply(tmp, z)
הקוד הזה מקביל בערך לקוד המדומה הבא:
for (i = 0; i < element_count; i++) {
tmp[i] = x[i] + y[i];
}
for (i = 0; i < element_count; i++) {
result[i] = tmp[i] * z[i];
}
במיזוג, הגישה למערך מתבצעת בו-זמנית:
for (i = 0; i < element_count; i++) {
result[i] = (x[i] + y[i]) * z[i];
}
בדוגמה הזו, מספר הגיחות הלוך ושוב לזיכרון מצטמצם, ו-XLA לא צריך להקצות מקום ל-tmp.
Fusion הוא אופטימיזציה קריטית שמועילה ל-Cloud TPU בכמה דרכים:
- הוא מפחית את העברות הזיכרון כי הוא מבטל את הצורך לאחסן תוצאות ביניים בזיכרון הראשי, שהוא איטי.
- היא מאפשרת ניצול טוב יותר של יחידות חומרה, שאחרת לא היה בהן שימוש.
- הוא יכול להקטין את ניצול הזיכרון של מודל, כי פחות מאגרי נתונים צריכים להיות פעילים בו-זמנית.
שידור
שידור מתרחש באופן מרומז כשמשלבים שני טנסורים עם צורות שונות, אבל תואמות.
לדוגמה, ב-tf.add(vector, matrix) הווקטור צריך להיות משודר לצורה של המטריצה. התוצאה של הפעולה היא באותה צורה כמו המטריצה. פרטים נוספים זמינים במדריך בנושא שידור מערכים.
לרוב אפשר למזג שידורים עם הצרכנים שלהם, אבל הפעלה של שידור בכוח עלולה לגרום לביצועים ירודים ולעלייה בשימוש בזיכרון.
בדוגמה הבאה, אי אפשר למזג את השידור שמשתמע מהוספה של וקטור ומטריצה עם argmax, ולכן מתקבל שידור מוחשי:
`tf.argmax(tf.add(vector, zero_matrix), axis=0)`
המלצות לשיפור הביצועים של ארכיטקטורת Ironwood עם שני צ'יפלטים
מודל התכנות של Ironwood מאפשר גישה לשני מכשירי TPU במקום לליבה לוגית אחת (שנקראת גם MegaCore) שבה נעשה שימוש בדורות הקודמים (TPU v4 ו-v5p). השינוי הזה משפר את היעילות הכלכלית של ייצור השבב. העיצוב החדש מייצג שינוי ארכיטקטוני, אבל הוא מאפשר לכם לעשות שימוש חוזר במודלים קיימים של תוכנה עם שינויים מינימליים.
כדי להשיג את הביצועים הטובים ביותר עם ארכיטקטורת ה-dual-chiplet, מומלץ להשתמש בגישות הבאות:
שימוש במקביליות טנסורית בין שבבים: הממשק D2D עם רוחב פס גבוה מיועד למקביליות טנסורית יעילה. מומלץ לפצל טנסורים בין שני המכשירים שבשבב.
שימוש במערכים היררכיים: כדי למקסם את יעילות התקשורת, כדאי לנצל את ההיררכיה של הרשת בשתי רמות: קישור D2D מהיר במיוחד בין שבבים על אותו שבב, וקישורי ICI מהירים בתוך פרוסת שבב. כשמשתמשים במקביליות אוטומטית עם SPMD (תוכנית אחת, נתונים מרובים), קומפיילר XLA מטפל בזה בשבילכם על ידי יצירה אוטומטית של פעולות היררכיות קולקטיביות. כשמבצעים חלוקה ידנית של המודל, צריך גם לעצב את דפוסי התקשורת בהתאם להיררכיה הזו. לתת עדיפות לתקשורת בין שני המכשירים באותו שבב לפני תקשורת עם מכשירים בשבבים אחרים.
חפיפה בין תקשורת לחישוב: כדי למקסם את השימוש בחומרה, כדאי להפחית עומס של פעולות תקשורת קולקטיביות, כמו all-reduce, ל-SparseCores. הפעולות האלה, שלא קשורות ליחידה של כפל מטריצות (MXU), יכולות להתבצע במקביל ב-SparseCores בזמן שה-TensorCores ממשיכות בחישוב שלהן. הטכניקה הזו יכולה לשחזר חלק מהיתרונות בביצועים שהיו טמונים בפעולות הממוזגות בארכיטקטורת MegaCore הקודמת.
הפחתת עומס ל-SparseCore לצורך הטמעה: בעיצוב של שני שבבים, אפשר לחלק את טבלאות ההטמעה בין ה-HBM של שני השבבים. כדי למנוע ירידה בביצועים בגלל חוסר בזיכרון משותף, כדאי להעביר את פעולות איסוף ההטמעה אל SparseCore. באסטרטגיה הזו נעשה שימוש בחיבור D2D מהיר במיוחד כדי להעביר ביעילות וקטורים של הטמעה בין הצ'יפלטים. מידע נוסף על SparseCore ועל מודלים של הטמעה זמין במאמר A deep dive into SparseCore for Large Embedding Models (LEM).
מידע נוסף על ארכיטקטורת Ironwood ב-TPU7x זמין במאמר בנושא TPU7x (Ironwood).