Постоянная автонастройка (только графический процессор)

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

Если слияний много, это может занять много времени, поэтому мы предоставляем способ загрузки результатов автонастройки, при этом остальные этапы компиляции выполняются в обычном режиме. Кэши автонастройки остаются полезными, даже если мы внесем несколько изменений: слияния, присутствующие в кэше, будут использовать кэш, а остальные будут автоматически настраиваться в обычном режиме.

--xla_gpu_per_fusion_autotune_cache_dir=your/directory

Используйте и поддерживайте кэш автонастройки для каждого слияния в указанной директории. Для каждого отдельного слияния будет создан один файл.

Главное преимущество этого подхода заключается в том, что вы можете использовать один и тот же каталог кэша для нескольких запусков XLA (разных моделей), и ваш кэш будет увеличиваться с каждым новым обнаруженным слиянием, что ускоряет последующие запуски. Также предусмотрена базовая поддержка одновременного запуска нескольких экземпляров XLA с одним и тем же каталогом кэша.

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

  • Каталог должен существовать до запуска XLA и быть доступен для записи.
  • Аннулирование кэша должно обрабатываться пользователем:
    • Пожалуйста, используйте пустую директорию, если хотите начать работу с пустым кэшем.
  • Проверка версии XLA должна выполняться пользователем:
    • Если вы хотите использовать отдельные кэши для разных версий XLA, пожалуйста, используйте разные каталоги.

Кэширование по умолчанию отключено (если параметр не указан).

Ограничение: Нет гарантии, что этот метод будет хорошо работать в сочетании с другими методами кэширования, описанными ниже.

Альтернативный вариант: Загрузка или сохранение всех результатов из заданного HLO в один файл.

Результаты автонастройки можно выгрузить/загрузить, используя следующие параметры:

--xla_gpu_dump_autotune_results_to=
--xla_gpu_load_autotune_results_from=

Если мы укажем файл с расширением .txt или .textproto, то кэш будет выгружен в формате textproto, в противном случае — в бинарном формате protobuf.

В тестах

Автоматическая настройка параметров также может использоваться в тестах. Рекомендуется использовать её, если тесты очень большие, особенно если производительность тестовой среды ограничена.

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

При проведении теста с использованием постоянной автонастройки

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

  1. Необходимо экспортировать результаты автонастройки из теста, например, указав следующие параметры в команде теста:

    --test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto
--test_sharding_strategy=disabled

    Для корректного получения единого кэша автонастройки для всех тестов необходимо отключить сегментирование.

  2. Затем нам нужно загрузить этот кэш в наш репозиторий кода.

  3. Затем нам нужно добавить кэш в зависимости данных нашего тестового объекта и загрузить его с помощью переменной окружения.

    data = ["test_autotune_cache.textproto"],
env = {"XLA_FLAGS": "--xla_gpu_load_autotune_results_from=" +
                    "$(execpath test_autotune_cache.textproto)"},

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

Также ознакомьтесь с примерами тестов в папке xla/backends/gpu/tests/BUILD :

  • load_autotune_results_using_execpath_test
  • load_autotune_results_from_test_workspace_test
  • dump_autotune_results_to_test_outputs_test

Устаревание кэша

Если в модель вносятся многочисленные изменения, возможно, кэш перестанет содержать все результаты слияния, что замедлит тестирование. В этом случае нам придется заново генерировать кэш автонастройки.

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

Кэш также может устареть, если компилятор XLA будет развиваться и создавать различные варианты слияния данных.