Часто задаваемые вопросы

Автор: Марк Сарофим

Поддерживает ли torch.compile обучение?

torch.compile поддерживает обучение, используя AOTAutograd для захвата обратного прохода:

1. Граф .forward() и optimizer.step() захватываются с помощью Python evalframe фронтенда TorchDynamo.
2. Для каждого сегмента .forward(), который захватывает torchdynamo, используется AOTAutograd для генерации сегмента графа обратного прохода.
3. Каждая пара графов прямого и обратного прохода (необязательно) разделяются с помощью минимального разреза, чтобы сохранить минимальное состояние между прямым и обратным проходами.
4. Пары прямого и обратного прохода оборачиваются в модули autograd.function.
5. Вызов кода пользователя .backward() все равно запускает движок автоградирования eager, который выполняет каждый скомпилированный граф обратного прохода как одну операцию, а также выполняет все функции .backward() нескомпилированных операций eager.

Поддерживаете ли вы распределённый код?

torch.compile поддерживает DistributedDataParallel (DDP). Поддержка других библиотек распределённого обучения рассматривается.

Основная причина, по которой распределённый код сложен с dynamo, заключается в том, что AOTAutograd развертывает как прямой, так и обратный проход и предоставляет 2 графа для оптимизации бэкендами. Это проблема для распределённого кода, потому что нам хотелось бы перекрывать операции связи с вычислениями. Eager pytorch достигает этого различными способами для DDP/FSDP – с помощью хуков автоградирования, хуков модулей и модификаций/мутаций состояний модулей. В простейшем применении dynamo хуки, которые должны выполняться непосредственно после операции во время обратного прохода, могут быть отложены до окончания всей скомпилированной области операций обратного прохода из-за того, как скомпилированные функции AOTAutograd взаимодействуют с хуками диспетчера.

Основная стратегия оптимизации DDP с Dynamo изложена в distributed.py, где основная идея будет заключаться в разрыве графа на границах ведер DDP.

Когда каждый узел в DDP должен синхронизировать свои веса с другими узлами, он организует свои градиенты и параметры в ведра, что уменьшает время связи и позволяет узлу передать часть своих градиентов другим ожидающим узлам.

Разрывы графа в распределённом коде означают, что вы можете ожидать, что dynamo и его бэкэнды оптимизируют вычислительную нагрузку распределённой программы, но не её нагрузку связи. Разрывы графа могут мешать ускорениям компиляции, если уменьшенный размер графа лишает компилятор возможностей слияния. Однако с увеличением размера графа существует убывающая отдача, поскольку большинство текущих оптимизаций вычислений представляют собой локальные слияния. Таким образом, на практике этот подход может быть достаточным.

Нужно ли мне по-прежнему экспортировать целые графы?

Для подавляющего большинства моделей, вероятно, нет, и вы можете использовать torch.compile() как есть, но есть несколько ситуаций, в которых нужны полные графы, и вы можете гарантировать полный граф, просто выполнив torch.compile(..., nopython=True). Эти ситуации включают:

• Крупномасштабные учебные забеги, такие как $250K+, которые требуют параллелизма обработки данных и других расширенных стратегий фрагментации.
• Оптимизаторы вывода, такие как TensorRT или AITemplate, которые полагаются на более агрессивное слияние, чем оптимизаторы обучения.
• Обучение или вывод на мобильных устройствах.

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

Почему мой код вылетает?

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

1. torch.compile(..., backend="eager") который выполняет только захват графа прямого прохода TorchDynamo, а затем выполняет захваченный граф с PyTorch. Если это не удается, то проблема в TorchDynamo.
2. torch.compile(..., backend="aot_eager") который выполняет TorchDynamo для захвата графа прямого прохода, а затем AOTAutograd для прослеживания графа обратного прохода без дополнительных шагов компилятора бэкенда. PyTorch eager будет использоваться для выполнения графов прямого и обратного прохода. Если это не удается, то проблема в AOTAutograd.
3. torch.compile(..., backend="inductor") который выполняет TorchDynamo для захвата графа прямого прохода, а затем AOTAutograd для прослеживания графа обратного прохода с компилятором TorchInductor. Если это не удается, то проблема в TorchInductor

Почему компиляция медленная?

• Компиляция Dynamo – TorchDynamo имеет встроенную функцию статистики для сбора и отображения времени, затраченного на каждом этапе компиляции. К этим статистическим данным можно получить доступ, вызвав torch._dynamo.utils.compile_times() после выполнения torch._dynamo. По умолчанию это возвращает строковое представление времени компиляции, затраченного на каждую функцию TorchDynamo по имени.
• Компиляция Inductor – TorchInductor имеет встроенные функции статистики и отслеживания для отображения времени, затраченного на каждом этапе компиляции, выходного кода, визуализации выходного графа и вывода IR. env TORCH_COMPILE_DEBUG=1 python repro.py. Это инструмент отладки, разработанный для облегчения отладки/понимания внутренних механизмов TorchInductor с выводом, который будет выглядеть примерно так, как этот. Каждый файл в этом отладочном трассировке можно включить/выключить с помощью torch._inductor.config.trace.*. Профиль и диаграмма по умолчанию отключены, так как их генерация дорогостоящая. Обратитесь к примерному выводу каталога отладки для получения дополнительных примеров.
• Избыточная перекомпиляция. Когда TorchDynamo компилирует функцию (или часть функции), он делает определённые предположения о локальных и глобальных переменных, чтобы позволить оптимизации компилятора, и выражает эти предположения в виде проверок, которые проверяют определённые значения во время выполнения. Если какая-либо из этих проверок терпит неудачу, Dynamo перекомпилирует эту функцию (или её часть) до torch._dynamo.config.cache_size_limit раз. Если ваша программа достигает предела кэша, вам сначала необходимо определить, какая проверка терпит неудачу и какая часть вашей программы её вызывает. Профилировщик перекомпиляций автоматизирует процесс установки предела кэша TorchDynamo в 1 и запуска вашей программы под наблюдением только компилятора, который записывает причины любых ошибок проверки. Вы должны быть уверены, что запускаете свою программу в течение как минимум такого же времени (такое же количество итераций), как и когда у вас возникли проблемы, и профилировщик будет накапливать статистические данные в течение этого времени.

from torch._dynamo.utils import CompileProfiler

def my_model():
    ...

with CompileProfiler() as prof:
    profiler_model = torch.compile(my_model, backend=prof)
    profiler_model()
    print(prof.report())

Почему вы перекомпилируете в производстве?

В некоторых случаях вы можете не захотеть неожиданных компиляций после прогрева программы. Например, если вы обслуживаете трафик производства в приложении, критичном для задержек. Для этого TorchDynamo предоставляет альтернативный режим, в котором используются ранее скомпилированные графы, но новые не генерируются:

frozen_toy_example = dynamo.run(toy_example)
frozen_toy_example(torch.randn(10), torch.randn(10))

Как вы ускоряете мой код?

1. Слияние ядер с помощью вертикальных слияний, которые сливают последовательные операции, чтобы избежать чрезмерных чтений/записей. Например, слияние 2 последовательных косинусов означает, что вы можете выполнить 1 чтение 1 запись вместо 2 чтений 2 записей 2. Горизонтальное слияние: простейшим примером является пакетное умножение, где одна матрица умножается на пакет примеров, но более общий сценарий – это групповое GEMM, где группа матричных умножений планируется вместе
2. Выполнение в произвольном порядке: общая оптимизация для компиляторов, анализируя зависимость данных в графе, мы можем определить наиболее подходящее время для выполнения узла и повторного использования буферов
3. Автоматическое распределение работ: аналогично пункту о выполнении в произвольном порядке, но подгон узлов графа к ресурсам, таким как физическое оборудование или память, позволяет разработать соответствующее расписание

Вышеперечисленные принципы являются общими принципами ускорения кода PyTorch, но различные бэкэнды будут делать различные компромиссы по оптимизации. Например, Inductor сначала позаботится о слиянии всего, что возможно, а только потом сгенерирует ядра Triton. Он также может

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

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

Почему я не вижу ускорения?

Разрывы графа

Основная причина, по которой вы не увидите желаемого ускорения с помощью dynamo, – это чрезмерные разрывы графа. Что такое разрыв графа?

Рассмотрим программу, подобную:

def some_fun(x):
    ...

torch.compile(some_fun)(x)
...

Torchdynamo попытается скомпилировать все операции torch/tensor внутри some_fun() в единый граф FX, но может не захватить всё в один граф.

Некоторые причины разрыва графа непреодолимы для TorchDynamo, например, вызов C-расширения, отличного от PyTorch, невидим для TorchDynamo и может делать произвольные вещи без возможности TorchDynamo внести необходимые проверки для обеспечения безопасности повторного использования скомпилированной программы.

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

Определение причины разрыва графа

Для определения всех разрывов графа в программе и связанных причин разрывов можно использовать torch._dynamo.explain. Этот инструмент выполняет TorchDynamo для предоставленной функции и агрегирует обнаруженные разрывы графа. Вот пример использования:

import torch
import torch._dynamo as dynamo
def toy_example(a, b):
    x = a / (torch.abs(a) + 1)
    print("woo")
    if b.sum() < 0:
        b = b * -1
    return x * b
explanation, out_guards, graphs, ops_per_graph = dynamo.explain(toy_example, torch.randn(10), torch.randn(10))
print(explanation)
"""
Dynamo produced 3 graphs, with 2 graph break and 6 ops.
 Break reasons:
1. call_function BuiltinVariable(print) [ConstantVariable(str)] {}
   File "t2.py", line 16, in toy_example
    print("woo")

2. generic_jump
   File "t2.py", line 17, in toy_example
    if b.sum() < 0:
 """

Для вывода ошибки при первом обнаруженном разрыве графа вы можете отключить обратную совместимость с Python, используя nopython=True, это должно быть знакомо, если вы работали с компиляторами на основе экспорта.

def toy_example(a, b):
   ...

torch.compile(toy_example, fullgraph=True, backend=<compiler>)

Почему мой код не перекомпилировался, когда я его изменил?

Если вы включили динамические формы, установив env TORCHDYNAMO_DYNAMIC_SHAPES=1 python model.py, то ваш код не будет перекомпилироваться при изменении форм. Мы добавили поддержку динамических форм, что позволяет избежать перекомпиляций в случае, если формы меняются менее чем вдвое. Это особенно полезно в таких сценариях, как изменение размеров изображений в компьютерном зрении или переменная длина последовательности в NLP. В сценариях вывода часто невозможно заранее узнать, каким будет размер пакета, потому что вы берете то, что можете получить от различных клиентских приложений.

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

Почему я получаю некорректные результаты?

Проблемы с точностью также можно минимизировать, если вы установите переменную окружения TORCHDYNAMO_REPRO_LEVEL=4, она работает по принципу модели git bisect, а полное воспроизведение может быть примерно таким: TORCHDYNAMO_REPRO_AFTER="aot" TORCHDYNAMO_REPRO_LEVEL=4 Причина, по которой нам это нужно, заключается в том, что компиляторы нижнего уровня генерируют код, будь то код Triton или C++ бэкэнд, арифметические операции этих компиляторов нижнего уровня могут быть немного разными, но при этом оказывать драматическое влияние на стабильность обучения. Таким образом, отладчик точности очень полезен для обнаружения ошибок в нашем кодогенераторе или компиляторе бэкенда.

Если вы хотите гарантировать, что генерация случайных чисел совпадает как в torch, так и в triton, вы можете включить torch._inductor.config.fallback_random = True

Почему у меня возникают ошибки OOM?

Dynamo — это всё ещё продукт альфа-версии, поэтому есть несколько источников ошибок OOM. Если вы видите ошибку OOM, попробуйте отключить следующие настройки в указанном порядке, а затем откройте вопрос на GitHub, чтобы мы могли решить основную проблему 1. Если вы используете динамические формы, попробуйте отключить их, мы отключили их по умолчанию: env TORCHDYNAMO_DYNAMIC_SHAPES=0 python model.py 2. Графики CUDA с Triton включены по умолчанию в inductor, но их отключение может устранить некоторые проблемы с OOM: torch._inductor.config.triton.cudagraphs = False

Работает ли torch.func с torch.compile (для grad и vmap преобразований)?

Применение преобразования torch.func к функции, использующей torch.compile, не работает:

import torch

@torch.compile
def f(x):
    return torch.sin(x)

def g(x):
    return torch.grad(f)(x)

x = torch.randn(2, 3)
g(x)

Этот код не будет работать. Есть проблема, которую вы можете отслеживать.

В качестве обходного решения используйте torch.compile вне функции torch.func:

import torch

torch._dynamo.config.capture_func_transforms=True

def f(x):
    return torch.sin(x)

@torch.compile
def g(x):
    return torch.vmap(f)(x)

x = torch.randn(2, 3)
g(x)

Вызов преобразования torch.func внутри функции, обработанной с помощью torch.compile

Компиляция torch.func.grad с torch.compile

import torch

torch._dynamo.config.capture_func_transforms=True

def wrapper_fn(x):
    return torch.func.grad(lambda x: x.sin().sum())(x)

x = torch.randn(3, 3, 3)
grad_x = torch.compile(wrapper_fn)(x)

Компиляция torch.vmap с torch.compile

import torch

torch._dynamo.config.capture_func_transforms=True

def my_fn(x):
    return torch.vmap(lambda x: x.sum(1))(x)

x = torch.randn(3, 3, 3)
output = torch.compile(my_fn)(x)

Ограничения

В настоящее время существуют несколько случаев, которые не поддерживаются и приводят к разрыву графа (то есть, torch.compile возвращается к eager-режиму PyTorch в этих случаях). Мы работаем над улучшением ситуации в следующем выпуске (PyTorch 2.2).

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

import torch

torch._dynamo.config.capture_func_transforms=True

def fn(x):
    x1, x2 = x
    return x1 + x2

def my_fn(x):
    return torch.func.vmap(fn)(x)

x1 = torch.randn(3, 3, 3)
x2 = torch.randn(3, 3, 3)
# Unsupported, falls back to eager-mode PyTorch
output = torch.compile(my_fn)((x1, x2))

2. Не поддерживаются ключевые аргументы.

import torch

torch._dynamo.config.capture_func_transforms=True

def fn(x, y):
    return (x + y).sum()

def my_fn(x, y):
    return torch.func.grad(fn)(x, y=y)

x = torch.randn(3, 3)
y = torch.randn(3, 3)
# Unsupported, falls back to eager-mode PyTorch
output = torch.compile(my_fn)(x, y)

3. Функции с наблюдаемыми побочными эффектами. Например, изменение списка, созданного внутри функции, допустимо, но изменение списка, созданного вне функции, недопустимо.

import torch

torch._dynamo.config.capture_func_transforms=True

some_list = []

def f(x, y):
    some_list.append(1)
    return x + y

def my_fn(x, y):
    return torch.func.vmap(f)(x, y)

x = torch.ones(2, 3)
y = torch.randn(2, 3)
# Unsupported, falls back to eager-mode PyTorch
output = torch.compile(my_fn)(x, y)

4. torch.vmap над функцией, которая вызывает один или несколько операторов из следующего списка.

import torch

torch._dynamo.config.capture_func_transforms=True

def bad_fn(x):
    x.stride()
    return x

def my_fn(x):
    return torch.func.vmap(bad_fn)(x)

x = torch.randn(3, 3, 3)
# Unsupported, falls back to eager-mode PyTorch
output = torch.compile(my_fn)(x)

Компиляция функций помимо поддерживаемых (обходной путь)

Для других преобразований в качестве обходного пути используйте torch._dynamo.allow_in_graph

allow_in_graph — это обходной путь. Если ваш код не работает с torch.compile, которое анализирует байткод Python, но вы считаете, что он будет работать с помощью подхода символического отслеживания (например, jax.jit), тогда используйте allow_in_graph

Используя allow_in_graph для аннотации функции, вы должны убедиться, что ваш код соответствует следующим требованиям:

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

import torch

@torch.compile
def f(x):
    return torch._dynamo.allow_in_graph(torch.vmap(torch.sum))(x)

x = torch.randn(2, 3)
f(x)

Распространённая ошибка — использование allow_in_graph для аннотации функции, вызывающей nn.Module. Это связано с тем, что теперь выходы зависят от параметров nn.Module. Для того, чтобы это работало, используйте torch.func.functional_call для извлечения состояния модуля.

Работает ли NumPy с torch.compile?

Начиная с версии 2.1, torch.compile понимает собственные программы NumPy, работающие с массивами NumPy, и смешанные программы PyTorch-NumPy, которые преобразуют из PyTorch в NumPy и обратно через x.numpy(), torch.from_numpy и родственные функции.

Какие функции NumPy поддерживает torch.compile?

NumPy внутри torch.compile соответствует предварительной версии NumPy 2.0.

В целом, torch.compile способен отслеживать большинство конструкций NumPy, а когда не может, он возвращается к eager-режиму и позволяет NumPy выполнить этот фрагмент кода. Даже тогда есть несколько функций, где семантика torch.compile немного отличается от семантики NumPy:

• Скаляры NumPy: мы моделируем их как массивы 0-й размерности. То есть, np.float32(3) возвращает массив 0-й размерности в torch.compile. Чтобы избежать разрыва графа, лучше всего использовать этот массив 0-й размерности. Если это нарушает ваш код, вы можете обойти это, преобразовав скаляр NumPy в соответствующий скалярный тип Python bool/int/float.
• Отрицательные шаги: np.flip и срезы с отрицательным шагом возвращают копию.
• Преобразование типов: преобразование типов NumPy изменится в NumPy 2.0. Новые правила описаны в NEP 50. torch.compile реализует NEP 50, а не текущие правила, которые скоро устареют.
• {tril,triu}_indices_from/{tril,triu}_indices возвращает массивы, а не кортеж массивов.

Есть и другие функции, для которых мы не поддерживаем отслеживание, и мы переходим к eager-режиму NumPy для их выполнения:

• Нечисленные типы данных, такие как даты, строки, символы, пустые, структурированные типы данных и recarrays.
• Типы данных long np.float128/np.complex256 и некоторые беззнаковые типы данных np.uint16/np.uint32/np.uint64.
• ndarray подклассы.
• Маскированные массивы.
• Эзотерические механизмы ufunc, такие как axes=[(n,k),(k,m)->(n,m)] и методы ufunc (например, np.add.reduce).
• Сортировка/упорядочивание complex64/complex128 массивов.
• NumPy np.poly1d и np.polynomial.
• Позиционные out1, out2 аргументы в функциях с 2 или более возвращаемыми значениями (out=tuple работает).
• __array_function__, __array_interface__ и __array_wrap__.
• ndarray.ctypes атрибут.

Могу ли я выполнить код NumPy на CUDA через torch.compile?

Да, можете! Для этого достаточно выполнить ваш код в контексте torch.device("cuda"). Рассмотрим пример

import torch
import numpy as np

@torch.compile
def numpy_fn(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
    return np.sum(X[:, :, None] * Y[:, None, :], axis=(-2, -1))

X = np.random.randn(1024, 64)
Y = np.random.randn(1024, 64)
with torch.device("cuda"):
    Z = numpy_fn(X, Y)

В этом примере numpy_fn будет выполняться на CUDA. Для этого torch.compile автоматически перемещает X и Y с CPU на CUDA, а затем перемещает результат Z с CUDA на CPU. Если мы выполняем эту функцию несколько раз в одном запуске программы, мы можем захотеть избежать всех этих довольно дорогих копий памяти. Для этого нам просто нужно немного изменить нашу numpy_fn таким образом, чтобы она принимала тензоры CUDA и возвращала тензоры:

@torch.compile
def numpy_fn(X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor:
    X, Y = X.numpy(), Y.numpy()
    Z = np.sum(X[:, :, None] * Y[:, None, :], axis=(-2, -1))
    return torch.from_numpy(Z)

X = torch.randn(1024, 64, device="cuda")
Y = torch.randn(1024, 64, device="cuda")
with torch.device("cuda"):
    Z = numpy_fn(X, Y)

Таким образом, мы явно создаём тензоры в памяти CUDA и храним их там. В этом случае X.numpy() и from_numpy() являются подсказками для компилятора, но никакой реальной перемещения данных не происходит. Обратите внимание, что исходная программа теперь не будет работать в режиме eager. Если вы хотите запустить её в режиме eager, вам нужно вызвать .numpy(force=True) выполнив Z = Z.cuda() перед возвратом Z. Конечно, выполнив это, программа будет выполняться в NumPy в режиме eager и на CPU.

Как отладить код NumPy под torch.compile?

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

Если вышеизложенного недостаточно для определения источника проблемы, есть ещё несколько инструментов NumPy, которые мы можем использовать. Мы можем определить, исходит ли ошибка полностью из кода PyTorch, отключив отслеживание функций NumPy:

from torch._dynamo import config
config.trace_numpy = False

Если ошибка находится в отслеживаемом коде NumPy, мы можем выполнить код NumPy в режиме eager (без torch.compile) с PyTorch в качестве бэкенда, импортировав import torch._numpy as np. Это следует использовать только для целей отладки и никоим образом не заменяет API PyTorch, так как оно намного менее производительно и, будучи приватным API, может измениться без предварительного уведомления. В любом случае, torch._numpy — это Python-реализация NumPy в терминах PyTorch, и она используется внутри torch.compile для преобразования кода NumPy в код PyTorch. Она довольно легко читается и изменяется, поэтому, если вы обнаружите в ней ошибку, смело отправляйте PR с её исправлением или просто откройте вопрос.

Если программа работает при импорте torch._numpy as np, скорее всего, ошибка в TorchDynamo. В этом случае, пожалуйста, откройте вопрос с минимальным воспроизводимым примером.

Я torch.compile некоторой код NumPy и не увидел ускорения.

Лучше всего начать с урока с общими рекомендациями по отладке проблем с torch.compile.

Некоторые разрывы графика могут возникнуть из-за использования недопустимых функций. Смотрите Какие функции NumPy поддерживает torch.compile?. В более общем плане следует помнить, что некоторые широко используемые функции NumPy несовместимы с компиляторами. Например, модификации на месте затрудняют рассуждения внутри компилятора и часто приводят к худшей производительности, чем их аналоги без модификации на месте. Поэтому лучше их избегать. То же относится и к использованию параметра out=. Вместо этого используйте операции без модификации на месте, и пусть torch.compile оптимизирует использование памяти. То же самое касается операций, зависящих от данных, таких как индексирование с маской через булевы маски или управляющие конструкции, зависящие от данных, такие как if или while.

Какой API использовать для отслеживания с высокой точностью?

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

Как прервать график на функции?

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

• Если вы хотите отключить компиляцию в этом кадре функции и рекурсивно вызванных кадрах, используйте torch._dynamo.disable.
• Если вы хотите, чтобы конкретный оператор, такой как fbgemm, использовал режим eager, используйте torch._dynamo.disallow_in_graph.

Некоторые из менее распространенных случаев использования включают:

• Если вы хотите отключить TorchDynamo в кадре функции, но включить его обратно в рекурсивно вызванных кадрах, используйте torch._dynamo.disable(recursive=False).
• Если вы хотите предотвратить встраивание кадра функции, используйте torch._dynamo.graph_break в начале функции, которую вы хотите предотвратить от встраивания.

В чем разница между torch._dynamo.disable и torch._dynamo.disallow_in_graph

Disallow-in-graph работает на уровне операторов или, точнее, операторов, которые вы видите в извлеченных из TorchDynamo графах.

Disable работает на уровне кадра функции и решает, должен ли TorchDynamo рассматривать кадр функции или нет.

В чем разница между torch._dynamo.disable и torch._dynamo_skip

Вероятнее всего, вам нужен torch._dynamo.disable. Но в маловероятном случае вам может потребоваться еще более точный контроль. Предположим, вы хотите отключить отслеживание только для функции a_fn, но хотите продолжить отслеживание обратно в aa_fn и ab_fn. Ниже представлен этот случай использования:

В этом случае вы можете использовать torch._dynamo.disable(recursive=False). В предыдущих версиях эта функциональность предоставлялась torch._dynamo.skip. Теперь это поддерживается флагом recursive внутри torch._dynamo.disable.