Снижение накладных расходов на проверки

1. Сократите время выполнения байт-кода до графа с помощью install_free_tensors

import torch._dynamo

with torch._dynamo.config.patch(install_free_tensors=True):
    # torch.compile your model / call the compiled function here
    ...

Этот параметр устанавливает свободные тензоры (например, параметры и буферы) в качестве атрибутов графа, а не входных данных графа, что сокращает объём работы байт-кода до графа, необходимой для их передачи в граф. В качестве побочного эффекта также меняется способ проверки этих параметров и буферов (для них больше не требуется сопоставление тензоров для каждого входного значения), поэтому накладные расходы на проверки могут немного снизиться. Само по себе это снижение часто невелико, поскольку большая часть затрат на каждый вызов приходится на доступ к проверяемому объекту, а не на саму проверку.

По умолчанию install_free_tensors имеет значение False (изначально этот параметр был создан для упрощения экспорта и получения стабильного количества входных данных графа), поэтому его нужно включить с помощью приведённого выше патча конфигурации.

2. Пропускайте проверки для nn.Modules с помощью guard_filter_fn

torch.compile принимает параметр guard_filter_fn, определяющий для каждой проверки, следует ли её сохранить. В torch.compiler предусмотрено несколько готовых фильтров. Наибольший эффект для сокращения накладных расходов на проверки даёт пропуск проверок атрибутов nn.Modules:

import torch

opt_mod = torch.compile(
    mod,
    options={"guard_filter_fn": torch.compiler.skip_guard_on_all_nn_modules_unsafe},
)

Обычно это значительно сокращает накладные расходы на проверки, поскольку модели, как правило, содержат множество проверяемых атрибутов модулей, которые никогда не изменяются. Обратите внимание на взаимодействие с байт-кодом до графа: пропуск этих проверок приводит к сбросу кэша указателей параметров и буферов, поэтому часть устранённых накладных расходов на проверки переносится в байт-код до графа. Чтобы получить полный эффект, сочетайте этот способ с install_free_tensors=True выше — применяйте оба способа (1) и (2) вместе.

Другие доступные фильтры в torch.compiler (все _unsafe, с теми же оговорками):

• skip_guard_on_inbuilt_nn_modules_unsafe — пропускает проверки только для встроенных модулей, например torch.nn.Linear.
• skip_guard_on_globals_unsafe — пропускает проверки глобальных переменных.
• keep_tensor_guards_unsafe — оставляет только проверки тензоров (при необходимости включая параметры).
• keep_portable_guards_unsafe — оставляет только проверки, переносимые между средами Python и без Python (состояния глобальных переменных, среды форм и неглобальные проверки тензоров).
• skip_all_guards_unsafe — удаляет все проверки и устраняет все гарантии безопасности; используйте с крайней осторожностью.

3. Сначала попробуйте use_recursive_dict_tags_for_guards

import torch._dynamo

torch._dynamo.config.use_recursive_dict_tags_for_guards = True

Это ускоряет выполнение проверок для вложенных nn.Modules: рекурсивно проверяются метки словарей, чтобы избежать выполнения полного набора проверок. Механизм основан на достаточно сложном использовании низкоуровневых возможностей CPython. В issues проекта OSS поднимались некоторые вопросы, поэтому по умолчанию этот параметр отключён, и его использование может быть пересмотрено. Стоит попробовать его до параметров (1) и (2): если он подходит для вашей модели, вам не придётся пропускать проверки nn.Module (2), хотя install_free_tensors (1) всё равно будет полезен.

4. Пропускайте оценку проверок после прогрева с помощью skip_guard_eval_unsafe

import torch

# 1. Warm up: run the compiled model with a sufficient variety of inputs until
#    no further recompilation occurs.
# 2. Then switch stance to run only the minimal set of differentiating guards.
with torch.compiler.set_stance(skip_guard_eval_unsafe=True):
    # steady-state inference / training iterations
    ...

После того как вы прогрели скомпилированную модель и убедились, что повторных компиляций больше не происходит, skip_guard_eval_unsafe выполняет только минимальный набор проверок, необходимый для различения уже имеющихся скомпилированных артефактов, пропуская остальные. В отличие от описанных выше параметров, этот параметр нельзя задать во время torch.compile — его нужно включить в цикле обслуживания/обучения после прогрева. Если предположение об отсутствии дальнейших компиляций окажется неверным (поступят действительно новые входные данные), существует риск незаметно получить неверные результаты, отсюда и unsafe в названии.

Итоговые рекомендации

• Сначала измерьте, на что уходит время при каждом вызове (на проверки или байт-код до графа), с помощью tlparse / TORCH_LOGS=guards.
• При желании сначала попробуйте use_recursive_dict_tags_for_guards=True.
• В противном случае используйте вместе install_free_tensors=True и guard_filter_fn, например skip_guard_on_all_nn_modules_unsafe: один способ сокращает накладные расходы на проверки, а другой не даёт этим сэкономленным затратам вновь появиться в виде накладных расходов на байт-код до графа.
• Для обслуживания в установившемся режиме после прогрева рассмотрите set_stance(skip_guard_eval_unsafe=True).

Дополнительная литература

Более подробное объяснение проверок — единицы компиляции (граф и соответствующий набор проверок), причины большого размера этого набора и примеры профилирования для каждого из описанных выше методов — см. в публикации для разработчиков Внутреннее устройство проверок torch.compile.