ארכיטקטורת TPU
יחידות לעיבוד טנסורים (TPU) הן מעגלים משולבים לאפליקציות ספציפיות (ASIC) שפותחו על ידי Google כדי להאיץ עומסי עבודה של למידת מכונה. Cloud TPU הוא Google Cloud שירות שמאפשר להשתמש ביחידות TPU כמשאב ניתן להרחבה.
יחידות TPU מיועדות לביצוע מהיר של פעולות מטריצה, ולכן הן אידיאליות לעומסי עבודה של למידת מכונה. אתם יכולים להריץ עומסי עבודה של למידת מכונה ב-TPU באמצעות מסגרות כמו PyTorch ו-JAX.
איך יחידות TPU פועלות?
כדי להבין איך יחידות TPU פועלות, כדאי להבין איך מאיצים אחרים מתמודדים עם האתגרים החישוביים של אימון מודלים של ML.
איך מעבד עובד
יחידת עיבוד מרכזית (CPU) היא מעבד למטרות כלליות שמבוסס על ארכיטקטורת פון נוימן. כלומר, מעבד (CPU) פועל עם תוכנה וזיכרון באופן הבא:
היתרון הכי גדול של מעבדים הוא הגמישות שלהם. אפשר לטעון כל סוג של תוכנה במעבד למגוון רחב של סוגי אפליקציות. לדוגמה, אפשר להשתמש במעבד כדי לערוך מסמכים במחשב, לשלוט במנועי רקטות, לבצע עסקאות בנקאיות או לסווג תמונות באמצעות רשת נוירונים.
מעבד טוען ערכים מהזיכרון, מבצע חישוב על הערכים ומאחסן את התוצאה בחזרה בזיכרון לכל חישוב. הגישה לזיכרון איטית בהשוואה למהירות החישוב, ויכולה להגביל את התפוקה הכוללת של המעבדים. הבעיה הזו נקראת לעיתים קרובות צוואר הבקבוק של פון נוימן.
איך GPU עובד
כדי להשיג תפוקה גבוהה יותר, מעבדי GPU מכילים אלפי יחידות לוגיות אריתמטיות (ALU) במעבד יחיד. מעבד GPU מודרני מכיל בדרך כלל בין 2,500 ל-5,000 יחידות ALU. מספר המעבדים הגדול מאפשר לבצע אלפי פעולות כפל וחיבור בו-זמנית.
ארכיטקטורת ה-GPU הזו מתאימה לאפליקציות עם מקביליות מסיבית, כמו פעולות מטריצה ברשת נוירונים. למעשה, בעומס עבודה אופייני של אימון ללמידה עמוקה, ה-GPU יכול לספק תפוקה גבוהה יותר מסדר גודל של CPU.
אבל ה-GPU הוא עדיין מעבד למטרות כלליות שצריך לתמוך בהרבה אפליקציות ותוכנות שונות. לכן, בעיות ב-GPU דומות לבעיות ב-CPU. לכל חישוב באלפי יחידות ALU, יחידת GPU צריכה לגשת לאוגרים או לזיכרון משותף כדי לקרוא אופרנדים ולאחסן את תוצאות החישוב הביניים.
איך TPU עובד
Google עיצבה את Cloud TPU כמעבד מטריצות שמתמחה בעומסי עבודה של רשתות נוירונים. אי אפשר להריץ מעבדי תמלילים, לשלוט במנועי רקטות או לבצע עסקאות בנקאיות במעבדי TPU, אבל הם יכולים לבצע פעולות מטריצה מסיביות שמשמשות ברשתות עצביות במהירות גבוהה.
המשימה העיקרית של יחידות TPU היא עיבוד מטריצות, שהוא שילוב של פעולות כפל וצבירה. יחידות TPU מכילות אלפי יחידות של מכפילים מצטברים שמחוברות ישירות זו לזו ויוצרות מטריצה פיזית גדולה. זה נקרא ארכיטקטורה של מערך סיסטולי. Cloud TPU v3, מכיל שני מערכים סיסטוליים של 128x128 יחידות ALU, במעבד יחיד.
המארח של ה-TPU מעביר נתונים בסטרימינג לתור של נתונים נכנסים. ה-TPU טוען נתונים מתור ה-infeed ומאחסן אותם בזיכרון עם רוחב פס גבוה (HBM). כשהחישוב מסתיים, ה-TPU טוען את התוצאות לתור של ה-outfeed. לאחר מכן, המארח של ה-TPU קורא את התוצאות מתור הנתונים היוצאים ומאחסן אותן בזיכרון של המארח.
כדי לבצע את פעולות המטריצה, ה-TPU טוען את הפרמטרים מ-HBM אל יחידת הכפל של המטריצה (MXU).
לאחר מכן, ה-TPU טוען נתונים מ-HBM. בכל פעם שמבצעים פעולת כפל, התוצאה מועברת למכפיל-צובר הבא. הפלט הוא סכום כל תוצאות הכפל בין הנתונים לפרמטרים. אין צורך בגישה לזיכרון במהלך תהליך הכפל של המטריצה.
כתוצאה מכך, יחידות TPU יכולות להשיג תפוקת מחשוב גבוהה בחישובים של רשתות נוירונים.
ארכיטקטורת מערכת TPU
בקטעים הבאים מוסברים המושגים המרכזיים של מערכת TPU. מידע נוסף על מונחים נפוצים בתחום למידת המכונה זמין במילון המונחים של למידת המכונה.
אם אתם חדשים ב-Cloud TPU, כדאי לעיין בדף הבית של התיעוד בנושא TPU.
שבב TPU
שבב TPU מכיל ליבת Tensor אחת או יותר. מספר ליבות ה-Tensor תלוי בגרסה של שבב ה-TPU. כל TensorCore מורכב מיחידה אחת או יותר של כפל מטריצות (MXU), יחידה וקטורית ויחידה סקלרית. מידע נוסף על TensorCores זמין במאמר A Domain-Specific Supercomputer for Training Deep Neural Networks.
יחידת MXU מורכבת מ-256 x 256 (TPU v6e ו-TPU7x (גרסת Preview)) או מ-128 x 128 (גרסאות TPU לפני v6e) של יחידות הצטברות כפל במערך סיסטולי. יחידות MXU מספקות את רוב כוח המחשוב ב-TensorCore. כל MXU יכול לבצע 16,000 פעולות של הכפלה וצבירה בכל מחזור. כל פעולות הכפל מקבלות קלט בפורמט bfloat16, אבל כל הצבירות מתבצעות בפורמט המספרי FP32.
יחידת הווקטור משמשת לחישובים כלליים כמו הפעלות ו-softmax. היחידה הסקלרית משמשת לבקרת זרימה, לחישוב כתובות זיכרון ולפעולות תחזוקה אחרות.
TPU Pod
TPU Pod הוא קבוצה רציפה של יחידות TPU שמקובצות יחד ברשת ייעודית. מספר שבבי ה-TPU באשכול TPU Pod תלוי בגירסת ה-TPU.
פרוסה
פרוסת TPU היא אוסף של שבבים שנמצאים כולם בתוך אותו TPU Pod ומחוברים באמצעות חיבורים מהירים בין השבבים (ICI). הפרוסות מתוארות במונחים של צ'יפים או TensorCores, בהתאם לגרסת ה-TPU.
המונחים צורת הצ׳יפ וטופולוגיית הצ׳יפ מתייחסים גם לצורות של פרוסות.
פרוסות מרובות לעומת פרוסה אחת
Multislice היא קבוצה של slices, שמרחיבה את הקישוריות של TPU מעבר לחיבורים בין-שבביים (ICI) וממנפת את הרשת של מרכז הנתונים (DCN) להעברת נתונים מעבר ל-slice. הנתונים בכל פלח עדיין מועברים על ידי ICI. באמצעות הקישוריות ההיברידית הזו, ריבוי-פרוסות (Multislice) מאפשרת טעינה במקביל בין פרוסות (slice) ומאפשרת להשתמש במספר גדול יותר של ליבות TPU למשימה אחת מאשר מה שאפשר להכיל בפרוסה (slice) יחידה.
אפשר להשתמש ב-TPU כדי להריץ משימה בפלח אחד או בכמה פלחים. פרטים נוספים זמינים במבוא לריבוי-פרוסות (Multislice).
קוביית TPU
טופולוגיה של 4x4x4 של שבבי TPU שמחוברים ביניהם. ההגדרה הזו רלוונטית רק לטופולוגיות תלת-ממדיות (החל מ-TPU v4).
SparseCore
מעבדי SparseCore הם מעבדי Dataflow שמאיצים מודלים באמצעות פעולות דלילות. תרחיש שימוש עיקרי הוא האצת מודלים של המלצות, שמסתמכים במידה רבה על הטמעות. ל-v5p ול-TPU7x (גרסת טרום-השקה) יש ארבע ליבות SparseCore לכל שבב, ול-v6e יש שתי ליבות SparseCore לכל שבב. הסבר מפורט על השימוש ב-SparseCores מופיע במאמר A deep dive into SparseCore for Large Embedding Models (LEM). אתם קובעים איך מהדר XLA משתמש ב-SparseCores באמצעות דגלי XLA. מידע נוסף זמין במאמר בנושא TPU XLA flags.
חוסן של Cloud TPU ICI
עמידות ה-ICI עוזרת לשפר את עמידות התקלות של קישורים אופטיים ושל מתגי מעגלים אופטיים (OCS) שמחברים יחידות TPU בין קוביות. (חיבורי ICI בתוך קובייה משתמשים בקישורי נחושת שלא מושפעים). הגמישות של ICI מאפשרת לנתב חיבורי ICI מסביב ל-OCS ולתקלות אופטיות ב-ICI. כתוצאה מכך, חל שיפור בזמינות של פריסות TPU, אבל יש פגיעה זמנית בביצועים של ICI.
ב-Cloud TPU v4, v5p ו-TPU7x (בתצוגה מקדימה), חוסן ה-ICI מופעל כברירת מחדל עבור פרוסות בגודל קובייה אחת או יותר, למשל:
- v5p-128 כשמציינים את סוג המאיץ
- 4x4x4 כשמציינים את תצורת המאיץ
גרסאות TPU
הארכיטקטורה המדויקת של שבב TPU תלויה בגרסת ה-TPU שבה אתם משתמשים. כל גרסה של TPU תומכת גם בגדלים ובהגדרות שונים של slice. מידע נוסף על ארכיטקטורת המערכת וההגדרות הנתמכות זמין בדפים הבאים:
ארכיטקטורת ענן TPU
Google Cloud TPU זמינים כמשאבי מחשוב באמצעות מכונות וירטואליות של TPU. אתם יכולים להשתמש ישירות במכונות וירטואליות של TPU לעומסי העבודה שלכם, או להשתמש בהן דרך Google Kubernetes Engine או Vertex AI. בקטעים הבאים מתוארים רכיבים מרכזיים בארכיטקטורת הענן של TPU.
ארכיטקטורת TPU VM
ארכיטקטורת TPU VM מאפשרת להתחבר ישירות למכונה הווירטואלית שמחוברת פיזית למכשיר ה-TPU באמצעות SSH. מכונת TPU וירטואלית, שנקראת גם worker, היא מכונה וירטואלית שמריצה Linux ויש לה גישה ל-TPU הבסיסי. יש לכם גישת root למכונה הווירטואלית, כך שאתם יכולים להריץ קוד שרירותי. אתם יכולים לגשת ליומני ניפוי באגים ולהודעות שגיאה של מהדר ושל זמן ריצה.
מארח יחיד, כמה מארחים ומארח משנה
מארח TPU הוא מכונה וירטואלית שפועלת במחשב פיזי שמחובר לחומרת TPU. עומסי עבודה של TPU יכולים להשתמש במארח אחד או יותר.
עומס עבודה של מארח יחיד מוגבל ל-TPU VM אחד. עומס עבודה עם כמה מארחים מחלק את האימון בין כמה מכונות וירטואליות של TPU. עומס עבודה של מארח משנה לא משתמש בכל הצ'יפים במכונה וירטואלית של TPU.