Использование torch._dynamo.nonstrict_trace

Создано: 28 июл. 2025 | Последнее обновление: 03 дек. 2025

Краткое описание:

• Используйте nonstrict_trace, чтобы трассировать функцию с помощью неточного режима трассировки внутри области torch.compile. Это может понадобиться, если граф Dynamo прерывается на каком-либо участке функции и вы уверены, что функция поддерживает неточную трассировку.

Рассмотрим следующий сценарий:

def get_magic_num():
    # This explicit graph break call is meant to emulate any kind of Dynamo
    # graph break, e.g., the function is implemented in C, or uses some python
    # language feature Dynamo doesn't yet support.
    torch._dynamo.graph_break()
    return torch.tensor([42])
@torch.compile(fullgraph=True)
def func(x):
    n = get_magic_num()
    return x + n
try:
    func(torch.rand(10))
except Exception as e:
    print(e)

При выполнении приведённого выше кода Dynamo выдаст ошибку, поскольку обнаружит разрыв графа, хотя пользователь указал fullgraph=True.

В таких ситуациях, если пользователь всё же хочет сохранить fullgraph=True, обычно у него есть несколько вариантов:

1. Разрыв графа вызван языковой конструкцией, которую Dynamo пока не поддерживает. В этом случае пользователь может переписать код или создать задачу на GitHub.
2. Разрыв графа вызван вызовом функции, реализованной на C. В этом случае пользователь может попробовать использовать пользовательскую операцию. Также можно попробовать предоставить полифилл (эталонную реализацию на Python), чтобы Dynamo мог выполнить трассировку этой функции.
3. В худшем случае возникает внутренняя ошибка компилятора. В этом случае пользователю, скорее всего, придётся создать задачу на GitHub.

Помимо всех этих вариантов, PyTorch предлагает альтернативный способ torch._dynamo.nonstrict_trace, если вызвавший разрыв графа вызов функции соответствует определённым требованиям:

• Требования общей неточной трассировки.
• Входные и выходные данные должны содержать либо базовые типы (например, int, float, list, dict, torch.Tensor), либо пользовательские типы, зарегистрированные в torch.utils._pytree.
• Функция должна быть определена вне области torch.compile.
• Все значения, не являющиеся входными и считываемые функцией, будут считаться константами (например, глобальный тензор), и для них не будут создаваться защитные проверки.

При трассировке вызова функции, помеченной torch._dynamo.nonstrict_trace, torch.compile переключается на неточную трассировку, и в конечном итоге граф FX будет содержать все необходимые тензорные операции, выполненные внутри этой функции.

В приведённом выше примере мы можем использовать torch._dynamo.nonstrict_trace to eliminate, чтобы устранить разрыв графа:

@torch._dynamo.nonstrict_trace
def get_magic_num():
    # This explicit graph break call is meant to emulate any kind of Dynamo
    # graph break, e.g., the function is implemented in C, or uses some python
    # language feature Dynamo doesn't yet support.
    torch._dynamo.graph_break()
    return torch.tensor([42])
@torch.compile(fullgraph=True)
def func(x):
    n = get_magic_num()
    return x + n
print(func(torch.rand(10)))
# No graph break and no error.

Обратите внимание, что это можно использовать и внутри области torch.compile:

def get_magic_num():
    # This explicit graph break call is meant to emulate any kind of Dynamo
    # graph break, e.g., the function is implemented in C, or uses some python
    # language feature Dynamo doesn't yet support.
    torch._dynamo.graph_break()
    return torch.tensor([42])
@torch.compile(fullgraph=True)
def func(x):
    n = torch._dynamo.nonstrict_trace(get_magic_num)()
    return x + n
print(func(torch.rand(10)))
# No graph break and no error.