torch.compiler.allow_in_graph

torch.compiler.allow_in_graph(fn) [исходный код]

Указывает фронтенду компилятора (Dynamo) пропустить символьный анализ функции и вместо этого напрямую записать её в граф при обнаружении.

Если вы используете torch.compile() (с backend=”inductor” (по умолчанию)) или torch.export.export() и пытаетесь превратить функцию Python в чёрный ящик на протяжении всей трассировки, не используйте этот API. Вместо этого создайте пользовательский оператор (см. Страница пользовательских операторов PyTorch).

Предупреждение

Если вы обычный пользователь torch.compile (например, применяете torch.compile к модели, чтобы ускорить её работу), скорее всего, вам не нужна эта функция. allow_in_graph() может привести к ошибкам, поскольку пропускает фронтенд компилятора (Dynamo), отвечающий за проверки безопасности (разрывы графа, обработку замыканий и т. д.). Неправильное использование приведёт к незаметным ошибкам, которые будет сложно отлаживать.

Если функция Python не помечена декоратором allow_in_graph, при обычном выполнении torch.compile трассировка проходит через неё. allow_in_graph() изменяет это поведение: фронтенд не трассирует содержимое функции, но бэкенд компилятора по-прежнему проходит через неё. В отличие от этого, пользовательские операторы рассматривают функцию как чёрный ящик на всём протяжении стека torch.compile. В таблице ниже сравниваются эти механизмы.

Механизм	Фронтенд (Dynamo)	Бэкенд (AOTAutograd+Inductor)
без декоратора	трассировка внутри	трассировка внутри
allow_in_graph	непрозрачный вызываемый объект	трассировка внутри
пользовательский оператор	непрозрачный вызываемый объект	непрозрачный вызываемый объект

Один из распространённых случаев применения allow_in_graph() — обходной путь для фронтенда компилятора: если вы знаете, что функция работает с последующими компонентами стека компиляции (AOTAutograd и Inductor), но ошибка в Dynamo не позволяет корректно выполнить её символьный анализ (или если ваш код написан на C/C++ и поэтому не может быть проанализирован Dynamo), эту функцию можно декорировать с помощью allow_in_graph(), чтобы обойти Dynamo.

Требуется, чтобы fn соблюдала следующие ограничения. Их несоблюдение приводит к неопределённому поведению:

• Входные данные fn должны иметь типы, которые можно представить как Proxy в графе FX. Допустимые типы: Tensor/int/bool/float/None/List[Tensor?]/List[int?]/List[float?] Tuple[Tensor?, …]/Tuple[int?, …]/Tuple[float?, …]/torch.dtype/torch.device
• Выходные данные fn должны иметь типы, которые можно представить как Proxy в графе FX (см. предыдущий пункт)
• Все тензоры, используемые внутри fn, должны передаваться непосредственно в качестве входных данных fn (а не захватываться в виде переменных).

См. также nonstrict_trace(), для которого действует немного меньше ограничений на входные данные.

Параметры:

fn – вызываемый объект, представляющий функцию, которую нужно включить в граф. Если fn — список или кортеж вызываемых объектов, декоратор рекурсивно применяет allow_in_graph() к каждой функции и возвращает новый список или кортеж с изменёнными функциями.

Пример:

torch.compiler.allow_in_graph(my_custom_function)


@torch.compile(...)
def fn(x):
    x = torch.add(x, 1)
    x = my_custom_function(x)
    x = torch.add(x, 1)
    return x


fn(...)

Захватит один граф, содержащий my_custom_function().