API TorchDynamo для точечного прослеживания
Примечание
В этом документе torch.compiler.compile и torch.compile используются взаимозаменяемо. Обе версии будут работать в вашем коде.
torch.compile выполняет прослеживание TorchDynamo для всей пользовательской модели. Однако возможно, что небольшой фрагмент кода модели не может быть обработан torch.compiler. В этом случае вы можете отключить компилятор для этой конкретной части, выполнив компиляцию остальной части модели. Этот раздел описывает существующие API, которые используются для определения частей вашего кода, в которых вы хотите пропустить компиляцию, и соответствующие примеры использования.
API, которые можно использовать для определения частей кода, для которых можно отключить компиляцию, перечислены в следующей таблице:
API | Описание | Когда использовать? |
|---|---|---|
| Отключает Dynamo для декорированной функции, а также для рекурсивно вызываемых функций. | Идеально подходит для освобождения пользователя, если небольшая часть модели не может быть обработана |
| Запрещает указанный оператор в графе TorchDynamo. TorchDynamo вызывает разрыв графа и выполняет оператор в режиме eager (без компиляции). Это подходит для операторов, в то время как | Этот API отлично подходит для отладки и освобождения от блокировок, если пользовательский оператор, например, |
| Аннотируемый вызываемый объект передаётся в граф TorchDynamo как есть. Например, чёрного ящика для TorchDynamo Dynamo.nnОбратите внимание, что AOT Autograd будет проходить через него, поэтому | Этот API полезен для частей модели, которые имеют известные сложноподдерживаемые TorchDynamo особенности, такие как хуки или |
| Добавляет разрыв графа. Код перед и после разрыва графа проходит через TorchDynamo. | Редко полезно для развертывания — если вы считаете, что вам это нужно, скорее всего, вам нужен либо |
torch.compiler.disable
torch.compiler.disable отключает компиляцию для декорированной функции и всех функций, рекурсивно вызываемых из декорированной функции.
TorchDynamo перехватывает выполнение каждой функции Python. Предположим, у вас есть структура кода (изображение ниже), где функция fn вызывает функции a_fn и b_fn . А a_fn вызывает aa_fn и ab_fn . Когда вы используете режим eager PyTorch вместо torch.compile, эти функции выполняются как есть. С torch.compile, TorchDynamo перехватывает каждую из этих функций (указано зелёным цветом):
Предположим, что функция a_fn вызывает проблемы с torch.compile. И это некритическая часть модели. Вы можете использовать compiler.disable для функции a_fn. Как показано выше, TorchDynamo прекратит просматривать фреймы, которые произошли от вызова a_fn (белый цвет указывает на исходное поведение Python).
Чтобы пропустить компиляцию, вы можете декорировать проблемную функцию с помощью @torch.compiler.disable.
Вы также можете использовать синтаксис без декораторов, если не хотите изменять исходный код. Однако мы рекомендуем избегать этого стиля, если возможно. Здесь вам нужно следить за тем, чтобы все пользователи исходной функции теперь использовали изменённую версию.
torch._dynamo.disallow_in_graph
torch._dynamo.disallow_in_graph запрещает оператор, но не функцию, присутствовать в извлечённом графе TorchDynamo. Обратите внимание, что это подходит для операторов, а не для общих функций, как в случае с _dynamo.disable.
Предположим, вы компилируете свою модель с PyTorch. TorchDynamo может извлечь граф, но затем вы видите, что компилятор в конце концов терпит неудачу. Например, отсутствует мета-ядро или неправильно задан некоторый ключ dispatch Autograd для определённого оператора. Тогда вы можете пометить этот оператор как disallow_in_graph, и TorchDynamo вызовет разрыв графа и выполнит этот оператор, используя режим eager PyTorch.
Суть в том, что вам нужно найти соответствующий оператор Dynamo-уровня, а не оператор ATen-уровня. Подробнее см. в разделе Ограничения документа.
Предупреждение
torch._dynamo.disallow_in_graph — это глобальный флаг. Если вы сравниваете разные компиляторы backends, вам может потребоваться вызвать allow_in_graph для запрещённого оператора при переключении на другой компилятор.
torch.compiler.allow_in_graph
torch.compiler.allow_in_graph полезно, когда связанный кадр функции имеет некоторые известные сложноподдерживаемые особенности TorchDynamo, такие как хуки и autograd.Function, и вы уверены, что последующие компоненты PyTorch, такие как AOTAutograd, могут безопасно отслеживать декорированную функцию. Когда функция декорируется с allow_in_graph, TorchDynamo рассматривает её как чёрный ящик и помещает её как есть в сгенерированный граф.
Предупреждение
allow_in_graph полностью пропускает TorchDynamo для декорированной функции, отключая все проверки безопасности TorchDynamo, включая разрывы графа, обработку замыканий и другие. Используйте allow_in_graph с осторожностью. Последующие компоненты PyTorch, такие как AOTAutograd, полагаются на TorchDynamo для обработки сложных функций Python, но allow_in_graph обходит TorchDynamo. Использование allow_in_graph может привести к проблемам с корректностью и трудноотлаживаемым ошибкам.
Ограничения
Все существующие API применяются на уровне TorchDynamo. Следовательно, эти API имеют доступ только к тому, что видит TorchDynamo. Это может привести к путанице.
Например, torch._dynamo.disallow_in_graph не будет работать для операторов ATen, так как они видны для AOT Autograd. Например, torch._dynamo.disallow_in_graph(torch.ops.aten.add) не будет работать в вышеуказанном примере.
© 2024, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://pytorch.org/docs/2.1/torch.compiler_fine_grain_apis.html
