הגדלת מודל במעבדי TPU

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

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

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

דרישות מוקדמות

צריכה להיות לכם הבנה בסיסית של מודלים גדולים של שפה (LLM) ושל ארכיטקטורת טרנספורמרים, אבל לא בהכרח הבנה של אופן הפעולה שלהם בהיקף גדול. כדאי להבין את היסודות של אימון מודלים גדולים של שפה (LLM), ורצוי שתהיה לכם היכרות בסיסית עם JAX. חומרי קריאה שימושיים בנושא ארכיטקטורת טרנספורמרים:

• Illustrated Transformer: פוסט בבלוג על ארכיטקטורת Transformer
• Attention is all you need: the original Transformer paper

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

חשיבות ההגדלה של המודל

מודלים של שפה גדולים (LLM) ורוב המודלים הקטנים פועלים היום קרוב מאוד למגבלות החומרה, ולכן כדי לפתח מודלים צריך לחשוב על יעילות בקנה מידה גדול. שיפור של 20% בהשוואה לביצועים של המתחרים לא רלוונטי אם הוא מגיע עם עלייה של 20% בעלות של יעילות השימוש ב-Google Cloud. ארכיטקטורות מבטיחות של מודלים נכשלות באופן שגרתי, או כי הן לא יכולות לפעול ביעילות בהיקף גדול, או כי לא נעשה מספיק מאמץ באופטימיזציה כדי שהן יוכלו לעשות זאת.

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

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

חלק 1: מושגים

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

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

חלק 2: הרחבת השימוש בטרנספורמרים

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

לבסוף, החלק הזה עוזר לכם לקבל תשובה לשאלה הבסיסית: בהינתן מודל בגודל מסוים ומספר מסוים של שבבים, איך אפשר להריץ את המודל במקביל כדי לשמור על תנאי ההתאמה החזקה. כדי לענות על השאלה הזו, בחלק הזה נדון בארבע טכניקות מקביליות עיקריות שמשמשות לפיצול מודלים על פני כמה שבבים: נתונים, טנסור, צינור ומומחה. בנוסף, מפורטות בו טכניקות אחרות לצמצום דרישות הזיכרון, כמו rematerialization,‏ ZeRO-powered model sharding,‏ host offload ו-gradient accumulation.

• מבוא לפעולות מתמטיות של Transformer: בחלק הזה מוסבר על המתמטיקה שמאחורי התשובות לשאלות לגבי מספר ה-FLOPs שמשמשים את ה-Transformer במהלך מעברים קדימה ואחורה, חישובים לחישוב מספר הפרמטרים וגודל מטמון ה-KV.
• מקביליות של Transformer לאימון: בחלק הזה מפורט התהליך למקסום יעילות האימון באמצעות תיאום של FSDP,‏ Megatron sharding ומקביליות של צינורות. במאמר הזה מוסבר איך לקבוע את ההפצה האופטימלית לגודל מודל ספציפי ולגודל אצווה ספציפי במספר קבוע של שבבים כדי להשיג את קצב העברת הנתונים המקסימלי.
  • אימון של Llama 3 ב-TPU: בחלק הזה מוסבר איך לאמן את Llama 3 ב-TPU, כמה זמן זה עשוי לקחת וכמה זה עשוי לעלות.
• שינוי קנה מידה של טרנספורמר להסקת מסקנות: אחרי שמודל מאומן, צריך להפעיל אותו. הסקת מסקנות מוסיפה שיקול חדש, זמן אחזור, ומשנה את נוף הזיכרון. בקטע הזה מוסבר איך פועל ניהול תנועה מפוצל ואיך כדאי לחשוב על מטמוני KV.
  • הפעלת Llama 3 ב-TPU: בקטע המשנה הזה מוסבר איך להפעיל את Llama 3 ב-TPU, כמה זה עשוי לעלות ומהם השיקולים לגבי זמן האחזור וקצב העברת הנתונים.

חלק 3: הטמעה בפועל

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

• יצירת פרופילים של תוכניות TPU: מודלים גדולים של שפה (LLM) הם מורכבים וקשה לפתח, לבצע אופטימיזציה ולנפות בהם באגים. בקטע הזה מוסבר על מחסנית JAX + XLA ואיך משתמשים בכלי ליצירת פרופילים של JAX/TensorBoard כדי לנפות באגים ולפתור בעיות אמיתיות.
• תכנות של TPU ב-JAX: בקטע הזה מוסבר איך להשתמש בממשקי ה-API של JAX כדי להקביל את החישוב.