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

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

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

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

--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/service/gpu/tests/BUILD :

  • load_autotune_results_using_execpath_test
  • load_autotune_results_from_test_workspace_test
  • dump_autotune_results_to_test_outputs_test

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

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

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

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