חלוקה לפלחים ללא תלות בניב

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

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

כללי פיצול

כלל חלוקה קובע איך אנחנו מפיצים את הפעולה. מכיוון ש-Shardy תלוי עכשיו ב-StableHLO, הוא מגדיר כללי חלוקה לפלחים לכל פעולת stablehlo. בנוסף, Shardy מספק את ShardingRuleOpInterface, שבעזרתו בעלי הדיאלקט יכולים להשתמש בפעולות שלהם כדי להגדיר כללי חלוקה לפלחים לפעולות שלהם. כל עוד הפעולה מיישמת את הממשק הזה, Shardy יוכל להפיץ אותה.

def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
  let methods = [
    InterfaceMethod<
      /*desc=*/[{
        Returns the sharding rule of the op.
      }],
      /*retType=*/"mlir::sdy::OpShardingRuleAttr",
      /*methodName=*/"getShardingRule"
    >,
  ];
}

פעולות של זרימת נתונים

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

def ShardableDataFlowOpInterface :
    OpInterface<"ShardableDataFlowOpInterface"> {
  (get|set)BlockArgumentEdgeOwnerShardings;
  (get|set)OpResultEdgeOwnerShardings;
  getBlockArgumentEdgeOwners;
  getOpResultEdgeOwners;
  getEdgeSources;
  // ...
}

מומלץ גם לעיין במאמר Data flow ops כדי לקבל סקירה כללית על האופן שבו אנחנו מטפלים ב-data flow ops.

ממשקים שעדיין לא יושמו

בעתיד נוסיף עוד ממשקים ומאפיינים כדי להפוך את Shardy לגמיש יותר ולא תלוי-דיאלקט. הרשימה מופיעה בהמשך.

פיצול קבוע

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

לדוגמה, באיור הבא, אם add מחולק למקטעים, זה לא אמור להשפיע על אופן החלוקה למקטעים של divide ו-subtract (בחלקים שונים של החישוב).

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

כדי לעשות זאת, משתמשי Shardy צריכים להגדיר: - העברה מ-your_dialect.constant אל sdy.constant; - מאפיין sdy::ConstantLike, כמו iota; - מאפיין mlir::Elementwise לפעולות על רכיבים, כמו add ו-multiply; - מאפיין sdy::ConstantFoldable לפעולות כמו slice/broadcast. מבחינה טכנית, אפשר לחשב את הפעולות האלה בזמן הידור, אם כל המשתנים שלהן הם קבועים.

עדיפויות תפעוליות

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

בתרשים הבא מוצג איך משתמשים בעדיפויות ב-GSPMD כדי להעביר את הפעולות בסדר הנכון.

במסמך GSPMD מוסבר למה חשוב להגדיר סדר עדיפויות לפעולות.

לא תלויים בניב

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