-sdy-close-shardings

Cierra los particionamientos de tensores y descarta los ejes replicados.

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

Entrada:

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>

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, hay un conflicto, ya que los tensores lhs y rhs se fragmentan en el eje "x" en sus dimensiones no contractivas. Aquí, el tensor rhs se vuelve a particionar, antes de la operación de punto, de forma explícita para particionarse solo en su primera dimensión y en el eje "x". De esta manera, la operación de punto se vuelve compatible.

-sdy-remove-sharding-groups

Quita ShardingGroupOps después de la propagación.

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

TODO(tomnatan): Considera mover el fragmentación a todos los destinos que puedan tener un fragmento adjunto.

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