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