-sdy-close-shardings
Ferme les fractionnements de tenseur et supprime les axes répliqués.
-sdy-constant-merger
Fusionnez les constantes identiques avec les shardings correspondants.
Effectue une CSE légère sur les constantes avec des fractionnements identiques.
Le pipeline d'importation divise et duplique les constantes de sorte que le fractionnement ne soit pas propagé entre les différentes utilisations d'un sous-calcul constant. Si les constantes ont les mêmes fractionnements après la propagation, cette étape les fusionne pour gagner du temps de compilation.
-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 shardings 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 tenseurs, ou lorsqu'un axe fractionne une dimension dans un tenseur, mais pas la dimension correspondante dans l'autre tenseur, 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 :
Entrée :
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>
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, lhs
et rhs
sont tous deux partitionnés sur l'axe "x" sur leurs dimensions non contractantes, ce qui est incompatible. La carte insère un reshard explicite sur rhs
avant l'opération de point, de sorte que l'opération de point dispose de shardings compatibles.
-sdy-remove-sharding-groups
Supprime ShardingGroupOps après la propagation.
-sdy-reshard-to-collectives
Convertit ReshardOp en diverses opérations collectives Shardy.
Correspond aux opérations de reshardage et les réécrit en différentes opérations collectives Shardy. Après cette étape, aucune opération de resharding ne reste dans le module. Cette étape suppose que des reshards explicites ont déjà été insérés (sdy-insert-explicit-reshards
).
Exemple :
Entrée :
mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.reshard %arg0 <@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>
Sortie :
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>
Dans l'exemple ci-dessus, le tenseur %0 : tensor<16x2xf32>
est partitionné en tant que \[{"x", "y", "z"}, {}\]
. Ensuite, une opération reshard
la re-partitionne en tant que \[{"x"}, {}\]
. Sur les premiers axes, comme le suffixe {"y", "z"}
est supprimé après le reshard, nous en déduisons que nous avons {"y", "z"}
tout rassemblé. La deuxième dimension n'est pas modifiée.
-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.
Options
-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
Insère des reshards explicites pour des optimisations spécifiques.
Cette carte est une solution temporaire jusqu'à ce que nous puissions activer la carte sdy-insert-explicit-reshards
par défaut.
Il nous permet d'insérer des reshards explicites sur des opérations spécifiques à des fins d'optimisation.
-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.