מדריך הביצועים של Cloud TPU

השלב הראשון בפתרון בעיות שקשורות לביצועים של 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) עלולה להוות צוואר בקבוק. כדי לשפר את רוחב הפס של הרשת לעומסי עבודה שמוגבלים על ידי הרשת, אפשר להשתמש ב-multi-NIC כדי להגדיל את מספר ממשקי הרשת במכונות הווירטואליות של TPU. כשמשתמשים ב-multi-NIC, לכל מכונת TPU VM מוקצים ממשקי רשת נוספים, שכל אחד מהם מחובר לרשת VPC ייחודית, וכך משפרים את התפוקה הכוללת של הרשת. כרטיסי ה-NIC הנוספים צריכים להיות בטווחים של כתובות IP שאינם חופפים.

למידע נוסף על הפעלת ריבוי רשתות כשמשתמשים ב-Google Kubernetes Engine ‏ (GKE), אפשר לעיין במאמר שיפור ביצועי הרשת ללא hostNetwork ב-TPU Trillium או Ironwood ‏(TPU7x).

אופטימיזציות של קומפיילר XLA

‫XLA הוא קומפיילר ללמידת מכונה שיכול ליצור קבצים בינאריים עבור מעבדי TPU, מעבדי CPU, מעבדי GPU ופלטפורמות אחרות. מודלים שנכתבו ב-PyTorch, ב-JAX או ב-TensorFlow יוצרים גרף של פעולות ברמה גבוהה (HLO) של XLA, ש-XLA קומפל לביצוע ב-TPU. מידע נוסף על XLA זמין במאמר XLA: Optimizing Compiler for Machine Learning.

מרווח פנימי

כדי להשתמש בזיכרון TPU בצורה יעילה, צריך לארגן את הנתונים כך שאפשר יהיה לחלק אותם לחלקים בגודל 128x8. כשהנתונים לחישוב מטריצה לא ממלאים נתח שלם בגודל 128x8, מהדר XLA מוסיף ריפוד לטנסורים. יש שני חסרונות לשימוש ב-padding:

  1. טנסורים עם ריפוד לא מנצלים את ליבת ה-TPU בצורה יעילה.
  2. הוספת ריפוד מגדילה את כמות הזיכרון שנדרשת לטנסור בשבב, ויכולה להוביל לשגיאה של חוסר זיכרון.

הקומפיילר של XLA מבצע ריפוד באופן אוטומטי כשצריך, אבל אפשר לקבוע את כמות הריפוד באמצעות הכלי Memory Viewer. כדי להימנע מריפוד, אפשר לבחור מימדים של טנסור שמתאימים ל-TPU.

מידות של טנסור

כדי להשיג ביצועים מקסימליים של FLOPs, המימדים של כפל מטריצות צריכים להיות גדולים יותר מגודל ה-MXU של גרסת ה-TPU שבה אתם משתמשים. הגודל של MXU הוא ‎256 x 256 ב-v6e וב-TPU7x‏ (Ironwood), ו-‎128 x 128 בגרסאות שקודמות ל-v6e. מידע נוסף זמין במאמר ארכיטקטורת מערכת Cloud TPU.

גודל האצווה

הקומפיילר XLA מעגל כלפי מעלה את הגדלים של טנסורים שמאוחסנים בזיכרון HBM של TPU כדי לבצע חישובים בצורה יעילה יותר. הריפוד הזה מתבצע באופן שקוף ברמת החומרה ולא משפיע על התוצאות. עם זאת, במקרים מסוימים, הריפוד יכול לגרום לשימוש מוגבר בזיכרון ולזמן ביצוע ארוך יותר.

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

  1. הגודל הכולל של אצווה צריך להיות כפולה של 64 (8 לכל ליבת TPU), וגודלי המאפיינים צריכים להיות כפולה של 128.

  2. הגודל הכולל של אצווה צריך להיות כפולה של 1,024 (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).

מדריך מקיף לאופטימיזציה של הביצועים ב-TPU7x זמין במאמר אופטימיזציה של הביצועים ב-TPU7x (Ironwood).