torch.compile

torch.compile(model=None, *, fullgraph=False, dynamic=None, backend='inductor', mode=None, options=None, disable=False) [source]

Оптимизирует заданную модель/функцию с использованием TorchDynamo и указанного бэкенда.

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

Параметры

• model (Callable) – Модуль/функция для оптимизации
• fullgraph (bool) – Разрешить разделение модели на несколько подграфов
• dynamic (bool или None) – Использовать динамическое отслеживание формы. Если True, мы предварительно попробуем сгенерировать ядро, которое будет максимально динамичным, чтобы избежать перекомпиляций при изменении размеров. Это может не всегда работать, так как некоторые операции/оптимизации принудят специализацию; используйте TORCH_LOGS=dynamic для отладки чрезмерной специализации. Если False, мы НИКОГДА не будем генерировать динамические ядра, мы всегда будем специализироваться. По умолчанию (None) мы автоматически обнаруживаем, произошла ли динамизация, и при перекомпиляции компилируем более динамичное ядро.

• backend (str или Callable) –

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

  • ”inductor” является бэкендом по умолчанию, который представляет собой хороший баланс между производительностью и накладными расходами
  • Неэкспериментальные бэкенды в составе библиотеки можно посмотреть с torch._dynamo.list_backends()
  • Экспериментальные или отладочные бэкенды в составе библиотеки можно посмотреть с torch._dynamo.list_backends(None)
  • Для регистрации пользовательского бэкенда вне библиотеки: https://pytorch.org/docs/main/compile/custom-backends.html

• mode (str) –

  Может быть «default», «reduce-overhead», «max-autotune» или «max-autotune-no-cudagraphs»

  • «default» — режим по умолчанию, который представляет собой хороший баланс между производительностью и накладными расходами
  • «reduce-overhead» — режим, который уменьшает накладные расходы Python с CUDA-графиками, полезен для небольших батчей. Уменьшение накладных расходов может произойти за счёт большего использования памяти, так как мы будем кешировать рабочую память, необходимую для вызова, чтобы не перераспределять её при последующих запусках. Уменьшение накладных расходов не гарантировано; сегодня мы уменьшаем накладные расходы только для CUDA-графиков, которые не изменяют входные данные. Есть и другие случаи, когда CUDA-графы неприменимы; используйте TORCH_LOG=perf_hints для отладки.
  • «max-autotune» — режим, использующий Triton для матричных умножений и свёртки. Он по умолчанию включает CUDA-графы.
  • «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 которая покажет вам график после слияния
  • Для inductor вы можете посмотреть полный список поддерживаемых конфигураций, вызвав torch._inductor.list_options()
• disable (bool) – Сделать torch.compile() бесполезной операцией для тестирования

Тип возвращаемого значения

Callable

Пример:

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