-sdy-close-shardings

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

-sdy-constant-or-scalar-merger

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

Выполняет упрощенное устранение общих подвыражений для констант с идентичным сегментированием.

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

-sdy-convert-global-to-local

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

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

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

Параметры

-per-dim-all-gather                     : Keep per-dimension all-gather without combining them into a single all-gather.
-combine-multi-dimension-reduce-scatter : Combine multi-dimension reduce-scatter into a single reduce-scatter.
-enable-rgv3                            : Use StableHLO ReplicaGroupV3 (mesh-axes based) for collectives.
-replica-count                          : Number of replicas (data parallelism).
-partition-count                        : Number of partitions (model parallelism).

-sdy-drop-sharding-and-mesh

Удаляет из программы операцию mesh и нотацию сегментирования.

Параметры

-keep-module-sharding-metadata : Whether to keep module-level sharding metadata (sdy.parameters_shardings and sdy.output_shardings).

-sdy-drop-sharding-rules

Выпуски OpShardingRuleAttr из всех зарегистрированных операций.

-sdy-export-named-computations

Перескажи звонки от пользователя NamedComputationOp.

Создает проход, который преобразует NamedComputationOp в CallOp с новой частной функцией, называемой name NamedComputationOp. Новые FuncOp и CallOp имеют такое же сегментирование, как и исходные операнды/результаты NamedComputationOp.

Если в модуле есть функция с тем же названием, что и NamedComputationOp, таблица символов MLIR изменит его на {name}_#.

-sdy-insert-explicit-reshards

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

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

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

Обратите внимание, что когда ось (или субось) используется для сегментирования несоответствующих измерений (например, несокращающихся измерений в 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.
-mark-partial-result-with-unreduced-axes : Whether to mark partial reduction results with unreduced axes.

-sdy-insert-func-call-reshards

Вставляет повторные шарды для конфликтов шардинга функций и вызовов.

Вставляет повторные сегменты для конфликтов сегментирования функций и вызовов в результатах.

-sdy-optimize-collectives

Оптимизирует коллективные коммуникационные операции.

Устраняет избыточные операции коллективной перестановки, предшествующие цепочкам "все ко всем", путем декомпозиции разделенных измерений и объединения в комбинированную операцию "все ко всем".

-sdy-pad-for-divisibility

Дополняет тензоры с неделимыми сегментами до делимых форм.

-sdy-per-instruction-partitioning

_Выборочно извлекает целевые инструкции, запускает для них конвейер разбиения, чтобы сгенерировать локальный код устройства, и упаковывает их в sdy.manualcomputation.

Для отладки и бисекции выборочно разбивает отдельные инструкции, используя автономный конвейер разбиения Shardy (resolve-permutation-factors, reshard-to-collectives, pad-for-divisibility, convert-global-to-local), и упаковывает полученный локальный код устройства в sdy.manual_computation. Это позволяет сохранить остальную часть глобальной программы, поэтому XLA SPMD разделяет только развернутые глобальные инструкции.

Параметры

-filter                   : Filter string for selective partitioning. Can be empty (all sharded ops), comma-separated op name substrings (e.g. 'dot, pad'), or key-values like 'selectLow=0, selectHigh=10, func=subroutine'.
-enable-halo-exchange     : Implement halo exchange logic for windowed operations inside the sub-pipeline.
-replica-count            : Number of replicas (data parallelism).
-partition-count          : Number of partitions (model parallelism).
-rng-bit-generator-unsafe : Whether to allow unsafe partitioning of stablehlo.rng_bit_generator.

-sdy-propagate-to-func-results

Распространение сегментов от терминатора функции до результатов функции.

Копирует сегментирование значений терминатора функции в соответствующие результаты func.func, сохраняя существующее сегментирование основной функции.

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

_Удаляет sdy.all_gather и sdy.reducescatter для CMV1.

Удаляет all-gather в шаблоне all-gather + точка. Удаляет reduce-scatter в шаблоне dot + reduce-scatter. Этот проход предназначен для совместимости с коллективным умножением матриц версии 1 (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, &q>uot;y"=<2, ">;z"=2
%0 <: tensor16x2xf32 {sdy.sharding>@mesh, \[{"x",< "y", &quo>t;z"<}, {}\]<>/span>
%1 = sdy.reshard %arg0 @mesh, \[{"x"}, {}\] : tensor16x2xf32

Вывод:

mesh = <"x"=2, &q>uot;y"=<2, ">;z"=2
%0 <: tensor16x2xf32 {sdy.sharding>@mesh, \[{"x", "y", "z"}, {}\]<
%1 = sdy.all_gather> \[{&quo<t;y">;, "z"}, {}\] %arg0 out_sharding=@mesh, \[{"x"}, {}\] : tensor16x2xf32

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

Параметры

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

-sdy-resolve-permutation-factors

Устраняет сегментирование по параметрам, сопоставленным с факторами kPermutation.

При сегментировании параметров с факторами kPermutation может потребоваться обмен данными между устройствами (например, обмен ореолами для окон или коллективные перестановки для обратных преобразований).

Если для параметра enableHaloExchange задано значение true, при создании карты будет использоваться доступная оптимизированная логика связи для разрешения факторов перестановки. В противном случае проход просто вставит операции sdy.reshard для репликации этих измерений. По умолчанию для параметра enableHaloExchange задано значение true.

Параметры

-enable-halo-exchange     : Implement halo exchange logic for windowed operations.
-replica-count            : Number of replicas (data parallelism).
-partition-count          : Number of partitions (model parallelism).
-rng-bit-generator-unsafe : Whether to allow unsafe partitioning of stablehlo.rng_bit_generator.

-sdy-resolve-single-device-sharding

Преобразует операции на одном устройстве в stablehlo.if, защищенные идентификатором целевого устройства.

Параметры

-replica-count   : Number of replicas per partition.
-partition-count : Number of partitions per replica.

-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-sink-func-data-flow-edges

Поглощает все FuncDataFlowEdgeOp.

Перемещает сегментирование каждого FuncDataFlowEdgeOp на его входные данные и заменяет операцию входными данными.

-sdy-unflatten-call-graph

Разворачивает граф вызовов.

Разворачивает график. Она удаляет дубликаты функций с одинаковыми входными и выходными сегментами и одинаковым источником, указанным в атрибуте original_func_name, прикрепленном к функциям.

Параметры

-dedup-functions-fully : If true, regardless of the input and output shardings of functions, it keeps one callee function for each caller function. The default is false, meaning it will deduplicate only if the input and output shardings are the same.

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

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

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

-sdy-verify-unreduced-axes

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

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