✨TpuGraphs: набор данных для прогнозирования производительности на больших тензорных вычислительных графах✨

Питчайя Мангпо Фотилимтана, Сами Абу-Эль-Хайджа, Кайди Цао, Бахаре Фатеми, Чарит Мендис, Брайан Пероцци

TpuGraphs представляет набор данных для прогнозирования производительности и базовые модели тензорных программ, представленных в виде вычислительных графов, работающих на тензорных процессорах (TPU).

Важные термины и предыстория

Прогноз эффективности

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

Компилятор

  • Компиляторы преобразуют программы, написанные людьми, в эффективные исполняемые файлы для компьютерного оборудования.
  • У компиляторов две задачи — перевод и оптимизация. Они должны переводить программы в исполняемые файлы, одновременно находя наиболее эффективный перевод.
  • Автонастройка компилятора позволяет компиляторам находить оптимальную конфигурацию программ с минимальным вмешательством пользователя или вообще без него.
  • Решая задачи оптимизации, компиляторы часто используют модели производительности для оценки производительности вместо сбора измерений производительности на реальном оборудовании. Последнее может быть дорогим, ограниченным или неосуществимым.
  • Компиляторы машинного обучения (ML) решают множество задач оптимизации, чтобы преобразовать программу ML, представленную в виде тензорного графа вычислений, в эффективный исполняемый файл для аппаратных целей.

XLA (ускоренная линейная алгебра)

  • Это специализированный эвристический компилятор производственного уровня для программ ML (линейная алгебра), способный генерировать код для различных аппаратных целей, включая процессоры, графические процессоры и TPU.
  • В XLA графы тензорных вычислений представляют собой операции высокого уровня (HLO). Каждый проход оптимизации преобразует граф HLO в функционально эквивалентный.
  • В XLA имеется автотюнер, который может настраивать конфигурации TPU как на уровне графа, так и на уровне ядра.

Оптимизация на уровне ядра

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

Оптимизация на уровне графика

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

Набор данных

  • TpuGraphs — это набор данных для прогнозирования производительности тензорных программ, представляющий вычисления программы ML (один или несколько шагов обучения или один шаг вывода) в виде вычислительных графов.
  • Данные разделены на обучающие, проверочные и тестовые наборы с использованием соотношения 80–10–10 по графикам в каждой коллекции.

Пример данных

  • Каждый образец данных содержит следующее: график вычислений HLO, конфигурацию TPU v3 и время выполнения на одноядерном TPU v3.
  • Конфигурация компиляции управляет тем, как компилятор XLA преобразует граф для определенного этапа оптимизации.
  • График HLO в каждой точке данных частично оптимизируется перед подачей на соответствующий проход оптимизации.
  • Для коллекции макетов граф HLO является входным графом для этапа назначения макета. Конфигурация макета — это набор решений по макету для каждого узла на настраиваемых узлах.
  • Для коллекции плиток граф HLO в каждой точке данных представляет собой объединенный подграф, представляющий ядро. Конфигурация плитки — это конфигурация всего подграфа, не специфичная для какого-либо конкретного узла.

Коллекции

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

Генерация данных

Графики HLO:

  • Источник XLA: комбинация регрессионного теста XLA и теста MLPerf.
  • Источник НЛП: различные BERT для обучения и вывода, с разными слоями, головами внимания и скрытыми размерами. Для каждой модели собирается самый большой скомпилированный график XLA.

Конфигурации:

  • Автотюнер XLA генерирует образцы данных из графиков.
  • Режим по умолчанию (для коллекций макетов): автотюнер исследует пространство поиска с помощью генетического алгоритма, начиная с конфигурации по умолчанию, выбранной эвристикой компилятора. Конфигурации из этого режима не слишком отличаются от настроек по умолчанию и имеют аналогичное время выполнения.
  • Случайный режим (для коллекций макетов): автотюнер исследует пространство поиска, выбирая случайных кандидатов. Конфигурации из этого режима очень разные и имеют разное время выполнения.
  • Для коллекций плиток: автотюнер вызывает компилятор для запуска оптимизации на уровне графа и получения объединенных подграфов (ядер). Для каждого подграфа автотюнер перечисляет все возможные размеры тайлов для ядра в случайном порядке (ограниченном временем ожидания), поскольку пространство поиска размера тайлов намного меньше.

Извлечение функций

  • TpuGraphs предоставляет данные в необработанном формате protobuf и массивах numpy, аналогичных формату OGBG.
  • Автотюнер выдает выходные результаты в формате protobuf, которые преобразуются в массивы numpy с помощью сценария предварительной обработки данных.

Функция извлечения признаков работает следующим образом:

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

Информация

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

Коллекция макетов:

  • Layout:XLA: Random и Layout:NLP: Random — самые сложные коллекции.
  • Метод обучения сегментов графа (GST) значительно улучшает качество модели по сравнению с полным обучением графа.
  • Качество модели падает, когда размер сегмента графика слишком мал.
  • Использование меньшего количества слоев нейронной сети графа (GNN) не оказывает существенного влияния на точность.
  • Выбор алгоритма разделения (METIS или Topo Partition) не имеет большого значения.
  • Парная шарнирная потеря значительно лучше представляет собой функцию потерь, чем среднеквадратическая ошибка (MSE).

Коллекция плиток:

  • Более раннее объединение функций конфигурации с функциями узла (раннее соединение) превосходит объединение функций конфигурации с более поздним внедрением сокращенного графа (позднее соединение).
  • Использование потери рейтинга (ListMLE) более эффективно, чем использование MSE.
  • GraphSAGE — лучший выбор для GNN, чем GCN.
  • GNN необходим для хорошей точности (более высокая производительность, чем у LSTM и Transformer, согласно предыдущей работе).

Как добраться

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

Для дополнительной информации:

📜Статья: https://arxiv.org/abs/2308.13490

⭐GitHub: https://github.com/google-research-datasets/tpu_graphs