TPU v5p

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

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

בקטע הזה מתוארת ארכיטקטורת המערכת הספציפית לגרסה v5p. לכל TensorCore יש ארבע יחידות Matrix Multiply‏ (MXU), יחידת וקטור ויחידה סקלרית.

ב-v5p Pod יש 8,960 צ'יפים. המשימה הגדולה ביותר שאפשר לתזמן היא משימה של 96 קוביות (6,144 צ'יפים).

בטבלה הבאה מוצגים המפרטים העיקריים של TPU v5p.

מפרט ערכים
מספר הצ'יפים בכל פוד 8960
שיא העיבוד לכל שבב (BF16) ‏(TFLOPs) 459
שיא העיבוד לכל שבב (FP8) ‏(TFLOPs) 459
קיבולת HBM לכל שבב (GiB) 95
רוחב פס של HBM לכל שבב (GBps) 2765
מספר המעבדים הווירטואליים (מכונה וירטואלית עם 4 שבבים) 208
‫RAM ‏ (GB) (מכונה וירטואלית עם 4 שבבים) 448
מספר ליבות TensorCore בכל שבב 2
מספר ליבות SparseCore בכל שבב 4
רוחב פס דו-כיווני של חיבור בין שבבים (ICI) לכל שבב (GBps) 1200
רוחב הפס של רשת מרכז הנתונים (DCN) לכל שבב (Gbps) 50
טופולוגיית Interconnect טורוס תלת-ממדי *

לכל שבב TPU v5p יש זיכרון פיזי ברוחב פס גבוה (HBM) בנפח 96GiB. בטבלה שלמעלה מפורטים 95GiB כדי להביא בחשבון את ה-HBM שמערכת ה-TPU וזמן הריצה של libtpu שומרים לאמינות ולתקורה של זמן הריצה. כלים ומסגרות כמו tpu-info, ‏ JAX ו-PyTorch מדווחים על HBM שניתן להקצאה על ידי המשתמש, ולא על הקיבולת הפיזית הנומינלית של 96GiB.

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

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

בטבלה הבאה מוצגות הצורות הנפוצות ביותר של פרוסה אחת שנתמכות ב-v5p, וגם רוב הצורות של קובייה מלאה (אבל לא כולן) שגדולות מקובייה אחת. הצורה המקסימלית של v5p היא 16x16x24 (6,144 צ'יפים, 96 קוביות).

טופולוגיה ליבות צ'יפים מארחים קוביות סוג המכונה תומך ב-Twisted?
2x2x1 8 4 1 לא רלוונטי ct5p-hightpu-4t לא רלוונטי
2x2x2 16 8 2 לא רלוונטי ct5p-hightpu-4t לא רלוונטי
2x4x4 64 32 8 לא רלוונטי ct5p-hightpu-4t לא רלוונטי
4x4x4 128 64 16 1 ct5p-hightpu-4t לא רלוונטי
4x4x8 256 128 32 2 ct5p-hightpu-4t כן
4x8x8 512 256 64 4 ct5p-hightpu-4t כן
8x8x8 1024 512 128 8 ct5p-hightpu-4t לא רלוונטי
8x8x16 2048 1024 256 16 ct5p-hightpu-4t כן
8x16x16 4096 2048 512 32 ct5p-hightpu-4t כן
16x16x16 8192 4096 1024 64 ct5p-hightpu-4t לא רלוונטי
16x16x24 12288 6144 1536 96 ct5p-hightpu-4t לא רלוונטי

אימון בפרוסה אחת נתמך עד 6,144 צ'יפים. אפשר להגדיל את המערך עד ל-18,432 צ'יפים באמצעות ריבוי-פרוסות (Multislice). למידע נוסף על Multislice, קראו את הסקירה הכללית על Cloud TPU Multislice.

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

בחלק מהצורות של פרוסות תלת-ממדיות, אפשר להשתמש בטופולוגיה של טורוס מעוות. הטופולוגיות האלה מציעות רוחב פס גבוה משמעותית לחציית הרשת. לדוגמה, טופולוגיה מפותלת של 4x4x8 מספקת עלייה תיאורטית של 70% ברוחב הפס של החצייה בהשוואה לפרוסה לא מפותלת של 4x4x8. רוחב הפס המוגדל הזה עוזר לעומסי עבודה שמשתמשים בדפוסי תקשורת גלובליים. טופולוגיות מעוותות יכולות לשפר את הביצועים של רוב המודלים, והיתרון הגדול ביותר הוא בעומסי עבודה גדולים של הטמעת TPU. תוכנת TPU תומכת בטופולוגיות מעוותות בפרוסות שבהן כל מימד שווה לגודל של המימד הקטן ביותר או כפול ממנו. לדוגמה, 4x4x8,‏ 4×8×8 או 12x12x24. טופולוגיות מעוותות נתמכות ב-TPU v4 וב-TPU v5p באמצעות Cloud TPU API.

בעומסי עבודה שמשתמשים במקביליות נתונים כאסטרטגיית המקביליות היחידה, יכול להיות שהביצועים של טופולוגיות מפותלות יהיו קצת יותר טובים. במודלים גדולים של שפה (LLM), הביצועים של טופולוגיית Twisted משתנים בהתאם לסוג המקביליות שמשתמשים בה (לדוגמה, מקביליות נתונים או מקביליות מודלים). כדי למצוא את הביצועים הטובים ביותר למודל, מאמנים את ה-LLM גם עם טופולוגיה מעוותת וגם בלי. בניסויים מסוימים במודל FSDP MaxText, נצפו שיפורים של 1-2 אחוזים בניצול ה-FLOP של המודל (MFU) כשמשתמשים בטופולוגיה מפותלת.

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

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

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

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

איור של טופולוגיית טורוס 4x2 לא מפותלת עם מעבדי TPU שמסומנים בקואורדינטות X ו-Y

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

פיתול הטופולוגיה הזו יוצר טופולוגיה סימטרית של טורוס מפותל בגודל 4×2:

איור של טופולוגיית טורוס 4x2 מעוותת עם מעבדי TPU שמסומנים בקואורדינטות X ו-Y

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

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

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

טופולוגיה עלייה תיאורטית ברוחב הפס של החיתוך לעומת טורוס לא מעוות
4×4×8 ~70%
8x8x16
12×12×24
4x8x8 ~40%
8×16×16

עמידות של Cloud TPU ICI

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

בדומה ל-Cloud TPU v4, עמידות ICI מופעלת כברירת מחדל לפרוסות v5p שהן קובייה אחת או יותר (טופולוגיה של 4x4x4).

מאפייני VM, מארח ופרוסה

מאפיין (property) הערך ב-TPU
‫# of v5p chips 4
מספר מעבדי ה-vCPU ‫208 (רק חצי מהם ניתנים לשימוש אם משתמשים בקישור NUMA כדי להימנע מפגיעה בביצועים של NUMA)
זיכרון RAM (ב-GB) ‫448 (אפשר להשתמש רק בחצי אם משתמשים בקישור NUMA כדי להימנע מפגיעה בביצועים של NUMA)
מספר צומתי NUMA 2
תפוקת NIC‏ (Gbps) 200

הקשר בין מספר ליבות ה-Tensor, הצ'יפים, המארחים/המכונות הווירטואליות והקוביות ב-Pod:

ליבות צ'יפים מארחים או מכונות וירטואליות קוביות
מארח 8 4 1
קובייה (ארון תקשורת) 128 64 16 1
הפרוסה הגדולה ביותר שנתמכת 12288 6144 1536 96
v5p full Pod 17920 8960 2240 140