-mpmd-convert-sdy-constants

Преобразует операции sdy.constant в stablehlo.constant .

Преобразует любую операцию типа sdy.constant , которая не подлежит свертыванию, в операцию типа stablehlo.constant . Нет причин препятствовать свертыванию констант, поскольку мы удаляем сегментацию из констант в mpmd-convert-sdy-shardings-to-mpmd-types .

-mpmd-convert-sdy-shardings-to-mpmd-types

Перемещает сегментацию из атрибутов операции в типы !mpmd.mesh_tensor .

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

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

-mpmd-enforce-user-shardings

Обеспечивает соблюдение заданного пользователем разделения входных и выходных данных.

Обеспечивает, чтобы входные и выходные сегменты фрагментов, принимающих аргументы функции или выдающих результаты функции соответственно, соответствовали сегментам, указанным пользователем, то есть входным и выходным сегментам функции.

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

Предварительное условие:

• Распределение пользователей по сегментам задается в аргументах и ​​результатах функции в виде атрибутов.
• Размер фрагментов данных задается атрибутами in_shardings и out_shardings .

-mpmd-extract-reshards-from-inter-mesh-transfers

Перемещает перераспределение данных SPMD между узлами сетки внутрь фрагмента.

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

Это необходимо, поскольку среды выполнения MPMD имеют ограничения в отношении поддерживаемого перераспределения сегментов.

Этот проход применяется только к функциям MPMD в глобальном представлении и с однородной топологией.

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

-mpmd-populate-unreduced-out-sharding

Заполняет несокращенные фрагменты данных.

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

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

Это необходимо, потому что конвейер распространения SDY не заполняет исходящие шарды без редукции, а только входящие. Поскольку технически Shardy не отвечает за распространение осей без редукции, этот проход выполняется в MPMD, а не в SDY.

-mpmd-simplify-program

Удаляет избыточные аргументы/результаты из фрагментов.

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