-sdy-close-shardings

Ferme les fractionnements de tenseur et supprime les axes répliqués.

-sdy-drop-sharding-rules

Supprime OpShardingRuleAttr de toutes les opérations enregistrées.

-sdy-insert-explicit-reshards

Insère des reshardings explicites pour que toutes les opérations disposent de shardings compatibles.

Un fractionnement compatible signifie essentiellement que l'opération peut accepter les opérandes fractionnés et produire un résultat fractionné sans nécessiter de communications de fractionnement (notez que l'opération peut toujours nécessiter une communication telle que la réduction globale ou les échanges de halo).

Après la propagation, il est possible que certaines opérations présentent toujours des fragmentations incompatibles.

Notez que lorsqu'un axe (ou un sous-axe) est utilisé pour fractionner des dimensions non correspondantes (par exemple, des dimensions non contractantes dans matmul) sur plusieurs tensors, ou lorsqu'un axe fractionne une dimension dans un tenser, mais pas la dimension correspondante dans l'autre tenser, on dit que l'opération présente un conflit de fractionnement. Par conséquent, après ce passage, les opérations ne sont plus en conflit.

Cette étape injecte explicitement des opérations de repartitionnement afin que, pour chaque opération, les dimensions correspondantes soient repartitionnées de la même manière pour tous les opérandes et résultats, et que chaque axe (ou sous-axe) ne puisse être utilisé que pour repartitionner un seul type de dimension.

Exemple pour clarifier:

Entrée :

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>

Sortie :

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>

Dans l'exemple ci-dessus, il existe un conflit, car les tenseurs lhs et rhs sont tous deux partitionnés sur l'axe "x" sur leurs dimensions non contractantes. Ici, le tenseur rhs est refragmenté, avant l'opération de point, de manière explicite pour ne le fragmenter que sur sa première dimension et sur l'axe "x". De cette manière, l'opération de point devient compatible.

-sdy-remove-sharding-groups

Supprime ShardingGroupOps après la propagation.

-sdy-sharding-constraint-to-reshard

Convertit ShardingConstraintOp en ReshardOp.

-sdy-sink-data-flow-edges

Envoie tous les DataFlowEdgeOp dans leur entrée.

Déplace le fractionnement de chaque DataFlowEdgeOp vers son entrée (la cible racine du bord) et remplace l'opération par son entrée.

À FAIRE(tomnatan): envisagez de déplacer le fractionnement vers toutes les cibles auxquelles un fractionnement peut être associé.

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

Répartit de manière uniforme les entrées/sorties FuncOp, ce qui élimine tout besoin de remplissage en raison de fragmentations non divisibles.

Les utilisateurs de Shardy s'attendent à ce que les entrées/sorties de la fonction soient divisibles/partitionnables de manière uniforme pour éviter d'avoir à ajouter un remplissage à leurs tenseurs. La propagation peut entraîner des fragmentations non divisibles des entrées/sorties. Cette étape les met donc à jour avec le plus grand préfixe de fragmentation de dimension de la fragmentation d'origine qui est fragmentée de manière uniforme.