Фиктивный тензор

Мотивация

При символьном вычислении Dynamo и выполнении проходов компилятора нам часто нужно запускать операции с тензорами, чтобы определить размеры выходных данных, типы данных и устройства, не выполняя сами эти операции (и не уничтожая существующие тензоры). Это замедлило бы работу (если выполняется много вычислений) и потребовало бы много памяти (плохо, если компилятору нужна память GPU во время компиляции программы). Фиктивный тензор во всех отношениях похож на настоящий, за исключением того, что в нём фактически нет данных. Например, при трассировке Dynamo нам нужно трассировать пользовательский код Tensor и отвечать на вопросы о промежуточных результатах (например, если пользователь выполняет условную проверку промежуточного тензора). Без фиктивных тензоров мы не смогли бы точно отвечать на такие запросы.

Аналогично, предположим, что вы хотите сохранить метаданные тензора, например, в узле FX IR (meta[‘val’]). Вместо этого можно сохранить непосредственно в узле фиктивный тензор, который предоставит все необходимые метаданные тензора, включая неочевидные свойства, которые вы, вероятно, не учли бы (например, отношения алиасинга).

Связанные работы

• Мета-тензор — это тензор с device=’meta’. Он предоставляет многое из того, что нужно от фиктивного тензора, однако мета-тензоры не моделируют устройства, а поведение stride иногда зависит от устройства. Поэтому фиктивные тензоры позволяют получить гораздо более точную информацию. Кроме того, мета-тензоры являются «глобальными» (они существуют сами по себе, подобно тензорам CPU/CUDA), тогда как фиктивные тензоры привязаны к FakeTensorMode.
• Подкласс тензора позволяет создавать подклассы torch.Tensor и настраивать их поведение. Фиктивные тензоры реализованы как подкласс тензора, поэтому почти вся их реализация написана на Python! Более простые примеры подклассов тензоров см. в subclass_zoo.
• Динамические формы позволяют создавать тензоры с символьными размерами, а не только с конкретными, и символьно распространять эти размеры при выполнении операций. Динамические формы хранят состояние в ShapeEnv, который всегда связан с FakeTensorMode (поэтому фиктивные тензоры также отвечают за управление символьными размерами). Как правило, при компиляции подграфа с помощью PT2 с этой компиляцией связан контекст трассировки, содержащий, помимо прочего, FakeTensorMode и (возможно) ShapeEnv.

Общая архитектура

Все фиктивные тензоры связаны с FakeTensorMode. Поскольку основное назначение фиктивных тензоров — анализ настоящих тензоров, обычно рабочий процесс выглядит так: у вас есть набор настоящих тензоров, вы создаёте FakeTensorMode, затем с помощью from_real_tensor преобразуете все эти настоящие тензоры в фиктивные и работаете с ними. В частности, FakeTensorMode постоянно поддерживает таблицу соответствий, связывающую тензоры (и хранилища) с одними и теми же хранилищами. Если несколько раз преобразовать один и тот же тензор в фиктивный, вы получите один и тот же фиктивный тензор; если преобразовать два тензора, которые являются алиасами друг друга, вы получите два фиктивных тензора, являющихся алиасами одного и того же фиктивного хранилища. FakeTensor — это подклассы тензоров, поэтому при выполнении операций с ними автоматически получаются фиктивные тензоры. Однако обычно операции над фиктивными тензорами (например, при выполнении прохода FX) следует выполнять с активным FakeTensorMode; операция над тензором автоматически активирует режим фиктивных тензоров и повторяет попытку.

Фиктивный тензор представлен как подкласс тензора meta-тензора с __torch_dispatch__. Это означает, что внутри фиктивные тензоры являются тензорами на устройстве meta; затем они используют дополнительные механизмы расширения, в частности dispatch_device, чтобы скрыть фактическое устройство тензора. На первых этапах разработки фиктивных тензоров это было одним из источников ошибок: иногда фиктивные тензоры слишком убедительно выдавали себя за CPU/CUDA-тензоры, и в итоге вызывалось ядро CPU, пытавшееся разыменовать указатель на данные фиктивного тензора, что, очевидно, не сработает. Если в коде фиктивных тензоров происходит ошибка сегментации, в первую очередь проверьте следующее: находится ли в обратной трассировке C++ ядро CPU (неожиданно!) или meta-ядро (ожидаемо!) Meta-ядро похоже на настоящее ядро, но только выделяет выходные данные и не выполняет вычисления над данными.

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

• Запустить meta-ядро на входных фиктивных тензорах, интерпретируя их как мета-тензоры. Это делается с помощью специального менеджера контекста in_kernel_invocation_manager, который указывает всему PyTorch рассматривать фиктивные тензоры как лежащие в их основе мета-тензоры, а не «разворачивать» фиктивные тензоры в мета-тензоры (фиктивный тензор и есть мета-тензор). Такое представление позволяет не поддерживать синхронно два набора метаданных (метаданные мета-тензора и метаданные фиктивного тензора); отношение «является» гарантирует наличие только одной канонической копии метаданных.
• Если это фабричная функция, вместо этого вызовите базовую фабричную функцию с device=’meta’.
• Преобразовать полученный мета-тензор в фиктивный тензор, вычислив, каким должно быть выходное устройство тензора (обычно это тривиально, но иногда нет, например, при продвижении скалярного значения CPU или операциях преобразования устройства).

API: основные моменты

Использование вне PT2 (дополнительные примеры см. в test/test_fake_tensor.py):

# Create a fake mode
from torch._subclasses.fake_tensor import FakeTensorMode
fake_mode = FakeTensorMode()
converter = fake_mode.fake_tensor_converter
# Fakeify some real tensors
fake_x = converter.from_real_tensor(fake_mode, x)
with fake_mode:
    # Do some operations on the fake tensors
    fake_y = fake_x * 2
    # Factory operations automatically get fakeified in the context manager
    fake_z = torch.empty(20)

В: Почему в качестве входных данных используются настоящие тензоры?

О: В контексте PT2 это связано с тем, что обычно компиляция выполняется «на лету», поэтому для всех входных данных компилируемого графа у вас уже есть «настоящие» входные данные: компиляция выполняется одновременно с запуском программы.

Использование PT2 до AOTAutograd (это необычный случай, и, скорее всего, вам так делать не нужно):

# Fake mode is not enabled!
from torch._guards import detect_fake_mode
fake_mode = detect_fake_mode(args)
# if fake_mode isn't None
converter = fake_mode.fake_tensor_converter
fake_args = [converter.from_real_tensor(fake_mode, arg) for arg in args]
with fake_mode:
    ... # do stuff with the fake args, if needed ...

detect_fake_mode выполняет поиск в нескольких местах, пытаясь найти «тот самый» режим фиктивных тензоров, связанный с жизненным циклом. Обычно он извлекается из контекста трассировки.

Использование PT2 после AOTAutograd:

# Fake mode is enabled! example_inputs is typically fake already
# TODO: we probably want to change this
# Still do this to access fake mode
fake_mode = detect_fake_mode(example_inputs)
# But in general you don't have to turn it on

Другие полезные сведения:

from torch._subclasses.fake_tensor import unset_fake_temporarily
with unset_fake_temporarily():
    ... # fake mode is disabled here, you can do real tensor compute

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

import FakeTensorProp from torch.fx.passes.fake_tensor_prop
gm: GraphModule
real_inputs: List[Tensor]
FakeTensorProp(gm).propagate(*real_inputs)
# This will populate meta['val'] on all the FX nodes with a fake tensor
# or if you have a preexisting fake mode, you should use it
FakeTensorProp(gm, mode=fake_mode).propagate(*real_inputs)
# There is also propagate_dont_convert_inputs if your inputs are already fake
fake_inputs: List[FakeTensor]
FakeTensorProp(gm, mode=fake_mode).propagate_dont_convert_inputs(*fake_inputs)

Подробности

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

with FakeTensorMode():
    real_tensor.t_()

Как должен работать этот код? Было бы неожиданно, если бы мы действительно изменили метаданные настоящего тензора. Но при этом нет очевидной возможности создать FakeTensor. Поэтому мы решили действовать консервативно и вызывать ошибку: «Вызов операторов с входными данными, не являющимися Fake Tensor, в FakeTensorMode пока не поддерживается. Сначала преобразуйте все тензоры в FakeTensors».

На практике эта ошибка весьма досадна. Например, предположим, что у вас есть настоящий nn.Module и вы хотите передать через него фиктивные тензоры. Для этого нужно каким-то образом преобразовать nn.Module в фиктивную форму. Это послужило причиной появления FakeCopyMode.

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

Изменение метаданных фиктивного тензора. Если у вас есть фиктивный тензор и вы вызываете для него t_(), его метаданные меняются. Само по себе это разумно, но иногда фиктивные тензоры также хранятся в качестве метаданных узлов FX; изменять фиктивный тензор нежелательно, поскольку это сделает устаревшими прежние метаданные!

На самом деле здесь существует фундаментальное противоречие: фиктивные тензоры поддерживают чрезвычайно точные метаданные о тензорах, вплоть до идентичности объектов. Если метаданные объекта изменяются со временем в графе FX, представить это изменение во времени невозможно. В большинстве случаев серьёзный анализ FX выполняется над функционализированными графами, у которых такой проблемы нет, но иногда анализ необходимо проводить над нефункционализированным графом. Возможно, было ошибкой помещать фиктивный тензор в meta[‘val’]

О подклассе тензора

В фиктивном тензоре используется шаблон с подклассом и режимом подкласса тензора: FakeTensor.__torch_dispatch__ активирует FakeTensorMode, связанный с фиктивным тензором, а затем выполняет повторную диспетчеризацию (полагаясь на FakeTensorMode в выполнении основной работы). Если в операциях с фиктивными тензорами встречается аргумент-подкласс, который она не распознаёт, она возвращает NotImplemented, давая другому подклассу возможность выполниться первым (в надежде, что он преобразует операцию в обычные операции с тензорами), прежде чем повторить попытку. Это может привести к бесконечным циклам.

Как реализован каждый отдельный оператор?

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

• Подклассы тензоров поддерживают ограниченное распространение констант, если число элементов очень мало (это помогает обрабатывать некоторые случаи, когда для таких тензоров сразу вызывается item()).
• Для повышения производительности некоторые операторы имеют быстрые реализации, полностью выполненные в коде фиктивных тензоров.
• Если вы используете @custom_op для создания пользовательского тензора, эти операции напрямую регистрируют impl_abstract для фиктивных тензоров.
• В самом фиктивном тензоре есть несколько жёстко заданных специальных случаев для операций преобразования устройства.
• Если нет ни реализации meta, ни декомпозиции, мы создаём настоящие тензоры, заполненные нулями, и пытаемся напрямую выполнить оператор, чтобы выяснить результат. Это может привести к ошибкам сегментации, если оператор пытается индексировать данные, поэтому для пользовательских операций эта возможность по умолчанию отключена.

Как работает преобразователь?

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

Характеристики производительности

Можно предположить, что фиктивные тензоры работают быстро, поскольку не выполняют вычислений над тензорами. Однако при небольших размерах тензоров издержки составляют практически всё время выполнения; кроме того, фиктивный тензор реализован на Python, и для одной операции с тензором часто выполняется ОЧЕНЬ много работы (поскольку операции реализованы в виде декомпозиций). Поэтому на практике фиктивные тензоры работают довольно медленно, особенно при использовании символьных форм. Сейчас в фиктивных тензорах есть два важных механизма ускорения, которые существенно влияют на производительность:

• Поточечные операции не проходят через декомпозиции: вместо этого для них вручную закодировано правило распространения.
• Если возможно, мы должны это делать.

Фиктивный тензор фиктивного тензора?

Рассматривается возможность передавать фиктивные тензоры в качестве пользовательских входных данных в стек PT2, что подразумевает необходимость создавать фиктивный тензор на основе фиктивного тензора. Сейчас такая возможность фактически не поддерживается, но, возможно, реализовать её будет не слишком сложно.

Взаимодействие с динамическими формами

Каждый FakeTensorMode содержит ShapeEnv, отслеживающий всю информацию о символьных формах. Обычно их жизненные циклы связаны: они создаются и уничтожаются вместе.

Поскольку FakeTensorMode содержит ShapeEnv (а реализации meta — нет), функции meta, зависящие от данных и требующие выделения unbacked SymInt, реализуются в фиктивных тензорах. Фиктивный тензор также отвечает за сохранение unbacked SymInt в таблице соответствий, чтобы, например, при двукратном вызове nonzero() для одного и того же фиктивного тензора возвращался один и тот же символьный размер.

Другие ресурсы

Учебное пособие Colab: использование FakeTensor для определения максимального размера пакета