-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
Перемещает шардинги из op attrs в типы !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 в глобальном представлении и с однородной топологией.
Условие: все шардинги указаны как атрибуты операции, а не в типах.