TPU גרסה 4

במאמר הזה מתוארת הארכיטקטורה של Cloud TPU v4 וההגדרות הנתמכות שלו.

ארכיטקטורת המערכת

כל שבב TPU v4 מכיל שני TensorCore. לכל TensorCore יש ארבע יחידות כפל מטריצות (MXU), יחידת וקטור ויחידה סקלרית. בטבלה הבאה מוצגים המפרטים העיקריים של TPU Pod מדור 4.

מפרט ערכים
שיא החישוב לכל שבב ‫275 טרה-פלופס (bf16 או int8)
הקיבולת ורוחב הפס של HBM2 ‫32GiB,‏ 1,200GBps
הספק מינימלי/ממוצע/מקסימלי שנמדד ‫90/170/192 ואט
גודל ה-TPU Pod ‫4,096 צ'יפס
טופולוגיית Interconnect רשת תלת-ממדית
שיא עוצמת החישוב לכל Pod ‫1.1 אקסה-פלופס (bf16 או int8)
רוחב הפס של כל הפחתה לכל Pod ‫1.1 PB/s
רוחב פס של חצייה לכל Pod ‫24TB/s

התרשים הבא מדגים שבב TPU v4.

תרשים של צ'יפ TPU v4

מידע נוסף על פרטי הארכיטקטורה ומאפייני הביצועים של TPU v4 זמין במאמר TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings.

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

ל-TPU מדור 4 יש חיבור ישיר לשבבים השכנים הקרובים ביותר בתלת-ממד, וכתוצאה מכך נוצרת רשת תלת-ממדית של חיבורים. אפשר להגדיר את החיבורים כטורוס תלת-ממדי בפרוסות שבהן הטופולוגיה, AxBxC, היא 2A=B=C או 2A=2B=C, כאשר כל מימד הוא כפולה של 4. לדוגמה, 4x4x8,‏ 4x8x8 או 12x12x24. באופן כללי, הביצועים של תצורת טורוס תלת-ממדית יהיו טובים יותר מאלה של תצורת רשת תלת-ממדית. מידע נוסף זמין במאמר בנושא טופולוגיות של טורי מעוותים.

היתרונות של TPU v4 לעומת v3 מבחינת ביצועים

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

מערכת זיכרון

גישה לזיכרון לא אחיד (NUMA) היא ארכיטקטורה של זיכרון מחשב למכונות עם כמה מעבדים. לכל מעבד יש גישה ישירה לבלוק של זיכרון מהיר. מעבד והזיכרון שלו נקראים צומת NUMA. צמתי NUMA מחוברים לצמתי NUMA אחרים שנמצאים בסמיכות ישירה זה לזה. מעבד מתוך צומת NUMA אחד יכול לגשת לזיכרון בצומת NUMA אחר, אבל הגישה הזו איטית יותר מגישה לזיכרון בתוך צומת NUMA.

תוכנה שפועלת במחשב עם כמה מעבדים יכולה למקם נתונים שדרושים למעבד בתוך צומת ה-NUMA שלו, וכך להגדיל את קצב העברת הנתונים בזיכרון. מידע נוסף על NUMA זמין במאמר Non Uniform Memory Access (גישה לזיכרון לא אחיד) ב-Wikipedia.

כדי ליהנות מהיתרונות של NUMA-locality, אפשר לקשור את סקריפט האימון ל-NUMA Node 0.

כדי להפעיל את הקישור של צומת NUMA:

  1. יצירת TPU v4.

  2. מתחברים ל-TPU באמצעות SSH.

  3. מתקינים את כלי שורת הפקודה numactl. בעזרת numactl אפשר להריץ תהליכים עם מדיניות תזמון NUMA ספציפית או מדיניות מיקום זיכרון.

     $ sudo apt-get update
 $ sudo apt-get install numactl

  4. מקשרים את קוד הסקריפט ל-NUMA Node 0. מחליפים את your-training-script בנתיב לסקריפט האימון. 

     $ numactl --cpunodebind=0 python3 your-training-script

הפעלת קישור של צומת NUMA אם:

  • אם עומס העבודה שלכם תלוי מאוד בעומסי עבודה של מעבד (לדוגמה, סיווג תמונות, עומסי עבודה של המלצות) ללא קשר למסגרת.
  • אם אתם משתמשים בגרסת זמן ריצה של TPU ללא הסיומת ‎-pod (לדוגמה, tpu-vm-tf-2.10.0-v4).

הבדלים אחרים במערכת הזיכרון:

  • לשבבי TPU מדור 4 יש מרחב זיכרון HBM מאוחד של 32 גיגה-בייט בכל השבב, מה שמאפשר תיאום טוב יותר בין שני ליבות TensorCore בשבב.
  • שיפור הביצועים של HBM באמצעות תקני זיכרון ומהירויות עדכניים.
  • פרופיל הביצועים של DMA שופר עם תמיכה מובנית בצעדים של 512B בביצועים גבוהים.

TensorCores

  • מספר כפול של MXU וקצב שעון גבוה יותר שמספקים 275 TFLOPS מקסימליים.
  • רוחב פס של 2x להחלפה ולתמורה.
  • מודל גישה לזיכרון מסוג Load-store לזיכרון משותף (Cmem).
  • רוחב פס מהיר יותר לטעינת משקלים של MXU ותמיכה במצב 8 ביט כדי לאפשר גדלים קטנים יותר של אצווה וזמן אחזור משופר של הסקה.

חיבור בין שבבים

שישה קישורים לחיבור בין שבבים כדי לאפשר טופולוגיות של רשתות עם קוטר רשת קטן יותר.

אחר

  • ממשק x16 PCIE gen3 למארח (חיבור ישיר).
  • מודל אבטחה משופר.
  • יעילות אנרגטית משופרת.

הגדרות אישיות

‫TPU v4 Pod מורכב מ-4,096 שבבים שמחוברים ביניהם באמצעות קישורים מהירים שניתנים להגדרה מחדש. הגמישות של רשת TPU v4 מאפשרת לחבר את הצ'יפים בפרוסת TPU באותו גודל בכמה דרכים. כשיוצרים פרוסת TPU, מציינים את גרסת ה-TPU ואת מספר משאבי ה-TPU שנדרשים.

מגדירים טופולוגיית TPU באמצעות 3-tuple,‏ AxBxC, כאשר A<=B<=C ו-A,‏ B,‏ C הם כולם <= 4 או כולם כפולות שלמות של 4. הערכים A,‏ B ו-C הם מספרי השבבים בכל אחד משלושת המאפיינים.

בטופולוגיות שבהן 2A=B=C או 2A=2B=C יש גם וריאציות של טופולוגיות שמיועדות לתקשורת בין כל הרכיבים, למשל 4x4x8,‏ 8x8x16 ו-12x12x24. הטופולוגיות האלה נקראות טופולוגיות של טורוסים מעוותים.

באיורים הבאים מוצגות כמה טופולוגיות נפוצות של TPU v4.

טופולוגיות של הגדרות נפוצות של TPU v4

אפשר ליצור פרוסות גדולות יותר מ'קוביות' של שבבים בגודל 4x4x4.

מידע נוסף על ניהול של TPU זמין במאמר ניהול של TPU. מידע נוסף על ארכיטקטורת המערכת של Cloud TPU זמין במאמר ארכיטקטורת המערכת.

טופולוגיות של טורי מפותלים

בחלק מהצורות של פרוסות טורוס תלת-ממדיות בגרסה 4, יש אפשרות להשתמש בטופולוגיה שנקראת טורוס מעוות. לדוגמה, אפשר לסדר שני קוביות v4 כפרוסה בגודל 4x4x8 או כ-4x4x8_twisted. טופולוגיות מעוותות מציעות רוחב פס של חצייה גבוה משמעותית. לדוגמה, פרוסת רשת עם טופולוגיה של 4x4x8_twisted מספקת עלייה תיאורטית של 70% ברוחב הפס של החיתוך הדו-צדדי בהשוואה לפרוסת רשת של 4x4x8 שאינה twisted. רוחב פס מוגדל של חציית הגרף שימושי לעומסי עבודה שמשתמשים בדפוסי תקשורת גלובליים. טופולוגיות מעוותות יכולות לשפר את הביצועים ברוב המודלים, והן מועילות במיוחד לעומסי עבודה גדולים של הטמעה ב-TPU.

עבור עומסי עבודה שמשתמשים במקביליות נתונים כאסטרטגיית המקביליות היחידה, יכול להיות ששימוש בטופולוגיות מסוג twisted יניב ביצועים טובים יותר. במודלי שפה גדולים, הביצועים באמצעות טופולוגיה מעוותת יכולים להשתנות בהתאם לסוג ההקבלה (DP,‏ MP וכו'). השיטה המומלצת היא לאמן את מודל ה-LLM עם טופולוגיה מעוותת וגם בלי, כדי לקבוע איזו מהן מספקת את הביצועים הטובים ביותר למודל. במסגרת כמה ניסויים במודל FSDP MaxText, נרשמו שיפורים של מיליון עד שני מיליון יחידות MFU באמצעות טופולוגיה מעוותת.

היתרון העיקרי של טופולוגיות מעוותות הוא שהן הופכות טופולוגיית טורוס לא סימטרית (לדוגמה, 4×4×8) לטופולוגיה סימטרית שקרובה אליה. לטופולוגיה הסימטרית יש הרבה יתרונות:

  • איזון עומסים משופר
  • רוחב פס גבוה יותר של חצייה
  • מסלולים קצרים יותר של חבילות

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

תוכנת ה-TPU תומכת בטורוסים מעוותים בפרוסות שבהן הגודל של כל מימד שווה לגודל של המימד הקטן ביותר או כפול ממנו. לדוגמה, 4x4x8,‏ 4×8×8 או 12x12x24.

לדוגמה, נניח שיש טופולוגיית טורוס 4x2 עם יחידות TPU שמסומנות בקואורדינטות (X,Y) שלהן בפרוסת ה-TPU:

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

שרטוט

אם נשנה את הטופולוגיה הזו, נקבל טופולוגיה סימטרית לחלוטין של טורוס מעוות בגודל 4x2:

שרטוט

ההבדל היחיד בין הדיאגרמה הזו לדיאגרמה הקודמת הוא הקצוות המעוגלים של Y. במקום להתחבר ל-TPU אחר עם אותו קואורדינטה X, הם הוסטו כדי להתחבר ל-TPU עם קואורדינטה X+2 mod 4.

אותו רעיון תקף גם לגדלים שונים של מאפיינים ולמספרים שונים של מאפיינים. הרשת שמתקבלת היא סימטרית, כל עוד כל מאפיין שווה למאפיין הקטן ביותר או גדול ממנו פי שניים.

בטבלה הבאה מוצגות טופולוגיות מעוותות נתמכות, ועלייה תיאורטית ברוחב הפס של החצייה בהשוואה לטופולוגיות לא מעוותות.

Twisted Topology הגדלה תיאורטית של רוחב הפס של חציית
הטורוס לעומת טורוס לא מסובב
4×4×8_twisted ‎~70%
8x8x16_twisted
12×12×24_twisted
4×8×8_twisted ‫~40%
8×16×16_twisted

וריאציות של טופולוגיות TPU v4

אפשר לסדר טופולוגיות מסוימות שמכילות את אותו מספר של צ'יפים בדרכים שונות. לדוגמה, אפשר להגדיר פרוסת TPU עם 512 שבבים (1, 024 TensorCores) באמצעות הטופולוגיות הבאות: 4x4x32, ‏ 4x8x16 או 8x8x8. פרוסת TPU עם 2,048 שבבים (4,096 TensorCores) מציעה עוד יותר אפשרויות טופולוגיה: 4x4x128,‏ 4x8x64,‏ 4x16x32 ו-8x16x16.

טופולוגיית ברירת המחדל שמשויכת למספר נתון של שבבים היא זו שהכי דומה לקובייה. הצורה הזו היא כנראה הבחירה הטובה ביותר לאימון מקביל של למידת מכונה (ML). טופולוגיות אחרות יכולות להיות שימושיות לעומסי עבודה עם כמה סוגים של מקביליות (לדוגמה, מקביליות של מודלים ונתונים, או חלוקה מרחבית של סימולציה). עומסי העבודה האלה פועלים בצורה הטובה ביותר אם הטופולוגיה תואמת למידת המקביליות שבה נעשה שימוש. לדוגמה, אם מציבים מקביליות של מודל 4-way במאפיין X ומקביליות של נתונים 256-way במאפיינים Y ו-Z, מתקבלת טופולוגיה של 4x16x16.

מודלים עם כמה מימדים של מקביליות מניבים את הביצועים הכי טובים כשמימדי המקביליות שלהם ממופים למימדי טופולוגיית ה-TPU. בדרך כלל מדובר במודלים גדולים של שפה (LLM) שמקבילים לנתונים ולמודלים. לדוגמה, עבור פרוסת TPU v4 עם טופולוגיה 8x16x16, המידות של טופולוגיית ה-TPU הן 8, ‏ 16 ו-16. יעיל יותר להשתמש במקביליות של מודל 8-way או 16-way (שממופה לאחד מהמאפיינים של טופולוגיית TPU פיזית). מקביליות מודל של 4 כיוונים לא תהיה אופטימלית בטופולוגיה הזו, כי היא לא תואמת לאף אחד מהממדים של טופולוגיית ה-TPU, אבל היא תהיה אופטימלית בטופולוגיה של 4x16x32 באותו מספר של שבבים.

התצורות של TPU v4 מחולקות לשתי קבוצות: תצורות עם טופולוגיות קטנות מ-64 שבבים (טופולוגיות קטנות) ותצורות עם טופולוגיות גדולות מ-64 שבבים (טופולוגיות גדולות).

טופולוגיות קטנות בגרסה 4

‫Cloud TPU תומך בפרוסות TPU v4 הבאות שקטנות מ-64 שבבים, קובייה בגודל 4x4x4. אפשר ליצור את הטופולוגיות הקטנות האלה של v4 באמצעות השם שלהן שמבוסס על TensorCore (לדוגמה, v4-32) או באמצעות הטופולוגיה שלהן (לדוגמה, 2x2x4):

שם (מבוסס על מספר ליבות Tensor) מספר הצ'יפים טופולוגיה
v4-8 4 2x2x1
v4-16 8 2x2x2
v4-32 16 ‫2x2x4
v4-64 32 2x4x4

טופולוגיות גדולות מגרסה 4

פרוסות TPU v4 זמינות במרווחים של 64 שבבים, עם צורות שהן כפולות של 4 בכל שלושת הממדים. המידות צריכות להיות בסדר עולה. בטבלה הבאה מוצגות כמה דוגמאות. חלק מהטופולוגיות האלה הן טופולוגיות 'מותאמות אישית' שאפשר ליצור רק באמצעות הדגלים --type ו---topology, כי יש יותר מדרך אחת לסדר את הצ'יפים.

שם (מבוסס על מספר ליבות Tensor) מספר הצ'יפים טופולוגיה
v4-128 64 4x4x4
v4-256 128 4x4x8
v4-512 256 4x8x8
טופולוגיה בהתאמה אישית: חובה להשתמש בדגלים --type ו---topology 256 4x4x16
v4-1024 512 8x8x8
v4-1536 768 8x8x12
v4-2048 1024 8x8x16
טופולוגיה בהתאמה אישית: חובה להשתמש בדגלים --type ו---topology 1024 4x16x16
v4-4096 2048 ‫8x16x16
… … …