-sdy-close-shardings
關閉張量分割作業,並捨棄複製的軸。
-sdy-constant-merger
合併具有相符分割的相同常數。
針對具有相同切割的常數執行輕量 CSE。
匯入管道會分割及複製常數,因此不會在常數子運算的不同用途之間傳播分割作業。如果常數在傳播後具有相同的切割,這個階段會將它們合併,以節省編譯時間。
-sdy-drop-sharding-rules
從所有已註冊的作業中移除 OpShardingRuleAttr
。
-sdy-insert-explicit-reshards
插入明確的 reshard,讓所有作業都具有相容的分割。
相容的分割作業基本上是指作業可以接受分割運算元,並產生分割結果,而不需要任何重新分割通訊 (請注意,作業可能仍需要通訊,例如 all-reduce 或 halo-swap)。
傳播後,部分作業可能仍有不相容的分割作業。
請注意,如果軸 (或子軸) 用於在多個張量中分割不相符的維度 (例如 matmul 中的非收縮維度),或是軸分割一個張量中的維度,但不分割另一個張量中的相符維度,則表示該運算有分割衝突。因此,在這個階段結束後,作業就不會發生衝突。
這個階段會明確插入重新分割作業,以便針對每項作業,在所有運算元和結果中以相同方式分割相應的維度,且每個軸 (或子軸) 只能用於分割單一維度類型。
範例:
輸入:
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
都會在其非收縮維度上以「x」軸分割,這兩者不相容。在點運算之前,傳遞會在 rhs
上插入明確的 reshard,讓點運算具有相容的分割。
-sdy-remove-sharding-groups
在傳播後移除 ShardingGroupOps。
-sdy-reshard-to-collectives
將 ReshardOp 轉換為各種 Shardy 集體作業。
比對 reshard 作業,並將這些作業重寫為各種 Shardy 集體作業。在這個階段結束後,模組中就不會保留任何 reshard 作業。這個階段假設已插入明確的 reshard (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"}, {}\]
。在第一個軸上,由於 reshard 後會移除後置字串 {"y", "z"}
,因此我們推斷已收集到所有 {"y", "z"}
。第二個維度不會變更。
-sdy-sharding-constraint-to-reshard
將 ShardingConstraintOp 轉換為 ReshardOp。
-sdy-sink-data-flow-edges
將所有 DataFlowEdgeOp
匯入輸入內容。
將每個 DataFlowEdgeOp
的分割作業移至其輸入內容 (邊緣的根目錄目標),並將作業替換為其輸入內容。
選項
-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
插入明確的 reshard,以便進行特定最佳化。
在我們預設啟用 sdy-insert-explicit-reshards
通行證之前,這項通行證是暫時的解決方案。
這可讓我們在特定運算中插入明確的 reshard,以利最佳化。
-sdy-update-non-divisible-input-output-shardings
讓 FuncOp 輸入/輸出均勻分割,消除因不可分割的分割而需要填充的情況。
Shardy 使用者預期函式輸入/輸出內容可平均分割/分割,以免需要填充張量。傳播可能會讓輸入/輸出端具有不可分割的分割,因此這個傳遞會將這些端更新為原始分割的最大維度分割前置字串,以便平均分割。