Устранение ошибок GuardOnDataDependentSymNode

Ошибки, зависящие от данных

Одна из основных проблем при работе с неподкреплёнными формами — обработка ветвлений: без подсказки компилятор не может определить, какой путь выбрать, и по умолчанию выдаёт ошибку GuardOnDataDependentSymNode.

Шаг 1. Изучите возможный источник проблемы (трассировку стека Python)

Исключение содержит трассировку стека, которая часто указывает на проблему. Поскольку трассировки стека PT2 могут быть длинными, в сообщении об ошибке также предлагается вероятный источник проблемы во фреймворке. Например:

Potential framework code culprit (scroll up for full backtrace):
  File "/data/users/ezyang/a/pytorch/torch/_prims_common/__init__.py", line 855, in infer_size
    if d == -1:

Шаг 2. Изучите трассировку стека C++

Если предполагаемый источник проблемы в коде фреймворка не даёт полезной информации, проверка может находиться в C++. Чтобы получить трассировку стека C++, запустите программу с TORCHDYNAMO_EXTENDED_DEBUG_CPP=1. В результате появится подробная трассировка C++, в которой чередуются кадры Python, CPython и C10/ATen/libtorch. Ищите символы в пространстве имён at:: или c10::, похожие на код для конкретного ядра и, вероятно, связанные с ядром, выполненным согласно трассировке стека Python. Если используется неотладочная сборка PyTorch, встраивание функций может привести к отсутствующим кадрам; тогда для поиска проблемы потребуется изучить исходный код. Например, см. https://github.com/pytorch/pytorch/pull/118579.

Ниже приведён пример трассировки стека C++ из сеанса отладки:

[2024-02-08 08:20:45,259] torch.fx.experimental.symbolic_shapes: [INFO]   File "../
__gen_aten__/out/RegisterCompositeImplicitAutograd.cpp", line 2025, in at::
(anonymous namespace)::(anonymous namespace)
::wrapper_CompositeImplicitAutograd_Tensor_narrow(at::Tensor const&, long,
at::Tensor const&, c10::SymInt) [2024-02-08 08:20:45,259] torch.fx.experimental.
symbolic_shapes: [INFO]   File "../aten/src/ATen/native/TensorShape.cpp", line 1410,
in at::native::narrow_tensor_symint(at::Tensor const&, long, at::Tensor const&,
c10::SymInt) [2024-02-08 08:20:45,259] torch.fx.experimental.symbolic_shapes:
[INFO]   File "../__gen_aten__/out/core/TensorMethods.cpp", line 52, in long
at::Tensor::item<long>() const [2024-02-08 08:20:45,259] torch.fx.experimental.
symbolic_shapes: [INFO]   File "../ATen/core/TensorBody.h", line 4274, in
at::Tensor::item() const

В этом примере at::native::narrow_tensor_symint вызывает item<long>, что приводит к проверке символа SymNode, зависящего от данных.

Подумайте над следующим:

• Имеет ли смысл, что это условие запускает проверку символа, зависящего от данных?
• Если уравнение содержит два разных символа, должны ли мы знать, что они на самом деле равны?
• Можно ли научить этот фрагмент кода обрабатывать входные данные обобщённым способом, который работает для любых форм?

Использование TORCH_LOGS=dynamic и изучение трассировки стека пользовательского кода крайне важны для понимания того, как исправить проблему: они помогают определить, как изменить пользовательскую программу.

[INFO] create_unbacked_symint u0 [-9223372036854775808, 9223372036854775807] (w.py:40 in custom_op_meta)

Это сообщение журнала указывает, где (w.py:40) был выделен неподкреплённый SymInt. Неподкреплённый SymInt может выделяться несколько раз, поэтому отслеживайте равенство этих значений:

[INFO] set_replacement u1 = u0 (trivial_lhs) ValueRanges(lower=0, upper=9223372036854775807, is_bool=False)

Шаг 1. Можно ли переписать код, чтобы использовать общий путь?

Лучшее решение — перестроить код так, чтобы ему вообще не требовалось ветвление по неподкреплённым символам. Спросите себя: Существует ли общий путь кода, работающий для любых форм?

Например, вместо следующего:

i = x.item()
if i > 4:
    return x * 2
else:
    return x + 3

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

from torch.fx.experimental.symbolic_shapes import statically_known_true

# Instead of: if x.numel() > 10:
if statically_known_true(x.numel() > 10):
    # optimization path
    ...
else:
    # general path (taken when unknown)
    ...

from torch.fx.experimental.symbolic_shapes import guard_or_false

# Instead of: if x == 0:
if guard_or_false(x == 0):
    return 1
else:
    torch._check(x != 0)  # runtime check for the general path
    return compute(x)

# Use ONLY for optimizations, not correctness-critical branches
if optimization_hint(x.numel(), fallback=0) > 1024:
    # optimized path for large tensors
    ...
else:
    # general path
    ...

Шаг 2. Известно ли, что всегда будет выбран один из путей?

Если устранить ветвление невозможно, спросите себя: Известно ли, что для моей конкретной модели всегда будет выбран один и тот же путь?

Если да, используйте torch._check, чтобы сообщить компилятору, какую ветвь выбрать:

i = x.item()
torch._check(i > 4)  # Assert that i > 4 is always true for your use case
if i > 4:
    return x * 2
else:
    return x + 3

Утверждая torch._check(i > 4), вы сообщаете системе символических вычислений, что i > 4 всегда имеет значение True, что позволяет разрешить ветвление без ошибки. С точки зрения компилятора ветвь else становится мёртвым кодом.

torch._check(cond, msg_fn)

torch._check — функция, используемая для проверки условий во время выполнения, в частности при работе с символическими целыми числами (SymInts) в PyTorch.

Пример использования:

torch._check(x.size(0) == y, lambda: f"size mismatch: {x.size(0)} != {y}")

Приведённый выше код выполняет следующие действия:

• Создаёт отложенную проверку во время выполнения вместо проверки во время компиляции
• Сообщает системе символических вычислений факты о ваших неподкреплённых SymInt
• Может устранять неподкреплённые символы, заменяя их эквивалентными выражениями
• Уточняет диапазоны значений символов
• Запоминает логические выражения, которые всегда истинны

Семантически функция работает как условная проверка:

if not cond:
    raise RuntimeError(msg_fn())

Однако есть несколько важных отличий:

• Условие всегда считается истинным во время компиляции, даже если оно включает неподкреплённые SymInts. Фактическая проверка откладывается до выполнения, что позволяет избежать ошибок во время компиляции. Вместо настройки проверки мы реализуем отложенное утверждение, которое проверяет условие во время выполнения. Во время компиляции мы предполагаем, что условие не приведёт к ошибке, поэтому нам не нужно определять, вычисляется ли оно в True или False.
• Если вы выполняете проверку на равенство u0 = RHS, мы пытаемся заменить все вхождения u0 правой частью выражения. Мы ВСЕГДА делаем это, если правая часть не содержит неподкреплённых символов, поскольку удаление неподкреплённых символов полезно: оно предотвращает создание GuardOnDataDependentSymNode. Даже если удалить u0 не удаётся, мы можем уточнить диапазон его значений. Диапазон значений задаёт множество возможных значений переменной. По умолчанию диапазон значений неподкреплённых SymInt, подобных размеру, равен [0, Inf]; если вы утверждаете, что значение равно выражению с уточнённым диапазоном, например [2, 20], диапазон значений u0 обновляется до [2, 20]. Поддерживается также ограниченное распространение диапазонов значений в обратном направлении.
• Если вы выполняете логическую проверку f(u0), мы запоминаем, что это выражение всегда вычисляется в True. Если вы вычислите выражение, содержащее его, мы заменим его на True. Также поддерживаются некоторые ограниченные рассуждения о логически эквивалентных утверждениях. Например, если вы torch._check(u0 < 4), мы также будем знать, что u0 >= 4 вычисляется в False, поэтому такая проверка в обычном условном операторе без check выполнится без ошибок.

Для задания ограничений и уточнения диапазонов значений также можно использовать torch._check. Например, torch._check(u0 >= 0) устанавливает, что u0 неотрицательно, уточняя его диапазон значений до [0, Inf]. Аналогично, torch._check(x > 7) ограничивает x значениями больше 7.

Когда неподкреплённые символы передаются фабричным функциям, например torch.empty, они автоматически распознаются как представляющие размеры.

Шаг 3. Невозможно ли исправить проблему?

Если во время выполнения действительно нужны обе ветви (то есть иногда i > 4, а иногда i <= 4), то никакой torch._check не поможет — выполнить трассировку в текущем виде невозможно. В таких случаях следует рассмотреть альтернативные подходы, например использовать torch.cond или дополнение до нужного размера.

Ещё один распространённый неисправимый шаблон — индексация списка Python с помощью значения, зависящего от данных:

return self.mlps[x.item()]

Здесь self.mlps — это список Python или ModuleList, а ветвление в коде зависит от значения, определяемого данными. Самое простое решение — вызвать разрыв графа перед операцией индексации.

Использование torch._check для проверок корректности в коде модели

Если в коде модели есть проверки корректности условий, можно использовать torch._check вместо инструкций if. torch._check обрабатывает зависимость от данных, откладывая проверки до выполнения, поэтому они не вызывают ошибок во время компиляции.

Примечание: для кода C++ используйте TORCH_SYM_CHECK — аналог torch._check для C++.

При объединении условий используйте sym_or, sym_and и т. д., чтобы выражения не вычислялись заранее (это вызвало бы ошибки, зависящие от данных):

# Instead of:
# if x != y or x > y:
#     raise RuntimeError("...")

# Use:
from torch.fx.experimental.symbolic_shapes import sym_or
torch._check(sym_or(x != y, x > y), lambda: "Validation failed: expected x != y or x > y")

u0 на самом деле равно u1, но мы этого не знаем

Несколько неподкреплённых SymInts могут быть заведомо равны во время компиляции:

i0 = x.sum().item()
i1 = x.sum().item()
return torch.randn(i0) + torch.randn(i1)

Если где-либо есть torch._check(i0 == i1) (в приведённом выше примере эта проверка выполняется внутри правила проверки формы для сложения), мы автоматически объединим два неподкреплённых SymInts и распознаем их равенство. Однако, если такое утверждение отсутствует, для этого объединения может потребоваться явно добавить проверку. Пример см. в https://github.com/pytorch/pytorch/issues/111950).

u0 — это тензор

Ещё одна причина избыточного выделения неподкреплённых SymInts — передача Tensor с расчётом на его неявное преобразование в целое число. Многие функции, принимающие целое число, также принимают Tensor и автоматически вызывают item() для целочисленного аргумента. Рекомендуется изучить TORCH_LOGS=dynamic, чтобы определить, ожидаемо ли количество неподкреплённых SymInts или оно чрезмерно. В этом случае новый SymInt будет выделен в строке вызова функции PyTorch.

Теперь эта проблема реже вызывает трудности, поскольку возвращаемое значение t.item() кэшируется: при повторных вызовах вы неизменно получаете один и тот же неподкреплённый SymInt.

Проблема чрезмерной специализации

В режиме нестрогого экспорта рассмотрим следующий код:

u0 = x.sum().item()
return y[:u0]

Этот код завершится ошибкой при попытке вычислить u0, поскольку при непосредственном использовании SymInt внутри среза Python (без Dynamo) Python принудительно специализирует целое число и выдаёт ошибку, если оно не подкреплено.

Чтобы решить эту проблему, можно переписать программу, избегая специализации. В приведённом выше примере проблему можно исправить, отказавшись от срезов:

u0 = x.sum().item()
return y.narrow(0, 0, u0)

Дополнительные сведения см. в связанной задаче: https://github.com/pytorch/pytorch/issues/111950.

Используйте длины вместо смещений

При работе с переменной длиной последовательностей часто используются тензоры, представляющие длины или смещения последовательностей. Например, для values = [[1, 2, 3], [4, 5], [6, 7, 8, 9]] у вас могут быть lengths = [3, 2, 4] и offsets = [0, 3, 5, 9]. Хотя эти представления взаимно преобразуемы, при работе с ними как с целыми числами (вызвав lengths.tolist()) лучше использовать длины, а не смещения.

Причина в том, что при выполнении torch.split() над тензором values нужно создавать тензоры для каждой подпоследовательности, например тензоры размеров 3, 2 и 4. Если размеры представлены неподкреплёнными SymInts, получаются u0, u1 и u2. Достаточно указать, что они подобны размерам. Если же смещения представлены неподкреплёнными SymInts, получаются u1 - u0, u2 - u1, u3 - u2, что усложняет задачу. Эти величины нельзя удобно пометить как подобные размерам, из-за чего могут возникнуть проблемы. Поскольку код можно сравнительно легко написать и с длинами, и со смещениями, предпочтительнее использовать длины.