-sdy-close-shardings
Cierra los particionamientos de tensores y descarta los ejes replicados.
-sdy-constant-merger
Combina constantes idénticas con particiones coincidentes.
Realiza una CSE liviana en constantes con fragmentaciones idénticas.
La canalización de importación divide y duplica las constantes de modo que el fragmentación no se propague entre los diferentes usos de una subcomputación constante. Si las constantes tienen los mismos fragmentos después de la propagación, este pase las combina para ahorrar tiempo de compilación.
-sdy-drop-sharding-rules
Quita OpShardingRuleAttr de todas las operaciones registradas.
-sdy-insert-explicit-reshards
Inserta reshards explícitos para que todas las operaciones tengan particiones compatibles.
En esencia, un fragmento compatible significa que la operación puede aceptar los operandos fragmentados y producir un resultado fragmentado sin requerir ninguna comunicación de fragmentación (ten en cuenta que la operación aún puede requerir comunicación, como la reducción total o los intercambios de halo).
Después de la propagación, es posible que algunas operaciones sigan teniendo particiones incompatibles.
Ten en cuenta que, cuando se usa un eje (o subeje) para dividir dimensiones que no corresponden (p.ej., dimensiones que no se contraen en matmul) en varios tensores, o cuando un eje divide una dimensión en un tensor, pero no la dimensión correspondiente en el otro tensor, se dice que la operación tiene un conflicto de división. Por lo tanto, después de este pase, las operaciones dejan de generar conflictos.
Este pase inserta operaciones de refragmentación de forma explícita para que, para cada operación, las dimensiones correspondientes se fragmenten de la misma manera en todos los operandos y resultados, y cada eje (o subeje) solo se pueda usar para fragmentar un solo tipo de dimensión.
Ejemplo:
Entrada:
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>
Resultado:
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>
En el ejemplo anterior, lhs y rhs se fragmentan en el eje "x" en sus dimensiones no contractivas, lo que es incompatible. El pase inserta un reescalado explícito en rhs antes de la operación de punto, de modo que la operación de punto tenga particiones compatibles.
-sdy-remove-sharding-groups
Quita ShardingGroupOps después de la propagación.
-sdy-reshard-to-collectives
Convierte ReshardOp en varias operaciones colectivas de Shardy.
Coincide con las operaciones de reshard y las reescribe en varias operaciones colectivas de Shardy. Después de este pase, no quedan operaciones de reshard en el módulo. Este pase supone que ya se insertaron reshards explícitos (sdy-insert-explicit-reshards).
Ejemplo:
Entrada:
mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.reshard %arg0 <@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>
Resultado:
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>
En el ejemplo anterior, el tensor %0 : tensor<16x2xf32> se particiona como \[{"x", "y", "z"}, {}\]. Luego, hay una operación reshard que lo vuelve a particionar como \[{"x"}, {}\]. En los primeros ejes, como se quita el sufijo {"y", "z"} después del reshard, inferimos que tenemos {"y", "z"} todo reunido. La segunda dimensión no cambia.
-sdy-sharding-constraint-to-reshard
Convierte ShardingConstraintOp en ReshardOp.
-sdy-sink-data-flow-edges
Envía todos los DataFlowEdgeOp a su entrada.
Mueve el fragmentación de cada DataFlowEdgeOp a su entrada (el destino raíz del borde) y reemplaza la operación por su entrada.
Opciones
-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
Inserta reshards explícitos para optimizaciones específicas.
Este pase es una solución temporal hasta que podamos habilitar el pase sdy-insert-explicit-reshards de forma predeterminada.
Nos permite insertar reshards explícitos en operaciones específicas para las optimizaciones.
-sdy-update-non-divisible-input-output-shardings
Hace que las entradas y salidas de FuncOp se fragmenten de manera uniforme, lo que elimina la necesidad de relleno debido a fragmentaciones no divisibles.
Los usuarios de Shardy esperan que las entradas y salidas de la función sean divisibles o particionables de manera uniforme para evitar que se requieran rellenos en sus tensores. La propagación puede hacer que las entradas o salidas tengan particiones no divisibles, por lo que este pase las actualiza al prefijo de partición de dimensión más grande del particionado original que está particionado de forma uniforme.