-sdy-close-shardings

Закрывает тензорные сегменты и удаляет реплицированные оси.

-sdy-drop-sharding-rules

Удаляет OpShardingRuleAttr из всех зарегистрированных операций.

-sdy-insert-explicit-reshards

Вставляет явные сегменты, чтобы все операции имели совместимые сегменты.

Совместимое сегментирование по сути означает, что операция может принимать сегментированные операнды и выдавать сегментированный результат, не требуя каких-либо дополнительных соединений (обратите внимание, что операция все равно может требовать связи, такой как all-reduce или halo-swaps).

После распространения некоторые операции могут по-прежнему иметь несовместимые сегменты.

Обратите внимание: когда ось (или подось) используется для сегментирования несоответствующих измерений (например, несжимающихся измерений в matmul) по нескольким тензорам или когда ось сегментирует измерение в одном тензоре, но не соответствующее измерение в другой тензор, говорят, что операция имеет конфликт шардинга. Следовательно, после этого прохода операции становятся бесконфликтными.

Этот проход явно вводит операции повторного разделения, так что для каждой операции соответствующие измерения сегментируются одинаковым образом для всех операндов и результатов, и каждая ось (или подось) может использоваться только для сегментирования одного типа измерения.

Поясняющий пример:

Вход:

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>

Выход:

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>

В приведенном выше примере возникает конфликт, поскольку оба тензора lhs и rhs сегментированы по оси «x» в своих несжимающихся измерениях. Здесь тензор rhs подвергается повторному сегментированию перед операцией с точкой явно для сегментирования только по его первому измерению и по оси «x». Таким образом, операция с точкой становится совместимой.

-sdy-remove-sharding-groups

Удаляет ShardingGroupOps после распространения.

-sdy-sharding-constraint-to-reshard

Преобразует ShardingConstraintOp в ReshardOp.

-sdy-sink-data-flow-edges

Помещает все DataFlowEdgeOp в свои входные данные.

Перемещает сегментацию каждого DataFlowEdgeOp на его вход (корневой целевой объект ребра) и заменяет операцию его входом.

TODO(tomnatan): рассмотрите возможность перемещения сегментирования на все цели, к которым может быть прикреплено сегментирование.

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

Делает входы/выходы FuncOp равномерно сегментированными, устраняя необходимость в заполнении из-за неделимых фрагментов.

Пользователи Shardy ожидают, что входные/выходные данные функции будут равномерно делимыми/разделяемыми, чтобы избежать необходимости заполнения их тензоров. Распространение может привести к тому, что входные/выходные данные будут иметь неделимые сегменты, поэтому на этом этапе они обновляются до префикса сегментирования наибольшего размера исходного сегментирования, который сегментируется равномерно.