torch.compiler.nested_compile_region

torch.compiler.nested_compile_region(fn=None, *, options=None, max_reuse_entries=8, reuse_hash_fn=None) [исходный код]

Сообщает ``torch.compile``, что отмеченный набор операций образует вложенную область компиляции (которая часто повторяется в полной модели), код которой можно скомпилировать один раз и безопасно использовать повторно. nested_compile_region также можно использовать в качестве декоратора.

Во время трассировки ``torch.compile`` компилятор применяет иерархическую компиляцию с nested_compile_region: при первом обнаружении отмеченной области он генерирует оптимизированный код, а при каждом последующем вызове повторно использует (или «вставляет») ранее скомпилированный код. Это может значительно сократить общее время компиляции для глубоко вложенных структурно идентичных компонентов, таких как слои трансформера большой языковой модели (LLM).

Вне контекста torch.compile — то есть при стандартном eager-выполнении — вызов ничего не делает, поэтому существующие рабочие процессы остаются без изменений.

Обратите внимание: nested_compile_region не гарантирует, что область будет скомпилирована ровно один раз. Если компилятор обнаружит, что новые условия входных данных (форма, dtype, устройство, шаг, глобальные переменные и т. д.) делают недействительным повторное использование кешированной версии, он прозрачно перекомпилирует область. Поэтому использовать эту возможность безопасно: корректность всегда сохраняется, а дополнительные затраты на компиляцию возникают только при необходимости.

Параметры:

• fn – Функция, которую нужно обернуть
• options (NestedCompileRegionOptions | None) – Необязательный бэкенд для компиляции подграфа. Предупреждение: эта экспериментальная функция находится в разработке и пока не готова к использованию.
• max_reuse_entries (int) – Максимальное количество записей кеша повторного использования для каждой функции, прежде чем будет выдана ошибка. При достижении этого ограничения проверки продолжают завершаться неудачей при последующих вызовах, и иерархическая компиляция неэффективна.
• reuse_hash_fn – Необязательная вызываемая функция, принимающая те же *args, **kwargs, что и обёрнутая функция, и возвращающая целочисленный хеш-ключ. Если она указана, Dynamo трассирует эту функцию, чтобы получить целочисленную константу, и использует её в качестве ключа кеша для повторного использования подграфа, обходя автоматический механизм создания отпечатков и проверок. Два вызова, возвращающие одинаковый хеш-ключ, используют один и тот же кешированный подграф. Функция хеширования должна полностью поддерживать трассировку (без разрывов графа) и возвращать целочисленную константу.