Распространённые разрывы графа

Некорректный код

В вашем коде могут быть ошибки (то есть он не выполняется даже без torch.compile). В примере ниже в вызове torch.sin допущена опечатка из-за лишнего аргумента. Всегда отключайте torch.compile, чтобы проверить, правильно ли выполняется код.

@torch.compile
def fn(x):
    y = torch.sin(x, x)
    return y

try:
    fn(torch.ones(3, 3))
except Exception as e:
    pass

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

Также можно использовать torch.compiler.set_stance("force_eager"), чтобы быстро отключить torch.compile, не изменяя вызов torch.compile:

@torch.compile
def fn(x):
    y = torch.sin(x, x)
    return y

try:
    with torch.compiler.set_stance("force_eager"):
        fn(torch.ones(3, 3))
except Exception as e:
    print(e)

Другие примеры использования set_stance для отладки см. на странице https://docs.pytorch.org/tutorials/recipes/torch_compiler_set_stance_tutorial.html#crashing-sooner.

Операции, зависящие от данных

torch.compile вызывает разрыв графа при операциях, зависящих от данных, например при управлении потоком, зависящем от данных (условные операторы if, циклы с тензорами), и прямом доступе к данным тензора (.item, .data_ptr).

@torch.compile
def fn(x):
    y = x.sum()
    if y > 0:
        return x + y.item()
    return x - y.item()

print(fn(torch.ones(3, 3)))

Общий способ обойти такие разрывы графа — избегать операций, зависящих от данных. Ниже приведены некоторые конкретные способы обхода:

• Если управление потоком на самом деле не зависит от значений данных, измените код так, чтобы управление потоком выполнялось на основе констант.

# old
x = torch.randn(3, 3)
@torch.compile
def fn(y):
    if x.sum() > 0:
        return y + x
    else:
        return y - x

print(fn(torch.ones(3, 3)))

# new
x = torch.randn(3, 3)
cond = (x.sum() > 0).item()
@torch.compile
def fn(y):
    if cond:
        return y + x
    else:
        return y - x

print(fn(torch.ones(3, 3)))

• Вместо управления потоком, зависящего от данных, используйте операции высшего порядка, например Управление потоком — Cond

# old
@torch.compile
def fn(x):
    if x.sum() > 0:
        return x + 1
    return x - 1

print(fn(torch.ones(3, 3)))

# new
@torch.compile
def fn(x):
    return torch.cond(
        x.sum() > 0,
        lambda x: x + 1,
        lambda x: x - 1,
        (x,),
    )

print(fn(torch.ones(3, 3)))

• Если вы используете вызов .item(), попробуйте torch._dynamo.config.capture_scalar_outputs = True или TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1.
• Оберните проблемные части функции в пользовательский оператор

Вывод и ведение журналов

Вывод, ведение журналов и выдача предупреждений приводят к разрыву графа. Можно попробовать обойти эту проблему с помощью torch._dynamo.config.reorderable_logging_functions. Эта настройка используется для переноса функций ведения журналов в конец трассируемой функции, что позволяет избежать разрыва графа. Однако содержимое журналов может отличаться, например, если происходит мутация.

Примечание: у reorderable_logging_functions есть ограничения: эти функции должны возвращать None, а их аргументы должны быть тензорами, константами или строками форматирования.

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