Работа с повторными компиляциями

Включены ли динамические формы?

В приведённом ниже примере компиляция выполняется повторно из-за несовпадения форм:

@torch.compile
def fn(x):
    return x + 1
fn(torch.ones(3))
fn(torch.ones(4))

Убедитесь, что параметр dynamic в torch.compile не установлен в False. Параметр по умолчанию, dynamic=None, будет пробовать использовать динамические формы только после первой компиляции. Можно установить dynamic=True, чтобы сразу компилировать с максимально возможным использованием динамических форм:

@torch.compile(dynamic=True)
def gn(x):
    return x + 1
gn(torch.ones(3))
gn(torch.ones(4))

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

Обёртывание констант в тензоры

По умолчанию переменные int / float считаются константами, и для них устанавливаются проверки на точное значение. В приведённом ниже примере повторная компиляция выполняется при каждом вызове функции.

@torch.compile
def fn(x, c):
    return x + c
for i in range(5):
    fn(torch.ones(i), 0.5 + i)

В частности, инициализация планировщиков LR константой может привести к повторным компиляциям:

mod = torch.nn.Linear(3, 3)
opt = torch.optim.Adam(mod.parameters(), lr=0.01)
sched = torch.optim.lr_scheduler.ExponentialLR(opt, 0.9)
@torch.compile
def gn(inp):
    opt.zero_grad(True)
    out = mod(inp).sum()
    out.backward()
    opt.step()
    sched.step()
for i in range(5):
    gn(torch.ones(3, 3))

В обоих примерах можно обернуть переменные float в тензоры, чтобы избежать повторных компиляций.

# first example
for i in range(5):
    fn(torch.ones(i), torch.tensor(0.5 + i))
# second example
opt = torch.optim.Adam(mod.parameters(), lr=torch.tensor(0.01))
sched = torch.optim.lr_scheduler.ExponentialLR(opt, torch.tensor(0.9))
for i in range(5):
    gn(torch.ones(3, 3))

Снятие специализации целочисленных атрибутов у nn.Modules

По умолчанию Dynamo специализируется на целочисленных атрибутах nn.Module (проверяет их точные значения). Поэтому целое число, которое меняется при каждом вызове, например счётчик шагов, увеличиваемый внутри forward, вызывает повторную компиляцию при каждом вызове. Вскоре достигается лимит повторных компиляций, и выполнение переключается на eager-режим.

class Mod(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.c = 0

    def forward(self, x):
        self.c += 1  # specialized on -> recompiles every call
        return x * self.c

mod = Mod()
opt_mod = torch.compile(mod, backend="eager")
for _ in range(5):
    opt_mod(torch.randn(4))

Если вам не нужна специализация по этому целому числу, установите флаг ниже, чтобы Dynamo считал его динамическим и не выполнял повторную компиляцию. Связанные флаги управления статическим и динамическим поведением см. в разделе Дополнительные параметры управления динамическим поведением.

torch._dynamo.config.allow_unspec_int_on_nn_module = True
mod = Mod()
opt_mod = torch.compile(mod, backend="eager")
for _ in range(5):
    opt_mod(torch.randn(4))
torch._dynamo.config.allow_unspec_int_on_nn_module = False

Изменение ограничения размера кэша

Количество повторных компиляций функции ограничено. Это ограничение определяется параметрами torch._dynamo.config.cache_size_limit и torch._dynamo.config.accumulated_cache_size_limit (точное различие между этими двумя значениями описано в torch/_dynamo/cache_size.py). При достижении ограничения кэша Dynamo все последующие попытки компиляции приведут к пропуску функции (она будет выполнена в eager-режиме). При этом Dynamo будет пытаться использовать ранее скомпилированный байт-код для будущих вызовов функции, если проверки пройдут. Обратите внимание: при достижении ограничения на повторные компиляции все вложенные вызовы функций БУДУТ пропущены (Dynamo попытается использовать ранее скомпилированный байт-код для вложенных функций). Dynamo также выдаст предупреждение с указанием затронутой функции и достигнутого ограничения. В приведённом ниже примере при каждом вызове функции предпринимается попытка повторной компиляции. При достижении ограничения размера кэша (по умолчанию — 8) попытки повторной компиляции прекращаются. (Обратите внимание: для демонстрации мы задаём dynamic=False, чтобы принудительно запускать повторную компиляцию при каждом вызове.)

@torch.compile(dynamic=False)
def fn(x):
    return x + 1
for i in range(1, 10):
    # recompile every time due to dynamic=False
    fn(torch.ones(i))

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

torch._dynamo.config.cache_size_limit = 16
@torch.compile(dynamic=False)
def gn(x):
    return x + 1
for i in range(1, 10):
    gn(torch.ones(i))

Сокращение числа повторных компиляций — один из способов уменьшить время компиляции. О других возможностях (региональная и иерархическая компиляция, измерение времени компиляции) см. в разделе Сокращение времени компиляции.

Разрыв графа для снижения затрат на повторную компиляцию

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

def very_large_function(x):
    return x + 1

@torch.compile(dynamic=False)
def fn(x, c):
    y = very_large_function(x)  # recompiled every time
    return y + c

for i in range(1, 5):
    fn(torch.ones(3), i)

@torch.compile(dynamic=False)
def gn(x, c):
    y = very_large_function(x)  # compiled only once
    torch._dynamo.graph_break()
    return y + c  # recompiled every time

for i in range(1, 5):
    gn(torch.ones(3), i)