-sdy-close-shardings

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

-sdy-constant-or-scalar-merger

Объединить одинаковые константы и скалярные разложения с соответствующими фрагментами.

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

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

-sdy-convert-global-to-local

Преобразует программу SDY из глобальных фигур в локальные фигуры.

Преобразует программу SDY из глобальных структур в локальные путем разделения логических измерений на основе атрибутов сегментирования.

На этом этапе используется преобразователь типов для сопоставления RankedTensorType с глобальными логическими формами с физическими формами, локальными для устройства.

-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 scatter для 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-remove-sub-axes-in-input-output-shardings

Удаляет подоси в сегментах ввода/вывода.

Некоторые пользователи Shardy ожидают, что входные/выходные данные функции будут иметь сегменты без подосей. Этот проход удаляет подоси и их конечные оси из входных/выходных сегментов. Этот проход обычно выполняется после sdy-update-non-divisible-input-output-shardings чтобы гарантировать, что удаление подосей не приведет к появлению неразделимых сегментов.

-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 ожидают, что входные/выходные данные функции будут делиться нацело/шардиться, чтобы избежать необходимости дополнения тензоров. Распространение может привести к тому, что входные/выходные данные будут иметь неделимые сегменты, поэтому этот проход обновляет их до префикса сегментации наибольшей размерности из исходной сегментации, которая является равномерно сегментированной.