Shardy היא מערכת חלוקה של טינסורים מבוססת-MLIR לכל הדיאלקטים. המערכת נוצרה בשיתוף פעולה בין צוות GSPMD לבין צוות PartIR, והיא משלבת את הטוב משתי המערכות ואת הניסיון המשותף של הצוותים והמשתמשים.
יתרונות
- שליטה רבה יותר וצפייה טובה יותר לנתונים של המשתמשים, על ידי שילוב של ההפצה של GSPMD עם חלוקה מצטברת של PartIR.
- תכונות חדשות שמבוססות על ניסיון משותף, למשל תמיכה חדשנית בשינוי הצורה של מודעות, שגורמת לצורך בתקשורת נוספת אם המשתמשים לא יודעים איך לעקוף אותה.
- נוחות שימוש משופרת ויכולת ניפוי באגים משופרת כדי להגביר את המהירות של משתמשי הקצה, למשל באמצעות ייצוג של חלוקה מבוססת-ציר.
- קוד בסיס פשוט בקוד פתוח שמשתמש ב-MLIR, עם קבוצה רחבה יותר של שותפים פעילים (פנימיים, חיצוניים ובאזורי זמן שונים) שתומכים במשתמשים.
רכיבים
- ייצוג של חלוקה לקטעים: ייצוג של חלוקה לקטעים שמבוסס על צירים, שמקושר לרשת לוגית ספציפית (מתוך כמה רשתות אפשריות) ותומך במגבלות על צירים וחלוקות לקטעים של מאפיינים, בחלוקת צירים לפעולות כמו שינוי צורה, בעדיפויות לחלוקה מצטברת ועוד.
- ממשקי API של מהדר: קבוצה של רכיבי מהדר שאפשר להשתמש בהם לצד הייצוג של חלוקת המשנה כדי להשפיע על ההפצה של חלוקת המשנה.
- חלוקות לחלקים של קלט/פלט – אפשר לצרף חלוקה לחלקים לקלט או לפלט של הפונקציה הראשית, כדי לציין שזו הדרך שבה צריך לפצל את הטנסור של הקלט/הפלט כשנותנים אותו לפונקציה או מקבלים אותו ממנה.
- מגבלת חלוקה למקטעים – מחברים חלוקה למקטעים לטרנספורמציה ביניים (למשל, התוצאה של matmul) כדי לציין שזו הדרך שבה צריך לפצל את הטרנספורמציה הזו, או קבוצת משנה של השימושים שלה.
- Shard As/Like – קיבוץ של כמה טינסורים לפי מזהה כדי לציין שצריך לפצל אותם באותו אופן.
- חישוב ידני – מקיף חישוב משנה שמחולק באופן ידני באמצעות קבוצת משנה של צירי רשת, כאשר החלוקות לצירים ידניים אלה מצוינות לכל הקלטות והפלטות, ובתוך החישוב המשנה סוגי הטנזורים הם מקומיים ביחס לחלוקות האלה.
- Sharding Propagation: אלגוריתם העברה שמשלב את העדיפויות של המשתמשים ואת האילוצים של חלוקת המידע לקטעים, עם מודלים של עלויות של מהדרים והיגוריית:
- עדיפויות מוגדרות על ידי משתמשים, למשל: ביצוע מקביליות באצווה ואז ZeRO
- עדיפויות מבוססות-פעולות, למשל פעולות לפי רכיבים קודם ואז פעולות כפל מטריצות וכו'.
- שיטות ניתוח נתונים (heuristics) מפורטות יותר, למשל, העדפה למאפייני קבוצות.
- חלוקה למחיצות של SPMD: רכיב שמפחית את מספר ההחלטות לגבי התפשטות של חלוקה למחיצות על ידי חלוקה של התוכנית לתוכנית SPMD, הוספת התנועה או הפורמט של הנתונים הנדרשים ופעולות קולקטיביות בתהליך.
- בטווח הקצר, בהטמעה הראשונית נעשה שימוש במחיצה הנוכחית של GSPMD SPMD.
- בטווח הארוך, אנחנו מתכננים ליצור מחלק חדש של SPMD שמבוסס על MLIR.
מאגר קוד
הפרויקט Shardy נמצא בפיתוח פעיל, ואנחנו מבקשים משוב מקהילת הקוד הפתוח. הקוד של Shardy זמין בכתובת https://github.com/openxla/shardy