Мы используем OpenAI Triton для генерации некоторых ядер графического процессора. Triton позволяет генерировать быстрые ядра графического процессора для определенных слияний, но нам необходимо настроить некоторые параметры для каждого такого слияния.
Если слияний много, это может занять много времени, поэтому мы предоставляем способ загрузки результатов автонастройки, при этом остальные этапы компиляции выполняются в обычном режиме. Автонастройка кешей по-прежнему будет полезна, если мы внесем несколько изменений: слияния, присутствующие в кеше, будут использовать кеш, а остальные будут автоматически настраиваться нормально.
Рекомендуется: каталог кэша.
--xla_gpu_per_fusion_autotune_cache_dir=your/directory
Используйте и поддерживайте кэш автонастройки per-fusion в заданном каталоге. Для каждого отдельного слияния будет один файл.
Основное преимущество этого подхода заключается в том, что вы можете использовать один и тот же каталог кэша для нескольких запусков XLA (разных моделей), и ваш кэш будет расти с каждым новым объединением, что ускоряет последующие запуски. Существует также базовая поддержка одновременного запуска нескольких экземпляров XLA с одним и тем же каталогом кэша.
XLA будет читать существующие результаты, когда они необходимы, и записывать новые результаты после того, как они будут определены.
- Каталог должен существовать до запуска XLA и быть доступен для записи.
- Аннулирование кеша должно обрабатываться пользователем:
- Пожалуйста, используйте пустой каталог, если вы хотите начать с пустого кеша.
- Проверку версии XLA должен выполнять пользователь:
- Если вы хотите использовать отдельные кэши для разных версий XLA, используйте разные каталоги.
Кэш отключен по умолчанию (если вы не указываете параметр).
Ограничение: не гарантируется хорошая работа в сочетании с другим методом кэширования, описанным ниже.
Альтернатива: загрузка или сброс всех результатов из данного HLO в один файл.
Результаты автонастройки можно выгрузить/загрузить, используя следующие параметры:
--xla_gpu_dump_autotune_results_to=
--xla_gpu_load_autotune_results_from=
Если мы укажем файл .txt или .textproto, то кеш будет сброшен в формате textproto, в противном случае — в двоичном формате protobuf.
В тестах
Постоянную автонастройку также можно использовать в тестах. Рекомендуется использовать его, если тесты очень большие, особенно если производительность тестовой среды ограничена.
Это работает хорошо только в том случае, если кэш автонастройки содержит результаты, сгенерированные на графическом процессоре того же типа, на котором выполняются тесты.
Выполнение теста с использованием постоянной автонастройки
А пока давайте предположим, что рассматриваемый тест всегда использует один и тот же тип графического процессора.
Нам нужно экспортировать результаты автонастройки из теста, например указав эти параметры в команде теста:
--test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto --test_sharding_strategy=disabledШардинг необходимо отключить, чтобы правильно получить единый кеш автонастройки для всех тестов.
Затем нам нужно загрузить этот кеш в наш репозиторий кода.
Затем нам нужно добавить кеш к зависимостям данных нашей тестовой цели и загрузить его с помощью переменной среды.
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 будет развиваться и генерировать различные слияния.