-sdy-close-shardings

Schließt Tensor-Shardings und verwirft replizierte Achsen.

-sdy-drop-sharding-rules

Löst OpShardingRuleAttr bei allen registrierten Geschäftsbereichen aus.

-sdy-insert-explicit-reshards

Fügt explizite Reshards ein, damit alle Vorgänge kompatible Shardings haben.

Eine kompatible Sharding-Technologie bedeutet im Wesentlichen, dass der Vorgang die geShardeten Operanden akzeptieren und ein geShardetes Ergebnis ohne Reshard-Kommunikation liefern kann. Beachten Sie, dass der Vorgang möglicherweise noch Kommunikation wie All-Reduce oder Halo-Swaps erfordert.

Nach der Replikation haben einige Vorgänge möglicherweise weiterhin inkompatible Shardings.

Wenn eine Achse (oder eine untergeordnete Achse) verwendet wird, um nicht übereinstimmende Dimensionen (z.B. nicht schrumpfende Dimensionen in matmul) über mehrere Tensoren zu trennen, oder wenn eine Achse eine Dimension in einem Tensor, aber nicht die entsprechende Dimension im anderen Tensor teilt, liegt ein Sharding-Konflikt vor. Nach diesem Pass werden die Vorgänge also konfliktfrei.

Mit diesem Pass werden Resharding-Vorgänge explizit eingefügt, damit für jeden Vorgang entsprechende Dimensionen auf dieselbe Weise über alle Operanden und Ergebnisse hinweg gesplittet werden. Außerdem kann jede Achse (oder untergeordnete Achse) nur zum Sharding eines einzelnen Dimensionstyps verwendet werden.

Ein Beispiel zur Verdeutlichung:

Eingabe:

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

Ausgabe:

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>

Im obigen Beispiel kommt es zu einem Konflikt, da die Tensoren lhs und rhs in ihren nicht schrumpfenden Dimensionen beide entlang der Achse „x“ geSharded sind. Hier wird der Tensor rhs vor dem Punktvorgang explizit neu segmentiert, damit er nur nach der ersten Dimension und nach der Achse „x“ segmentiert wird. So wird die Punktoperation kompatibel.

-sdy-remove-sharding-groups

Entfernt ShardingGroupOps nach der Replikation.

-sdy-sharding-constraint-to-reshard

Konvertiert ShardingConstraintOp in ReshardOp.

-sdy-sink-data-flow-edges

Leitet alle DataFlowEdgeOp in die Eingabe weiter.

Verschiebt das Sharding jeder DataFlowEdgeOp an ihre Eingabe (das Stammziel der Kante) und ersetzt die Operation durch ihre Eingabe.

TODO(tomnatan): Sharding auf alle Ziele umstellen, denen ein Sharding zugeordnet werden kann.

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

Die FuncOp-Eingaben/-Ausgaben werden gleichmäßig geSharded, sodass kein Padding aufgrund nicht teilbarer Shardings erforderlich ist.

Nutzer von Shardy erwarten, dass die Funktionseingaben/-ausgaben gleichmäßig teilbar/shardbar sind, damit keine Tensoren aufgefüllt werden müssen. Durch die Replikation können Eingaben/Ausgaben nicht teilbare Shardings haben. Daher werden sie in diesem Pass auf das größte Dimensions-Sharding-Präfix des ursprünglichen Shardings aktualisiert, das gleichmäßig geSharded ist.