-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>

在上例中,lhsrhs 都會在其非收縮維度上以「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 使用者預期函式輸入/輸出內容可平均分割/分割,以免需要填充張量。傳播可能會讓輸入/輸出端具有不可分割的分割,因此這個傳遞會將這些端更新為原始分割的最大維度分割前置字串,以便平均分割。