Используйте fullgraph=True для выявления и устранения разрывов графа

Стратегия 1: перепишите неподдерживаемый код, используя возможности, поддерживаемые Dynamo

Во многих сообщениях об ошибках, связанных с разрывом графа, приводятся рекомендации по переписыванию кода, чтобы избежать разрыва. Если устранить разрыв графа по-прежнему сложно, перейдите к следующей стратегии или создайте задачу в репозитории PyTorch на GitHub.

Дополнительные примеры разрывов графа и способы их устранения приведены в разделе Распространённые разрывы графа.

Пример: Dynamo не поддерживает вызов next для объекта list_iterator, который был входным аргументом функции, проходящей компиляцию.

@torch.compile(fullgraph=True)
def f(xs):
    a = next(xs)
    b = next(xs)
    return a + b

xs = [torch.tensor(1.), torch.tensor(2.)]
try:
    out = f(iter(xs))
except Exception as e:
    print(e)

Вместо этого перепишите компилируемую функцию так, чтобы она принимала список.

@torch.compile(fullgraph=True)
def f_rewritten(xs):
    it = iter(xs)
    a = next(it)
    b = next(it)
    return a + b

f_rewritten(xs)

Стратегия 2: чистые функции всегда можно компилировать с помощью обходного механизма.

Краткое описание: пространство всех функций Python огромно, поэтому Dynamo нецелесообразно трассировать каждую функцию Python без разрывов графа. Для функций Python, которые считаются «чистыми» и которые Dynamo не может трассировать без разрывов графа, предусмотрены обходные механизмы, позволяющие попытаться трассировать такие функции:

1. Используйте custom_op или triton_op для чистых ядер Triton.
2. Используйте nonstrict_trace для чистых функций, использующих только операции с тензорами PyTorch.
3. Используйте custom_op для всех остальных чистых функций.

«Чистая функция» обладает следующими свойствами:

• Детерминированность. При одинаковых входных данных чистая функция всегда возвращает одинаковый результат.
• Отсутствие внешних побочных эффектов. Чистая функция не имеет побочных эффектов, видимых извне, таких как изменение внешнего состояния или выполнение операций ввода-вывода. Допускаются побочные эффекты, остающиеся внутри функции (например, изменение промежуточных тензоров). Одно примечательное исключение: обычно допускаются операции torch.*, изменяющие входные тензоры функции.
• Явные входные и выходные данные. Все входные данные должны передаваться через параметры функции, а все результаты должны возвращаться функцией.

Примеры см. в разделе Чистые функции.

Теоретически Dynamo способен обрабатывать широкий спектр нечистых функций, но для некоторых возможностей языка Python поддержка может быть недостаточной. Однако чистые функции всегда можно компилировать с помощью обходного механизма.

Если у вас возник разрыв графа, возможно, код вокруг него удастся переработать в чистую функцию и использовать обходной механизм, позволяющий обойти трассировку Dynamo:

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

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

Пример: если Dynamo не поддерживает некоторую возможность или API Python, допускающие нестрогую трассировку (например, использующие операции PyTorch), используйте torch._dynamo.nonstrict_trace для их захвата.

# this is a function that Dynamo doesn't support (due to the graph_break() call).
def g(x):
    y = x.sin()
    torch._dynamo.graph_break()
    z = y.sin()
    return z

@torch.compile(fullgraph=True)
def f(x):
    w = x.sin()
    return g(w)

x = torch.randn(3)
try:
    f(x)  # Graph Break: there was a call to torch._dynamo.graph_break()
except Exception as e:
    print(e)

@torch.compile(fullgraph=True)
def f_rewritten(x):
    w = x.sin()
    return torch._dynamo.nonstrict_trace(g)(w)
f_rewritten(x)  # works

Пример: используйте пользовательские операторы, чтобы создавать функции, непрозрачные для torch.compile

from torch.utils.cpp_extension import load_inline

# C++ source code for the square operation
cpp_source = """
torch::Tensor square_cpu(torch::Tensor input) {
    // Check that input is a CPU tensor
    TORCH_CHECK(input.device().is_cpu(), "Input must be a CPU tensor");

    // Create output tensor with same shape and dtype as input
    torch::Tensor output = torch::empty_like(input);

    // Get data pointers
    float* input_data = input.data_ptr<float>();
    float* output_data = output.data_ptr<float>();

    // Get total number of elements
    int64_t numel = input.numel();

    // For loop to compute square of each element
    for (int64_t i = 0; i < numel; i++) {
        output_data[i] = input_data[i] * input_data[i];
    }

    return output;
}
"""

# Load the extension inline
square_module = load_inline(
    name="square_cpu_kernel",
    cpp_sources=cpp_source,
    functions=["square_cpu"],
    verbose=True
)

def square(x):
    return square_module.square_cpu(x)

@torch.compile(fullgraph=True)
def f(x):
    return square(x)

try:
    f(torch.randn(3, 3))  # graph break
except Exception as e:
    print(e)

# Use torch.library.custom_op to define a new custom operator.
# Custom operators are opaque with respect to torch.compile:
# that is, torch.compile does not peek into them.

@torch.library.custom_op("mylib::square", mutates_args=())
def square(x: torch.Tensor) -> torch.Tensor:
    return square_module.square_cpu(x)

# Use register_fake to add a ``FakeTensor`` kernel for the operator
@square.register_fake
def _(x):
    return x.new_empty(x.size())

print(f(torch.randn(3, 3)))  # no graph break

Дополнительные сведения о triton_op для пользовательских ядер Triton см. в руководстве по пользовательским ядрам Triton.

Стратегия 3: не компилируйте код

Не весь код подходит для компиляции. torch.compile — это компилятор вычислений с тензорами; он не сможет оптимизировать такие операции, как ввод-вывод с диска. Попробуйте переработать код так, чтобы неподдерживаемый код не вызывался в компилируемой области.

@torch.compile(fullgraph=True)
def f(x):
    y = x ** 2  / 2
    torch.save(y, "foo.pt")
    z = y ** 3 / 6
    return z

x = torch.randn(3)
try:
    f(x)  # Graph Break: torch.save not supported
except Exception as e:
    print(e)

def f_rewritten(x):
    y = g(x)
    torch.save(y, "foo.pt")
    z = h(y)
    return z

@torch.compile(fullgraph=True)
def g(x):
    y = x ** 2  / 2
    return y

@torch.compile(fullgraph=True)
def h(y):
    z = y ** 3 / 6
    return z

f_rewritten(x)

Если проблемную функцию не нужно выполнять во время компиляции, рассмотрите возможность использовать torch.compiler.is_compiling(), чтобы пропустить её.

@torch.compile(fullgraph=True)
def f(x):
    y = x ** 2  / 2
    if not torch.compiler.is_compiling():
        torch.save(y, "foo.pt")
    z = y ** 3 / 6
    return z

x = torch.randn(3)
f(x)  # torch.save is not called

Если функция вызывается во многих местах и вы не возражаете против того, чтобы torch.compile всегда пропускал её, добавьте её в torch._dynamo.config.ignore_logging_functions.

def bad_fn(y):
    torch.save(y, "foo.pt")

torch._dynamo.config.ignore_logging_functions.add(bad_fn)

@torch.compile(fullgraph=True)
def f(x):
    y = x ** 2  / 2
    bad_fn()
    z = y ** 3 / 6
    return z

x = torch.randn(3)
f(x)  # torch.save is not called

Обратите внимание, что существуют ограничения на типы функций, которые можно добавлять в ignore_logging_functions. В частности:

• Функция может принимать любые аргументы, но ДОЛЖНА возвращать None.
• Функции должны быть функциями на уровне модуля, logging.Logger.<method> (игнорирует этот метод для всех экземпляров logging.Logger) или logger_obj.<method> (игнорирует этот метод только для конкретного экземпляра logger_obj).

Другие функции могут игнорироваться или не игнорироваться в зависимости от деталей реализации. Если вы хотите игнорировать функцию, но ignore_logging_functions не удаётся её игнорировать, создайте задачу.