-sdy-close-shardings

Fecha as fragmentações de tensor e descarta os eixos replicados.

-sdy-constant-or-scalar-merger

Mescla constantes idênticas e expansões escalares com fragmentações correspondentes.

Realiza uma CSE leve em constantes com fragmentações idênticas.

O pipeline de importação divide e duplica as constantes e expansões escalares para que a fragmentação não seja propagada entre diferentes usos de um subcálculo constante. Se as constantes tiverem as mesmas fragmentações após a propagação, essa passagem as mesclará para economizar tempo de compilação. Consulte -sdy-constant-or-scalar-splitter para mais informações.

-sdy-convert-global-to-local

Converte um programa SDY de formas globais para locais.

Converte um programa SDY de formas globais para locais, particionando dimensões lógicas com base em atributos de fragmentação.

Essa passagem aproveita um conversor de tipo para mapear RankedTensorType de formas lógicas globais para formas físicas locais do dispositivo.

-sdy-drop-sharding-rules

Descarta OpShardingRuleAttr de todas as operações registradas.

-sdy-insert-explicit-reshards

Insere novas fragmentações explícitas para que todas as operações tenham fragmentações compatíveis.

Uma fragmentação compatível significa que a operação pode aceitar os operandos fragmentados e produzir um resultado fragmentado sem exigir comunicações de nova fragmentação. 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 fragmentações incompatíveis.

Quando um eixo (ou subeixo) é usado para fragmentar dimensões não correspondentes (por exemplo, dimensões não contratuais em matmul) em vários tensores ou quando um eixo fragmenta uma dimensão em um tensor, mas não a dimensão correspondente no outro tensor, dizemos que a operação tem um conflito de fragmentação. Portanto, após essa passagem, as operações ficam sem conflitos.

Essa passagem injeta operações de nova fragmentação explicitamente 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 fragmentar 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 fragmentados no eixo "x" nas dimensões não contratuais, o que é incompatível. A passagem insere uma nova fragmentação explícita em rhs antes da operação de ponto, para que ela tenha fragmentações compatíveis.

Opções

-enable-full-version                  : Enable full version.
-avoid-reshards-on-named-computations : Avoid explicit reshards/collectives on named computations.

-sdy-remove-all-gather-reduce-scatter-for-cmv1

_Remove sdy.all_gather e sdy.reducescatter para CMV1.

Remove a coleta total no padrão de coleta total + ponto. Remove a redução de dispersão no padrão de ponto + redução de dispersão. Essa passagem é para compatibilidade com o matmul coletivo V1 (CMV1). É uma solução temporária para b/432019089.

-sdy-remove-propagation-debug-info

Remove informações de depuração de propagação (bordas de propagação e fragmentações de origem) durante a exportação.

-sdy-remove-sharding-groups

Remove ShardingGroupOps após a propagação.

-sdy-remove-sub-axes-in-input-output-shardings

Remove subeixos em fragmentações de entrada/saída.

Alguns usuários do Shardy esperam que as entradas/saídas de função tenham fragmentações sem subeixos. Essa passagem remove subeixos e os eixos finais das fragmentações de entrada/saída. Essa passagem geralmente ocorre após sdy-update-non-divisible-input-output-shardings para garantir que a remoção de subeixos não introduza fragmentações não divisíveis.

-sdy-reshard-to-collectives

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

Corresponde a operações de nova fragmentação e as reescreve em várias operações coletivas do Shardy. Após essa passagem, nenhuma operação de nova fragmentação permanece no módulo.

Opcionalmente, se keepRedundantReshards for verdadeiro, as únicas operações de nova fragmentação que permanecem são as redundantes. Por padrão, ele pressupõe que as novas fragmentações explícitas já foram inseridas (sdy-insert-explicit-reshards) e não mantém as redundantes. Ele deve manter as novas fragmentações redundantes se as explícitas ainda não tiverem sido inseridas.

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> é fragmentado como \[{"x", "y", "z"}, {}\]. Em seguida, há uma operação reshard que o fragmenta novamente como \[{"x"}, {}\]. Nos primeiros eixos, como o sufixo {"y", "z"} é removido após a nova fragmentação, inferimos que coletamos {"y", "z"}. A segunda dimensão não é alterada.

Opções

-keep-redundant-reshards : Whether it keeps redundant reshards or removes.

-sdy-sharding-constraint-to-reshard

Converte ShardingConstraintOp em ReshardOp.

-sdy-sink-data-flow-edges

Afunda todos os DataFlowEdgeOp na entrada.

Move a 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-update-non-divisible-input-output-shardings

Faz com que as entradas/saídas do FuncOp sejam fragmentadas uniformemente, removendo a necessidade de preenchimento devido a fragmentações não divisíveis.

Os usuários do Shardy esperam que as entradas/saídas de função sejam divisíveis/fragmentáveis uniformemente para evitar a necessidade de preencher os tensores. A propagação pode fazer com que as entradas/saídas tenham fragmentações não divisíveis. Portanto, essa passagem as atualiza para o maior prefixo de fragmentação de dimensão da fragmentação original que é fragmentada uniformemente.