-sdy-close-shardings

Fecha partições de tensor e exclui eixos replicados.

-sdy-constant-merger

Mesclar constantes idênticas com particionamentos correspondentes.

Realiza uma CSE leve em constantes com particionamentos idênticas.

O pipeline de importação divide e duplica as constantes para que o fragmentação não seja propagado entre diferentes usos de uma subcomputação constante. Se as constantes tiverem os mesmos shardings após a propagação, esse cartão vai mesclar para economizar tempo de compilação.

-sdy-drop-sharding-rules

Exclui OpShardingRuleAttr de todas as operações registradas.

-sdy-insert-explicit-reshards

Insere reshards explícitos para que todas as operações tenham particionamentos compatíveis.

Um fragmentação compatível significa que a operação pode aceitar os operandos fragmentados e produzir um resultado fragmentado sem exigir nenhuma comunicação de resfragmentação. No entanto, a operação ainda pode exigir comunicação, como redução total ou trocas de halo.

Após a propagação, algumas operações ainda podem ter divisões incompatíveis.

Observe que, quando um eixo (ou subeixo) é usado para dividir dimensões não correspondentes (por exemplo, dimensões não contratuais em matmul) em vários tensores ou quando um eixo divide uma dimensão em um tensor, mas não a dimensão correspondente no outro tensor, diz-se que a operação tem um conflito de fragmentação. Portanto, depois dessa passagem, as operações ficam sem conflitos.

Esse cartão injeta operações de reshard de forma explícita para que, em cada operação, as dimensões correspondentes sejam fragmentadas da mesma maneira em todos os operandos e resultados, e cada eixo (ou subeixo) só possa ser usado para dividir um único tipo de dimensão.

Exemplo:

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>

Saída:

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>

No exemplo acima, lhs e rhs são divididos no eixo "x" nas dimensões não contratuais, o que é incompatível. O cartão insere um reshard explícito em rhs antes da operação de ponto, para que a operação de ponto tenha particionamentos compatíveis.

-sdy-remove-sharding-groups

Remove ShardingGroupOps após a propagação.

-sdy-reshard-to-collectives

Converte ReshardOp em várias operações coletivas Shardy.

Combina operações de reshard e as reescreve em várias operações coletivas de Shardy. Após essa passagem, nenhuma operação de reshard permanece no módulo. Este cartão pressupõe que os reshards explícitos já foram inseridos (sdy-insert-explicit-reshards).

Exemplo:

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>

Saída:

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>

No exemplo acima, o tensor %0 : tensor<16x2xf32> é dividido como \[{"x", "y", "z"}, {}\]. Em seguida, há uma operação reshard que o redimensiona como \[{"x"}, {}\]. Nos primeiros eixos, como o sufixo {"y", "z"} é removido após a reshard, inferimos que reunimos todos os {"y", "z"}. A segunda dimensão não é alterada.

-sdy-sharding-constraint-to-reshard

Converte ShardingConstraintOp em ReshardOp.

-sdy-sink-data-flow-edges

Acumula todos os DataFlowEdgeOp na entrada.

Move o fragmentação de cada DataFlowEdgeOp para a entrada (o destino raiz da borda) e substitui a operação pela entrada.

Opções

-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

Insere reshards explícitos para otimizações específicas.

Esse cartão é uma solução temporária até que possamos ativar o cartão sdy-insert-explicit-reshards por padrão.

Isso nos permite inserir reshards explícitos em operações específicas para otimizações.

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

Faz com que as entradas/saídas de FuncOp sejam divididas de maneira uniforme, removendo a necessidade de preenchimento devido a particionamentos não divisíveis.

Os usuários do Shardy esperam que as entradas/saídas da função sejam divisíveis/fragmentáveis de maneira uniforme para evitar o preenchimento dos tensores. A propagação pode fazer com que as entradas/saídas tenham particionamentos não divisíveis. Portanto, esse processo as atualiza para o maior prefixo de particionamento de dimensão do particionamento original que é dividido de maneira uniforme.