-sdy-close-shardings

סגירת חלוקות טינסורים ושחרור של צירים שהועתקו.

-sdy-constant-merger

מיזוג של קבועים זהים עם חלוקות תואמות.

ביצוע CSE קל משקל על קבועים עם חלוקות זהות.

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

-sdy-drop-sharding-rules

הסרת OpShardingRuleAttr מכל התפעולים הרשומים.

-sdy-insert-explicit-reshards

הוספה של חלוקות מחדש מפורשות כדי שכל הפעולות יהיו עם חלוקות תואמות.

חלוקה תואמת למשנה (shard) היא למעשה היכולת של הפעולה לקבל את אופרטנדים המחולקים וליצור תוצאה מחולקת בלי צורך בתקשורת של חלוקה מחדש (שימו לב: יכול להיות שהפעולה עדיין תצטרך תקשורת, כמו all-reduce או halo-swaps).

אחרי ההעברה, יכול להיות שחלק מהפעולות עדיין יכללו חלוקות לא תואמות.

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

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

דוגמה:

קלט:

mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"x"}, {"y"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
stablehlo.dot %lhs, %rhs {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

פלט:

sdy.mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"x"}, {"y"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
%0 = sdy.reshard %rhs <@mesh, \[{"y"}, {}\]> : tensor<32x16xf32>
stablehlo.dot %lhs, %0 {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

בדוגמה שלמעלה, גם lhs וגם rhs מחולקים לקטעים (shards) בציר 'x' לפי המאפיינים שלהם שלא מתכווצים, וזוהי התאמה לא תואמת. הכרטיס מוסיף חלוקה מחדש (reshard) מפורשת של rhs לפני פעולת הנקודה, כדי שלפעולת הנקודה יהיו חלוקות תואמות.

-sdy-remove-sharding-groups

הסרת ShardingGroupOps אחרי העברה (propagation).

-sdy-reshard-to-collectives

המרת ReshardOp לפעולות קולקטיביות שונות של Shardy.

התאמה של פעולות חלוקה מחדש של שריד (reshard) וכתיבה מחדש שלהן לפעולות קולקטיביות שונות של Shardy. אחרי השלבים האלה, לא נותרות פעולות של חלוקה מחדש של קטעי נתונים במודול. בכרטיס הזה אנחנו יוצאים מנקודת הנחה שכבר הוכנסו חלוקות מחדש מפורשות (sdy-insert-explicit-reshards).

דוגמה:

קלט:

mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.reshard %arg0 <@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>

פלט:

mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.all_gather \[{"y", "z"}, {}\] %arg0 out_sharding=<@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>

בדוגמה שלמעלה, הטנזור %0 : tensor<16x2xf32> מחולק ל-\[{"x", "y", "z"}, {}\]. לאחר מכן, יש פעולת reshard שמבצעת חלוקה מחדש של המרחב האחסון כ-\[{"x"}, {}\]. בציורים הראשונים, מאחר שהסיומת {"y", "z"} מוסרת אחרי חלוקה מחדש, אפשר להסיק שאנחנו אוספים את כל הערכים של {"y", "z"}. המאפיין השני לא משתנה.

-sdy-sharding-constraint-to-reshard

המרת ShardingConstraintOp ל-ReshardOp.

-sdy-sink-data-flow-edges

מזין את כל DataFlowEdgeOp לקלט שלו.

העברת הפיצול של כל DataFlowEdgeOp לקלט שלו (יעד הבסיס של הקצה) והחלפת הפעולה (op) בקלט שלו.

אפשרויות

-sink-debug-sharding-origins          : Whether to sink the debug sharding origins info. See `debug-sharding-origins` option in propagation for more info.
-sink-debug-propagation-edge-sharding : Whether to sink the debug propagation edge sharding info. See `debug-propagation-edge-sharding` option in propagation for more info.

-sdy-temp-explicit-reshards-for-optimizations

הוספה של חלוקות מחדש מפורשות לאופטימיזציות ספציפיות.

הכרטיס הזה הוא פתרון זמני עד שנוכל להפעיל את הכרטיס sdy-insert-explicit-reshards כברירת מחדל.

כך אנחנו יכולים להוסיף חלוקות מחדש מפורשות של פעולות ספציפיות לצורך אופטימיזציה.

-sdy-update-non-divisible-input-output-shardings

הפעולה הזו מאפשרת לחלק באופן שווה את הקלט/הפלט של FuncOp, ומבטלת את הצורך במילוי בגלל חלוקות שלא ניתן לחלק.

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