Динамические формы

Код: symbolic_shapes.py

См. также: Руководство по динамическим формам

Мотивация

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

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

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

Краткое описание публичного API

По умолчанию в PyTorch 2.1 используется динамическое поведение:

• PT2 по умолчанию предполагает, что все статично
• Если мы перекомпилируем из-за изменения размера, мы вместо этого попытаемся перекомпилировать этот размер как динамический (вероятность изменения изменённых размеров в будущем). Эта обобщение может потерпеть неудачу (например, из-за того, что пользовательский код выполняет условный переход по соответствующему размеру или отсутствует поддержка динамических форм в PT2). Если вы пытаетесь понять, почему PT2 чрезмерно специализировал некоторый код, запустите с TORCH_LOGS=dynamic и найдите записи «eval», которые показывают, когда добавляются блокировки и почему.
• Если заранее известно, что что-то будет динамическим, вы можете пропустить первую перекомпиляцию с помощью torch._dynamo.mark_dynamic(tensor, dim).
• Если вы укажете torch.compile(dynamic=False), мы отключим автоматическую динамику форм при перекомпиляции и всегда будем перекомпилировать для каждого отдельного размера. И наоборот, если вы укажете torch.compile(dynamic=True), мы постараемся сделать все максимально динамичным. Это в основном полезно для небольших операторов; если вы попробуете это на большой модели, то (1) вероятно, она вызовет сбой PT2, и (2) будет работать медленно без видимых причин.

Модель защиты

При рассмотрении способов добавления поддержки динамических форм в TorchDynamo и TorchInductor мы приняли важное архитектурное решение: для повторного использования декомпозиций и другого предварительно написанного кода на Python/C++, ориентированного на API PyTorch, мы должны иметь возможность проследить динамические формы. В отличие от полностью символической системы, которая может захватывать оба разветвления условного оператора, мы всегда выбираем одно разветвление и специализируем наше прослеживание с предположением, что мы используем это прослеживание только тогда, когда мы бы сделали тот же выбор для этого разветвления в будущем. Для этого мы поддерживаем «подсказку» для каждого символического размера, указывающую его конкретное значение во время компиляции (поскольку TorchDynamo является компилятором Just-in-Time, он всегда знает фактические размеры входных данных). Когда мы выполняем условное выражение над тензором, мы просто обращаемся к подсказке, чтобы узнать, какое разветвление выбрать.

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

def f(x, y):
    z = torch.cat([x, y])
    if z.size(0) > 2:
        return z.mul(2)
    else:
        return z.add(2)

Конечная IR, которую мы будем компилировать с помощью TorchInductor, будет либо torch.cat([x, y]).add(2), либо torch.cat([x, y]).mul(2) (при условном выражении, удалённом), но чтобы определить, в каком разветвлении мы находимся, нам нужно знать размер z, промежуточного значения. Поскольку TorchDynamo должен заранее знать, является ли скомпилированный след допустимым (мы не поддерживаем откаты, как некоторые компиляторы JIT), мы должны уметь сократить z.size(0) в качестве выражения в терминах входных данных, x.size(0) + y.size(0). Это делается путём написания мета-функций для всех операторов в PyTorch, которые могут распространять информацию о размере к выходу тензора, не выполняя фактическое вычисление в узле.

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

Рабочий процесс символических форм:

1. Когда мы начинаем компиляцию фрейма в Dynamo, мы выделяем ShapeEnv (прикреплённый к FakeTensorMode), который отслеживает состояние символических форм.
2. Мы выделяем символические размеры для тензоров при входе (то, что статично или динамично, является политическим решением с некоторыми переключателями).
3. Мы распространяем символические размеры через операторы, сохраняя как (1) FX IR, чтобы мы могли корректно экспортировать символическое вычисление, и (2) выражения Sympy, представляющие переменные размера, чтобы мы могли рассуждать о них.
4. При условном выражении над символическими размерами, как при прослеживании Dynamo, так и при оптимизации Inductor, мы добавляем защитные механизмы на основе условного оператора. Эти механизмы могут быть индуцированы как из Python, так и из C++.
5. Эти блокировки могут вызвать дальнейшие упрощения символических переменных. Например, если вы утверждаете s0 == 4, мы можем теперь заменить все вхождения s0 на 4.
6. Когда мы закончим прослеживание и оптимизацию, мы устанавливаем все эти блокировки со скомпилированным кодом; скомпилированный код является повторно используемым только если все блокировки имеют истинное значение.

Важные файлы:

• API SymInt на C++: c10/core/SymInt.h, SymFloat.h, SymBool.h
• API SymInt на Python: torch/__init__.py (поищите SymInt/SymFloat/SymBool)
• C++ утилиты: c10/core/SymNodeImpl.h, torch/csrc/utils/python_symnode.h, torch/csrc/jit/python/init.cpp
• Инфраструктура Python: torch/fx/experimental/symbolic_shapes.py
• Другие важные файлы: torch/_subclasses/fake_tensor.py, torch/_meta_registrations.py, декомпозиции, ссылки PrimTorch

Краткое описание внутреннего API

Понимание иерархии классов Python:

• SymInt/SymFloat/SymBool: это видимые для пользователя классы, которые имитируют свои аналоги int/float/bool. Если вы складываете два SymInt, мы дадим вам новый SymInt, который символически отслеживает, что сложение целых чисел произошло.
• SymNode: это внутренняя структура (доступная, например, через symint.node), которая содержит фактическую информацию о символическом отслеживании. SymNode имеет тип-очистку; это делает его более удобным для представления операций смешанных типов. Обратите внимание, что технически вам не обязательно вызывать SymNode Python из SymInt; например, C++ XLA SymNodeImpl возьмёт на себя роль SymNode.
• ShapeEnv: состояние контекста на время компиляции, которое отслеживает все свободные символы и защитные механизмы, которые мы накопили до сих пор. Каждый SymNode записывает свой ShapeEnv (но не наоборот; SymNode используются только если они участвуют в защитных механизмах).

C++ довольно похож:

• c10::SymInt/SymFloat/SymBool: видимые для пользователя классы, которые имитируют int/float/bool.
• c10::SymNode/SymNodeImpl: аналогично SymNode
• ShapeEnv в C++ отсутствует; для удобства отладки весь аппарат символического рассуждения находится в Python.

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

Политика DimDynamic

Символическое рассуждение:

• Диапазоны значений
• Примечания по использованию Sympy
• Ограничения
• DimDynamic/Constraint

Неподтверждённые SymInt

Для разрешения потока управления мы проверяем подсказку, т.е. фактическое значение, символического целого числа, чтобы определить, по какому пути идти. Однако в некоторых случаях у нас может не быть подсказки: так называемые неподтверждённые символические целые числа возникают, когда переменная размера появляется из операции, зависящей от данных, такой как .nonzero() или .item(). Незаконно выполнять поток управления по этим символическим целым числам, поэтому мы должны прервать построение графа на этих операциях.

При простом реализации это слишком ограничивает: большинство программ PyTorch немедленно завершатся неудачей, если вы попытаетесь сделать что-либо с неподтверждёнными символическими целыми числами. Вот самые важные улучшения, чтобы сделать это действительно работающим:

• При создании тензора PyTorch предварительно вычисляет много данных о тензоре; например, если вы используете empty_strided для создания тензора, мы будем усердно сортировать шаги и определять, является ли тензор неперекрывающимся и плотным. Сортировки производят много защитных механизмов. Однако чаще всего тензор создаётся напрямую с помощью API более высокого уровня, например empty, который гарантирует создание неперекрывающегося и плотного тензора. Мы изменили PyTorch, чтобы избежать ненужного повторного вычисления этих свойств.
• Даже если требуется нетривиальное вычисление, иногда свойство вообще не запрашивается. Делание этих предварительно вычисленных свойств ленивыми позволяет нам избежать защитных механизмов на неподтверждённом символическом целом числе, пока оно фактически не понадобится.
• Данные в целочисленном тензоре, как правило, не известны как неотрицательные. Однако мы предоставляем API constrain_range, посредством которого пользователь может указать, что размер ограничен сверху и снизу известными пределами.

В будущих версиях PT2 (после PT2.1) мы расширим нашу систему рассуждения, чтобы сделать вывод, что неподтверждённое символическое целое число является размером на основе использования. Например, если вы передадите результат вызова .item() в функцию-фабрику, такую как torch.empty, мы автоматически выведем, что результат является размером (потому что в противном случае это было бы ошибкой). Это предположение будет проверяться во время выполнения, вызывая ошибку, если это не будет выполнено.