Динамические формы

mark_dynamic(tensor, dim, min=min, max=max)

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

torch._dynamo.mark_dynamic() нельзя вызывать внутри функции, компилируемой с помощью torch.compile() (например, метода forward() модели или любой вызываемой им функции).

Эта функция является API времени трассировки. Если вызвать её из скомпилированного кода, Dynamo выдаст ошибку, например:

AssertionError: Attempt to trace forbidden callable

Правильный способ использования — вызвать mark_dynamic для входных тензоров до вызова torch.compile, например:

torch._dynamo.mark_dynamic(x, 0)
compiled_model = torch.compile(model)

Функция torch._dynamo.mark_dynamic() помечает размерность тензора как динамическую и завершится с ошибкой, если размерность будет специализирована. Она не работает с целыми числами. Используйте эту функцию, только если уверены, что все графы в кадре, использующие этот вход, сойдутся к одному динамическому графу. В противном случае может появиться вводящее в заблуждение сообщение об ошибке нарушения ограничений. В таких случаях рассмотрите возможность использования torch._dynamo.maybe_mark_dynamic(). В настоящее время torch._dynamo.mark_dynamic() не имеет приоритета над force_parameter_static_shapes = True или force_nn_module_property_static_shapes = True.

Если заранее известно, что определённая размерность будет динамической, можно избежать первоначальной повторной компиляции, используя torch._dynamo.mark_dynamic(tensor, dim)(). Кроме того, если уже известны минимальное и максимальное возможные значения этой размерности, их можно указать с помощью torch._dynamo.mark_dynamic(tensor, dim, min=min, max=max)().

Вот краткий пример:

import torch

@torch.compile
def f(x):
    return x * x.size()[0]

x = torch.randn(10)
torch._dynamo.mark_dynamic(x, 0)

# first invocation we give it is a tensor marked as dynamic
f(x)
# rest of these invocations will use dynamically compiled code
f(torch.randn(20))
f(torch.randn(30))
f(torch.randn(40))

maybe_mark_dynamic(tensor, dim)

Функция torch._dynamo.maybe_mark_dynamic() обладает теми же свойствами, что и torch._dynamo.mark_dynamic(), но не завершится с ошибкой, если размер будет специализирован. Используйте её для входных данных, общих для нескольких графов, или если число графов для определённого кадра не сходится к одному. Например, в приведённом выше примере используйте torch._dynamo.maybe_mark_dynamic(), поскольку графы для размеров 0 и 1 будут специализированы. Однако с помощью torch._dynamo.mark_dynamic() можно гарантировать, что специализация никогда не произойдёт.

mark_unbacked(tensor, dim)

Функция torch._dynamo.decorators.mark_unbacked() помечает размерность тензора как неподкреплённую. Вероятно, это не тот инструмент, который вам нужен, однако он может оказаться полезным, если специализация происходит внутри условия guard_size_oblivious(x) и использование этой функции позволяет её устранить. Убедитесь, что она устраняет специализацию и не приводит к ошибке зависимости от данных, из-за которой в месте специализации, которое вы пытаетесь обойти, или перед ним происходит разрыв графа. Возможно, лучше использовать следующий вариант.

Список разрешённых динамических источников (DYNAMIC_SOURCES)

Используйте переменную окружения TORCH_COMPILE_DYNAMIC_SOURCES, чтобы передать список имён источников, которые следует пометить как динамические. Например: TORCH_COMPILE_DYNAMIC_SOURCES=L[‘x’],L[‘y’] Проще всего найти имена динамических источников в артефакте PGO, расположенном в tlparse. Имена динамических источников можно скопировать из артефакта PGO. Этот способ работает с целыми числами и размерами тензоров и имеет наивысший приоритет среди всех остальных флагов, принудительно включающих статические формы. Ошибка не возникнет, если помеченное как динамическое значение будет специализировано или если указанный вход не существует.

Вот пример:

import torch

@torch.compile()
def f(x):
     return x * x.size()[0]

with torch.compiler.config.patch(dynamic_sources="L['x']"):
    f(torch.rand(10))
f(torch.rand(20))
f(torch.rand(30))
f(torch.rand(40))

torch.compiler.set_stance ("eager_then_compile")

Иногда бывает сложно определить, какие именно входные данные нужно пометить как динамические. Если вы готовы пожертвовать производительностью первого пакета, можно воспользоваться режимами eager_then_compile, которые автоматически определяют динамические входные данные. Подробнее см. в описании torch.compiler.set_stance() и в руководстве Управление динамической компиляцией с помощью torch.compiler.set_stance.