-mpmd-delay-inferred-fragments
Задержки подразумевали выполнение фрагментов как можно позже.
В результате выполнения фрагментов предполагалось, что они будут обработаны как можно позже, то есть непосредственно перед первым пользователем.
-mpmd-delay-transfers-from-cpu
Задерживает выполнение передач данных с ЦП до максимально возможной задержки.
Перемещает операции передачи данных с процессора на устройство непосредственно перед их первым потребителем. Это означает отсрочку выделения памяти для передаваемых буферов, что может быть полезно для использования HBM.
-mpmd-lower-to-fragment-calls
Преобразует фрагменты MPMD в вызовы фрагментов и функции.
Заменяет все фрагменты вызовами фрагментов.
На этом этапе создается функция для каждой группы фрагментов, имеющих одинаковую форму тела и сетки, с именем первого встреченного фрагмента в группе в качестве имени символа, и добавляется она в таблицу символов.
Тело каждой функции извлекается из тела первого встреченного фрагмента в соответствующей группе, а форма сетки берется из топологии по имени сетки этого фрагмента. Если фрагмент имеет атрибуты аргументов, касающиеся псевдонимов ввода-вывода, они будут присвоены атрибутам аргументов пониженной функции.
Поскольку функции должны иметь уникальные имена, на этом этапе к имени всех функций, кроме первой, имеющих то же исходное имя, добавляется индекс, то есть i-я функция с именем "some_name" для i > 0 будет иметь имя "some_name_i".
Параметры
-verbose-logging : Whether to enable verbose logging
-mpmd-mark-aliasing-and-donation
Каждый фрагмент помечается псевдонимом или информацией о пожертвовании.
Устанавливает атрибут arg_attrs для операций фрагментации, если любой из их входных сигналов может быть сопоставлен с выходным сигналом или передан в качестве донора. Каждый входной сигнал, который может быть сопоставлен, будет иметь атрибут tf.aliasing_output . В противном случае — атрибут jax.buffer_donor = true . Например, {arg_attrs =\[{tf.aliasing_output = 0 : i32}, {jax.buffer_donor = true}, {}\]} показывает, что первый вход может быть псевдонимизирован с выходом 0, второй вход может быть передан таким образом, чтобы XLA обнаружил псевдонимизированный выход, а третий вход не может быть ни псевдонимизирован, ни передан.
-mpmd-mark-fragment-reserved-memory
Укажите для каждого фрагмента объем памяти, который необходимо зарезервировать для компиляции.
Присваивает каждому фрагменту атрибут xla_tpu_user_reserved_hbm_bytes , который указывает XLA, сколько байтов следует сохранить во время компиляции каждого фрагмента. Отслеживая активные тензоры на сетке, XLA будет знать фактическое минимальное использование памяти во время выполнения, и мы можем предотвратить применение оптимизаций в исполняемом файле, которые увеличили бы использование памяти сверх возможностей устройства. ПРИМЕЧАНИЕ: этот проход предполагает, что фрагменты выполняются в порядке следования программ.
-mpmd-mark-input-output-with-layouts
Передает параметры компоновки из аргументов/результатов функции в аргументы/результаты фрагмента.
Передает атрибуты mhlo.layout_mode от входных данных программы к фрагментам, которые являются потребителями входных данных программы, и передает атрибуты mhlo.layout_mode от выходных данных программы к фрагментам, которые являются производителями выходных данных. Если возвращается аргумент программы, то mhlo.layout_mode передается в/из результатов программы. Входные данные, выходные данные и аргументы/результаты фрагментов, связанные с операцией передачи, всегда перезаписываются макетом DEFAULT, поскольку передачи поддерживают только макеты по умолчанию.
Если для аргумента программы задана компоновка AUTO и он используется в нескольких фрагментах, то мы устанавливаем для него компоновку DEFAULT, чтобы обеспечить единообразную компоновку во всех фрагментах.
Если для выходных данных программы, являющихся результатом работы фрагмента, установлен режим автоматической компоновки (AUTO) и они используются в других фрагментах в качестве входных данных, мы также устанавливаем для них режим компоновки по умолчанию (DEFAULT), чтобы обеспечить единообразную компоновку во всех фрагментах.
-mpmd-mark-offloaded-input-output
Помечает выгруженные входные и выходные значения, чтобы компилятор знал, что они находятся в памяти хоста.
Помечает аргументы и результаты фрагмента атрибутами, чтобы определить, какие значения находятся в памяти хоста, и чтобы эта информация могла быть использована XLA. Также помечает аргументы и результаты функции точки входа, чтобы Pathways мог использовать эту информацию.
-mpmd-validate-no-reshards
Проверяет, что фрагментов, содержащих только данные для повторного оцифровки, не существует.
Фрагмент, содержащий только перераспределение сегментов, — это фрагмент, который содержит только результат mpmd.return одного из своих аргументов. Такие фрагменты обычно указывают на неожиданное перераспределение сегментов.
По умолчанию этот проход выдает предупреждение для каждого такого фрагмента. Если задан параметр fail-on-reshard-only-fragments , он выдает ошибку и завершается с ошибкой.
Параметры
-fail-on-reshard-only-fragments : Whether to emit an error (and fail) instead of a warning.