-sdy-close-shardings

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

-sdy-constant-or-scalar-merger

Объедините идентичные константы и скалярные расширения с соответствующими сегментами.

Выполняет облегченный CSE для констант с идентичными шардингами.

Конвейер импорта разделяет и дублирует константы и скалярные разложения, чтобы не допустить распространения шардинга между различными вариантами использования константного подвычисления. Если после распространения шардинги констант совпадают, этот проход объединяет их для экономии времени компиляции. Подробнее см. в описании -sdy-constant-or-scalar-splitter.

-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, \[{"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>

Выход:

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 перед операцией «точка», так что операция «точка» имеет совместимые шардинги.

Параметры

-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

_Удаляет разброс sdy.all_gather и sdy.reduce для CMV1.

Удаляет all-gather в шаблоне all-gather + dot. Удаляет reduce-scatter в шаблоне dot + reduce-scatter. Этот проход предназначен для совместимости с Collective Matmul V1 (CMV1). Это временное решение для b/432019089.

-sdy-remove-propagation-debug-info

Удаляет отладочную информацию распространения (ребра распространения и исходные сегменты) во время экспорта.

-sdy-remove-sharding-groups

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

-sdy-reshard-to-collectives

Конвертирует ReshardOp в различные коллективные операции Shardy.

Совмещает операции перешарда и переписывает их в различные коллективные операции Shardy. После этого прохода в модуле не остаётся операций перешарда.

При необходимости, если keepRedundantReshards имеет значение true, то останутся только избыточные операции перешардирования. По умолчанию предполагается, что явные перешарды уже были вставлены ( sdy-insert-explicit-reshards ), и избыточные перешарды не сохраняются. Следует сохранять избыточные перешарды, если явные перешарды ещё не были вставлены.

Пример:

Вход:

mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.reshard %arg0 <@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>

Выход:

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>

В приведенном выше примере тензор %0 : tensor<16x2xf32> сегментируется как\[{"x", "y", "z"}, {}\] . Затем, есть операция reshard , которая перешардирует его как\[{"x"}, {}\] . На первой оси, поскольку суффикс {"y", "z"} удаляется после перегруппировки, мы заключаем, что мы собрали все {"y", "z"} . Второе измерение не меняется.

Параметры

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

-sdy-sharding-constraint-to-reshard

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

-sdy-sink-data-flow-edges

Направляет все DataFlowEdgeOp на их вход.

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

Параметры

-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

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

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