Ложный тензор

Код: fake_tensor.py

Мотивация

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

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

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

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

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

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

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

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

• Запустите мета-ядро на входных ложных тензорах, переинтерпретируя их как мета-тензоры. Это делается с помощью магического менеджера контекста 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()
# Fakeify some real tensors
fake_x = fake_mode.from_real_tensor(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)
fake_args = [fake_mode.from_real_tensor(arg) for arg in args]
with fake_mode:
... do stuff with the fake args, if needed ...

detect_fake_mode будет искать ряд мест, чтобы попытаться найти «режим» ложного тензора, связанный с жизненным циклом. Как правило, он будет извлечен из контекста отслеживания.

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

# Режим Fake включен! example_inputs обычно уже ложный # TODO: возможно, нам нужно это изменить # Все еще делаем это, чтобы получить доступ к режиму Fake fake_mode = detect_fake_mode(example_inputs) # Но в общем случае вам не нужно его включать

Другие полезные вещи:

from torch.fx.experimental.proxy_tensor import maybe_disable_fake_tensor_mode
with maybe_disable_fake_tensor_mode():
    # fake mode is disabled here, you can do real tensor compute

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

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_()

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

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

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

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

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

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

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

Как реализуется каждая отдельная операция?

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

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

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

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

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

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

• Точечные операции не проходят через PrimTorch-декомпозиции, вместо этого мы разработали правила их распространения.
• Если возможно, мы должны.

Ложный тензор ложного тензора?

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

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

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

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

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

Обучающее пособие по использованию FakeTensor для определения максимального размера пакетной обработки