torch.compile

torch.compile(model: Callable[[_InputT], _RetT], *, fullgraph: bool = False, dynamic: bool | None = None, backend: str | Callable[[...], Any] = 'inductor', mode: str | None = None, options: dict[str, str | int | bool | Callable[[...], Any]] | None = None, name: str | None = None, disable: bool = False, dynamic_shapes: Any = None) → Callable[[_InputT], _RetT] [исходный код]

torch.compile(model:None=None, *, fullgraph:bool=False, dynamic:bool|None=None, backend:str|Callable[[...],Any]='inductor', mode:str|None=None, options:dict[str,str|int|bool|Callable[[...],Any]]|None=None, name:str|None=None, disable:bool=False, dynamic_shapes:Any=None) → Callable[[Callable[[_InputT],_RetT]],Callable[[_InputT],_RetT]]

Оптимизирует заданную модель/функцию с помощью TorchDynamo и указанного бэкенда. Если вы компилируете torch.nn.Module, можно также использовать torch.nn.Module.compile(), чтобы скомпилировать модуль на месте, не изменяя его структуру.

В частности, для каждого кадра, выполняемого внутри скомпилированной области, мы попытаемся скомпилировать его и кэшировать результат компиляции в объекте кода для дальнейшего использования. Один кадр может быть скомпилирован несколько раз, если предыдущие результаты компиляции неприменимы к последующим вызовам (это называется «сбоем проверки»); для отладки таких ситуаций можно использовать TORCH_LOGS=guards. С кадром может быть связано несколько результатов компиляции — до torch._dynamo.config.recompile_limit, значение по умолчанию — 8; после этого выполнение будет переключено на eager-режим. Обратите внимание, что кэши компиляции относятся к объекту кода, а не к кадру; если вы динамически создаёте несколько копий функции, все они будут использовать один и тот же кэш кода.

Параметры:

• model (Callable or None) – Модуль/функция для оптимизации
• fullgraph (bool) – Если False (значение по умолчанию), torch.compile пытается найти в функции компилируемые области, которые можно оптимизировать. Если True, требуется, чтобы всю функцию можно было захватить в один граф. Если это невозможно (то есть если возникают разрывы графа), будет вызвана ошибка. Этот параметр также включает семантику unbacked; в частности, по умолчанию включаются capture_scalar_outputs и capture_dynamic_output_shape_ops.
• dynamic (bool or None) – Использовать трассировку динамических форм. Если задано True, мы заранее попытаемся создать максимально динамичное ядро, чтобы избежать повторной компиляции при изменении размеров. Это не всегда возможно, поскольку некоторые операции/оптимизации принудительно требуют специализации; для отладки чрезмерной специализации используйте TORCH_LOGS=dynamic. Если задано False, динамические ядра НЕ создаются — всегда выполняется специализация. По умолчанию (None) мы автоматически определяем появление динамичности и при повторной компиляции создаём более динамичное ядро.

• backend (str or Callable) –

  Используемый бэкенд

  • «inductor» — бэкенд по умолчанию, обеспечивающий хороший баланс между производительностью и накладными расходами
  • Нестандартные бэкенды, включённые в основную поставку и не являющиеся экспериментальными, можно просмотреть с помощью torch._dynamo.list_backends()
  • Экспериментальные бэкенды или бэкенды для отладки, включённые в основную поставку, можно просмотреть с помощью torch._dynamo.list_backends(None)
  • Чтобы зарегистрировать пользовательский бэкенд вне основной поставки: https://docs.pytorch.org/docs/main/user_guide/torch_compiler/torch.compiler_custom_backends.html#registering-custom-backends

• mode (str) –

  Может принимать значения «default», «reduce-overhead», «max-autotune» или «max-autotune-no-cudagraphs»

  • «default» — режим по умолчанию, обеспечивающий хороший баланс между производительностью и накладными расходами
  • «reduce-overhead» — режим, уменьшающий накладные расходы Python с помощью графов CUDA; полезен для небольших пакетов. Снижение накладных расходов может привести к увеличению использования памяти, поскольку рабочая память, необходимая для вызова, будет кэшироваться, чтобы не выделять её повторно при последующих запусках. Снижение накладных расходов не гарантируется; на данный момент оно применяется только к графам, использующим исключительно CUDA и не изменяющим входные данные. Существуют и другие случаи, когда графы CUDA неприменимы; для отладки используйте TORCH_LOGS=perf_hints.
  • «max-autotune» — режим, использующий умножение матриц на основе Triton или шаблонов на поддерживаемых устройствах, а также свёртки на основе Triton на GPU. По умолчанию он включает графы CUDA на GPU.
  • «max-autotune-no-cudagraphs» — режим, аналогичный «max-autotune», но без графов CUDA
  • Чтобы увидеть точные конфигурации, задаваемые каждым режимом, можно вызвать torch._inductor.list_mode_options()

• options (dict) –

  Словарь параметров, передаваемых бэкенду. Среди параметров, которые стоит попробовать:

  • epilogue_fusion, который объединяет поэлементные операции в шаблоны. Для этого также требуется задать max_autotune
  • max_autotune, который выполняет профилирование для выбора наилучшей конфигурации matmul
  • fallback_random, полезный при отладке проблем с точностью
  • shape_padding, который дополняет формы матриц, чтобы лучше выравнивать загрузку данных на GPU, особенно для тензорных ядер
  • triton.cudagraphs, который уменьшает накладные расходы Python с помощью графов CUDA
  • trace.enabled — наиболее полезный флаг для включения при отладке
  • trace.graph_diagram, который покажет изображение графа после слияния
  • guard_filter_fn, управляющий тем, какие guards dynamo сохраняются вместе с результатами компиляции. Это небезопасная функция, и обратная совместимость для guards dynamo в качестве типов данных не гарантируется. Описание стабильных вспомогательных функций см. в документации по адресу torch.compiler, например: - torch.compiler.skip_guard_on_inbuilt_nn_modules_unsafe - torch.compiler.skip_guard_on_all_nn_modules_unsafe - torch.compiler.keep_tensor_guards_unsafe
  • Полный список поддерживаемых конфигураций inductor можно получить, вызвав torch._inductor.list_options()
• name (str or None) – Необязательный идентификатор скомпилированной области. Если последующие инструменты это поддерживают, идентификатор отображается в обёрнутых операторах высшего порядка скомпилированной области и других метаданных отладки.
• disable (bool) – Превращает torch.compile() в пустую операцию для тестирования
• recompile_limit (int or None) – Максимальное число повторных компиляций для этого вызова torch.compile() до переключения на eager-режим. Если указано None (значение по умолчанию), используется глобальный параметр torch._dynamo.config.recompile_limit (по умолчанию 8). При использовании fullgraph=True превышение лимита вызывает FailOnRecompileLimitHit.
• isolate_recompiles (bool) – Если задано True, этот вызов torch.compile() отслеживает повторные компиляции независимо. По умолчанию все вызовы torch.compile() одной и той же функции используют общий набор скомпилированных записей, поэтому повторные компиляции одного вызова учитываются в лимите всех остальных вызовов. При использовании isolate_recompiles=True каждый вызов получает отдельный набор записей. Поиск для изолированного вызова компиляции по-прежнему обращается к записям неизолированных вызовов компиляции (повторное использование, совместимое с обратной совместимостью), однако новые результаты компиляции сохраняются отдельно. recompile_limit проверяется для каждой области; accumulated_recompile_limit — это глобальное ограничение для всех областей. Поведение стратегии по умолчанию асимметрично: решения SKIP (из skip_code, @torch._dynamo.skip, кода, сгенерированного FX, и т. д.) наследуются изолированными областями — выполнение в них по-прежнему пропускается. RUN_ONLY, сохранённый неизолированной областью при достижении лимита повторных компиляций, НЕ распространяется на изолированные области — каждая область управляет собственным состоянием RUN_ONLY. Значение по умолчанию — False.

Пример:

@torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
def foo(x):
    return torch.sin(x) + torch.cos(x)