טסלה רוצה לקחת סיליקון למידת מכונה לדוג'ו

כדי להרוות את הצמא לדגמי AI ולמידת מכונה גדולים מתמיד, טסלה חשפה שפע של פרטים ב-Hot Chips 34 על ארכיטקטורת מחשוב העל המותאמת אישית שלהם בשם Dojo.

המערכת היא בעצם מחשב-על עצום הניתן לחיבור, אם כי בניגוד למה שאנו רואים במחשב 500 למעלה, הוא בנוי מארכיטקטורה מותאמת אישית לחלוטין המשתרעת על סיליקון המחשוב, הרשת והקלט/פלט (I/O) ועד לארכיטקטורת ערכת הוראות (ISA), אספקת חשמל, אריזה וקירור. כל זה נעשה במטרה מפורשת להפעיל אלגוריתמים מותאמים וספציפיים לאימון מכונה בקנה מידה.

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

עם זאת, הוא טען ששיטות מסורתיות להגדלת עומסי עבודה מבוזרים לא הצליחו להאיץ בקצב הדרוש כדי לעמוד בקצב הדרישות של למידת מכונה. למעשה, חוק מור לא חותך את זה וגם לא המערכות זמינות לאימון AI/ML בקנה מידה, כלומר שילוב כלשהו של CPU/GPU או בנסיבות נדירות יותר באמצעות מאיצי AI מיוחדים. 

"באופן מסורתי אנחנו בונים שבבים, אנחנו שמים אותם על חבילות, חבילות עוברות על PCB, שנכנסים למערכות. מערכות נכנסות למתלים", אמר Venkataramanan. הבעיה היא בכל פעם שהנתונים עוברים מהשבב לחבילה ומחוץ לחבילה, הם גוררים קנס של חביון ורוחב פס.

סנדוויץ' במרכז הנתונים

אז כדי לעקוף את המגבלות, Venkataramanan והצוות שלו התחילו מאפס.

"ממש מהראיון שלי עם אילון, הוא שאל אותי מה אתה יכול לעשות ששונה ממעבדים ומעבדים גרפיים עבור AI. אני מרגיש שכל הצוות עדיין עונה על השאלה הזו".

זה הוביל לפיתוח של אריח האימון Dojo, אשכול מחשוב עצמאי התופס חצי רגל מעוקב המסוגל ל-556 TFLOPS של ביצועי FP32 באריזה מקוררת נוזלים של 15kW.

כל אריח מצויד ב-11GB של SRAM ומחובר על בד של 9TB/s באמצעות פרוטוקול הובלה מותאם אישית לאורך כל הערימה.

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

בלב אריח האימון נמצא ה-D1 של טסלה, קוביית טרנזיסטור של 50 מיליארד, המבוססת על תהליך 7nm של TSMC. טסלה אומרת שכל D1 מסוגל ל-22 TFLOPS של ביצועי FP32 ב-TDP של 400W. עם זאת, טסלה מציינת כי השבב מסוגל להריץ מגוון רחב של חישובי נקודה צפה כולל כמה מותאמים אישית.

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

לאחר מכן טסלה לקחה 25 D1s, אספה אותם עבור קוביות טובות ידועות, ואז ארזה אותם באמצעות טכנולוגיית system-on-wafer של TSMC כדי "להשיג כמות עצומה של אינטגרציה מחשוב בהשהיה נמוכה מאוד ורוחב פס גבוה מאוד", אמר.

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

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

לשים את הכל ביחד

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

לשם כך פיתחה טסלה גם את מעבד ממשק ה-Dojo (DIP), שמתפקד כגשר בין המעבד המארח למעבדי האימון. ה-DIP משמש גם כמקור לזיכרון משותף ברוחב פס גבוה (HBM) וכ-NIC מהיר של 400Gbit/sec.

כל DIP כולל 32GB של HBM ועד חמישה מהכרטיסים הללו יכולים להיות מחוברים לאריח אימון במהירות 900GB/s עבור מצטבר של 4.5TB/s למארח עבור סך של 160GB של HBM לכל אריח.

צמדי תצורת V1 של טסלה של אריחים אלה - או 150 קוביות D1 - במערך תמכו בארבעה מעבדים מארחים, שכל אחד מהם מצויד בחמישה כרטיסי DIP כדי להשיג ביצועים נטענים של BF16 או CFP8.

יחד, Venkataramanan אומר שהארכיטקטורה - מפורטת לעומק כאן by הפלטפורמה הבאה - מאפשר לטסלה להתגבר על המגבלות הקשורות למאיצים מסורתיים כמו Nvidia ו-AMD.

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

זו לא בעיה חדשה, הוא ציין. מתג ה-NV של Nvidia, למשל, מאפשר לאגד את הזיכרון על פני גדות גדולות של GPUs. עם זאת, Venkataramanan טוען שזה לא רק מוסיף מורכבות, אלא מציג חביון ומתפשר על רוחב הפס.

"חשבנו על זה כבר מההתחלה. אריחי המחשוב שלנו וכל אחד מהדפים נוצרו כדי להתאים דגמים גדולים", אמר Venkataramanan.

תוכנה

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

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

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

"מה שעשינו זה שהשתמשנו ב-PiTorch. יצרנו שכבת ביניים, שעוזרת לנו לבצע מקבילות כדי להרחיב את החומרה מתחתיה. מתחת לכל דבר יש קוד קומפילציה", אמר. "זו הדרך היחידה ליצור ערימות תוכנה הניתנות להתאמה לכל אותם עומסי עבודה עתידיים."

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

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