Сокращение времени компиляции

Избегайте ненужных перекомпиляций

Каждая перекомпиляция требует повторной оплаты затрат на компиляцию, поэтому самая распространённая причина чрезмерного времени компиляции — лишние перекомпиляции. Выявите перекомпиляции с помощью tlparse или TORCH_LOGS=recompiles, а затем устраните их причины. Полный порядок действий см. в разделе Работа с перекомпиляциями.

Используйте динамические формы, чтобы избежать перекомпиляций из-за изменения форм

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

Попытки использовать динамические формы управляются аргументом dynamic функции torch.compile; описание поведения настроек dynamic=None/True/False см. в разделе Включены ли динамические формы?. Чтобы определённое измерение было динамическим уже при первой компиляции, явно аннотируйте его с помощью mark_dynamic:

import torch

# Mark dim 0 of x as dynamic, optionally with a min/max range, BEFORE compiling.
torch._dynamo.mark_dynamic(x, 0, min=1, max=1024)

opt_fn = torch.compile(fn)
opt_fn(x)

mark_dynamic необходимо вызывать для входных тензоров до их передачи в скомпилированную функцию (не внутри скомпилированной области). Полное описание понятия «динамический», автоматической динамики и других аннотаций см. в разделе Динамические формы.

Региональная компиляция

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

Пример см. в руководстве по региональной компиляции.

Иерархическая компиляция с помощью nested_compile_region

torch.compiler.nested_compile_region помечает набор операций — обычно повторяющуюся структурную единицу, например слой трансформера LLM, — как вложенный регион компиляции. При полной компиляции модели torch.compile компилятор при первом появлении региона генерирует для него оптимизированный код, а при каждом последующем появлении повторно использует («тиражирует») этот скомпилированный код, вместо того чтобы компилировать регион заново. Это позволяет существенно сократить общее время компиляции компонентов с глубокой вложенностью и идентичной структурой.

import torch

@torch.compiler.nested_compile_region
def transformer_layer(x):
    ...

@torch.compile
def model(x):
    for _ in range(num_layers):
        x = transformer_layer(x)
    return x

В отличие от региональной компиляции, этот подход работает внутри одного вызова torch.compile и не требует менять способ применения torch.compile. Он также безопасен: нет гарантии, что регион будет скомпилирован ровно один раз. Если новые входные условия (форма, dtype, устройство, шаг, глобальные переменные и т. д.) сделают кэшированный регион недействительным, компилятор прозрачно скомпилирует его повторно, поэтому корректность всегда сохраняется, а дополнительные затраты на компиляцию возникают только при необходимости. Вне контекста torch.compile вызов ничего не делает.

Измерение времени компиляции

Перед оптимизацией измерьте, на что уходит время компиляции:

• torch._dynamo.utils.compile_times() выводит время, затраченное на каждый этап компиляции Dynamo.
• TORCH_COMPILE_DEBUG=1 показывает длительность этапов Inductor и артефакты.
• tlparse / TORCH_TRACE формирует отчёт по каждой компиляции; см. раздел tlparse / TORCH_TRACE.

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