Обслуживание модели с отслеживанием состояния: как мы ускоряем вывод с помощью среды выполнения ONNX

Есть разница между обслуживанием моделей без сохранения состояния и с сохранением состояния.

Авторы Лестер Солбаккен из Verizon Media и Пранав Шарма из Microsoft.

Есть разница между обслуживанием моделей без сохранения состояния и с сохранением состояния.

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

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

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

В Vespa.ai мы заботимся об эффективной оценке модели с отслеживанием состояния. Vespa.ai - это платформа с открытым исходным кодом для создания приложений, которые в реальном времени обрабатывают большие наборы данных. Разработанный для обеспечения высокой производительности и масштабируемости в Интернете, он используется для таких разнообразных задач, как поиск, персонализация, рекомендации, реклама, автозаполнение, поиск изображений и сходств, ранжирование комментариев и даже для поиска любви.

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

В конце концов мы решили использовать ONNX Runtime (ORT) для этой задачи. ONNX Runtime - это ускоритель вывода модели. Это значительно расширило возможности Vespa.ai по оценке больших моделей как по характеристикам, так и по типам моделей, которые мы поддерживаем. Возможности ONNX Runtime в области аппаратного ускорения и оптимизации моделей, таких как квантование, позволили эффективно оценивать большие модели NLP, такие как BERT и другие модели Transformer в Vespa.ai.

В этом посте мы расскажем о том, почему и как мы в конечном итоге выбрали ONNX Runtime, и поделимся своим опытом работы с ней.

О Vespa.ai

Vespa.ai имеет богатую историю. Его происхождение происходит от поисковой системы, появившейся в 1997 году. Первоначально она использовалась для поиска в Интернете на alltheweb.com, но была достаточно гибкой для использования в различных более специализированных продуктах или вертикалях, таких как поиск документов, мобильный поиск, желтые страницы и банковское дело. . Эта гибкость вертикальной поисковой платформы в конечном итоге дала название Vespa.

Эта технология была приобретена Yahoo в 2003 году. Там Vespa зарекомендовала себя как ключевой элемент технологии, обеспечивающий работу сотен приложений, включая многие из наиболее важных сервисов Yahoo. Мы открыли исходный код Vespa в 2017 году, и сегодня он обслуживает сотни тысяч запросов в секунду по всему миру в любой момент времени с миллиардами элементов контента для сотен миллионов пользователей.

Хотя в конечном итоге Yahoo была приобретена Verizon, интересно отметить, что наша команда на протяжении многих лет оставалась на удивление стабильной. Действительно, несколько инженеров, которые начали работать над первым двигателем более 20 лет назад, все еще здесь. Наша команда насчитывает около 30 разработчиков, и мы находимся в Тронхейме в Норвегии.

Основываясь на опыте, накопленном за многие годы, Vespa.ai претерпела существенные изменения и стала тем, чем является сегодня. Теперь он выступает в качестве проверенного в боях общего движка для вычислений в реальном времени над большими наборами данных. Он имеет множество функций, которые делают его подходящим для веб-приложений. Он хранит и индексирует данные с мгновенной записью, так что запросы, выбор и обработка данных могут выполняться эффективно во время обслуживания. Он эластичен и отказоустойчив, поэтому узлы можно добавлять, удалять или заменять в реальном времени без потери данных. Его легко настроить, использовать и добавить собственную логику. Важно отметить, что он содержит встроенные возможности для расширенных вычислений, включая модели машинного обучения.

Приложения Vespa.ai

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

Пакеты приложений содержат одну или несколько схем документов. Схема в основном состоит из:

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

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

После того, как приложение развернуто, оно готово к обработке записи данных и запросов. Каналы данных сначала обрабатываются на уровне без сохранения состояния, прежде чем контент будет распределен (с избыточностью) на узлы контента. Точно так же запросы проходят через уровень без сохранения состояния перед тем, как разветвляться на узлы контента, где обрабатываются зависимые от данных вычисления. Они возвращают свои результаты обратно на уровень без сохранения состояния, где определяются глобально лучшие результаты, и в конечном итоге возвращается ответ.

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

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

ONNX в Vespa.ai

В последние несколько лет для Vespa.ai становится все более важным поддерживать различные типы моделей машинного обучения из разных фреймворков. Это привело к тому, что в 2018 году мы представили начальную поддержку моделей ONNX.

Open Neural Network Exchange (ONNX) - это открытый стандарт для распределения моделей машинного обучения между различными системами. Целью ONNX является взаимодействие между структурами обучения моделей и механизмами вывода, избегая привязки к какой-либо привязке к поставщику. Например, библиотека Transformer от HuggingFace включает экспорт в ONNX, PyTorch имеет собственный экспорт в ONNX, а модели TensorFlow можно конвертировать в ONNX. С нашей точки зрения, поддержка ONNX, очевидно, интересна, поскольку она позволит максимально расширить диапазон поддерживаемых моделей.

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

document my_document {
  field text_embedding type tensor(x[769]) {
    indexing: attribute | index  
    attribute {
      distance-metric: euclidean
    }
  }
  field text_tokens type tensor(d0[256]) {
    indexing: summary | attribute
  }
}
onnx-model my_model {
  file: files/my_model.onnx
  input input_0: ...
  input input_1: ...
  output output_0: ...
}
rank-profile my_profile {
  first-phase {
    expression: closeness(field, text_embedding)
  }
  second-phase {
    rerank-count: 10
    expression: onnx(my_model)
  }
}

Это пример настройки Vespa.ai для вычисления евклидова расстояния между вектором запроса и сохраненным вектором text_embedding на первом этапе. Обычно это используется вместе с приблизительным поиском ближайшего соседа. На втором этапе 10 лучших кандидатов отправляются в модель ONNX. Обратите внимание, что это для каждого узла содержимого, поэтому с 10 узлами содержимого модель эффективно работает на 100 кандидатах.

Модель настраивается в разделе onnx-model. Файл относится к модели ONNX где-то в пакете приложения. Входные данные в модель, хотя на самом деле здесь не показаны для краткости, могут поступать из различных источников, таких как константы, запрос, документ или некоторая комбинация, выраженная через пользовательскую функцию. В то время как выходные данные моделей являются тензорами, результирующее значение выражения первой или второй фазы должно быть одним скаляром, поскольку документы сортируются в соответствии с этой оценкой перед возвратом.

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

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

Изначально это работало нормально. Мы реализовали различные операторы ONNX, используя доступные тензорные операции. Однако сначала мы поддерживали только подмножество из 150+ операторов ONNX, так как считали, что только определенные типы моделей могут быть пригодны для использования в Vespa.ai из-за его требований к низкой задержке. Например, язык выражений ранжирования не поддерживает итерации, что затрудняет реализацию операторов, используемых в сверточных или рекуррентных нейронных сетях. Вместо этого мы решили постоянно добавлять поддержку операторов, поскольку в Vespa.ai использовались новые типы моделей.

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

ONNX Runtime

К сожалению, мы столкнулись с проблемами, когда начали разработку поддержки моделей Transformer. Наша первая попытка поддержать 12-слойную базовую модель BERT не удалась. Это была модель, преобразованная из TensorFlow в ONNX. Результат оценки был неверным, с относительно низкой производительностью.

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

Во время разработки мы обратились за справкой к ONNX Runtime. ONNX Runtime очень прост в использовании:

import onnxruntime as ort
session = ort.InferenceSession(“model.onnx”)
session.run( output_names=[...], input_feed={...} )

Это было бесценно, так как давало нам ссылку на правильность и целевую производительность.

В какой-то момент мы начали обдумывать идею фактического использования ONNX Runtime непосредственно для вывода модели вместо преобразования ее в выражения Vespa.ai. Узел содержимого Vespa.ai написан на C ++, поэтому для этого потребовалась интеграция с интерфейсом C ++ среды выполнения ONNX. Следует отметить, что добавление зависимостей в Vespa.ai - это не то, что мы часто делаем, поскольку мы предпочитаем избегать зависимостей и, таким образом, владеть всем стеком.

В течение пары недель у нас была испытанная концепция, которая показала многообещающие результаты. Поэтому мы решили пойти дальше и начать использовать ONNX Runtime для оценки всех моделей ONNX в Vespa.ai.

Это оказалось для нас переломным моментом. Это значительно увеличивает возможности оценки больших моделей глубокого обучения в Vespa.ai с точки зрения типов моделей, которые мы поддерживаем, и производительности оценки. Мы можем использовать использование ONNX Runtime MLAS, вычислительной библиотеки, содержащей ядра, оптимизированные для процессора. ONNX Runtime также содержит оптимизацию для моделей BERT (например, слияние нескольких узлов внимания) и упрощает оценку моделей с пониженной точностью путем квантования для еще более эффективного вывода.

ONNX Runtime в Vespa.ai

Учтите следующее:

onnx-model my_model {
  file: files/my_model.onnx
  input input_0: query(my_query_input)
  input input_1: attribute(my_doc_field)
}
rank-profile my_profile {
  first-phase {
    expression: sum( onnx(my_model) )
  }
}

Здесь у нас есть одна модель ONNX с двумя входами. Во время развертывания приложения Vespa.ai распространяет эту модель ONNX на все узлы содержимого. Там выражения ранжирования анализируются, а исполнители функций, реализующие функции ранжирования, настраиваются для подготовки к обработке трафика. Здесь у нас есть 4 функции:

• query(...), который извлекает тензор из запроса.
• attribute(...), который извлекает тензор из поля, хранящегося в документе.
• onnx(...), который оценивает модель ONNX.
• sum(...), который уменьшает и суммирует тензор аргумента до единственного скалярного значения.

Эти функции связаны друг с другом во время инициализации, поэтому выходы query и attribute используются как входы для onnx, а выход функции onnx является входом для функции sum. Функция onnx в основном настраивает среду выполнения ONNX для оценки модели.

Система оценки Vespa.ai написана на C ++, поэтому мы используем API C / C ++, предоставляемый ONNX Runtime. Хотя интеграция с ONNX Runtime работала без сбоев, здесь стоит упомянуть две области: многопоточность и распределение тензоров ввода / вывода.

Многопоточность

Во время установки мы инициализируем сеанс выполнения ONNX для каждой функции и потока onnx:

#include <onnxruntime/onnxruntime_cxx_api.h>
Ort::Env shared_env;
Ort::SessionOptions options;
options.SetIntraOpNumThreads(1);
options.SetInterOpNumThreads(1);
options.SetGraphOptimizationLevel(ORT_ENABLE_ALL);
Ort::Session session = Ort::Session(shared_env, “local_file_path”, options);

В сеансе есть опции для управления потоками. ONNX Runtime поддерживает 2 режима выполнения: последовательный и параллельный. Это контролирует, будут ли операторы на графе работать последовательно или параллельно. Параллельное выполнение операторов планируется в пуле межоперационных потоков. Выполнение отдельного оператора распараллеливается с использованием пула потоков внутри операции. Сильно оптимизированный вариант пула собственных потоков используется для межоперационного параллелизма, в то время как OpenMP используется для внутриоперационного.

Vespa.ai обрабатывает несколько запросов параллельно. Кроме того, Vespa.ai можно настроить на использование нескольких потоков для каждого запроса. Из-за этого Vespa.ai необходимо жестко контролировать использование потоков. Использование дополнительных потоков внутри сеанса ONNX Runtime приводит к тому, что пропускная способность на системном уровне становится непредсказуемой, с большими отклонениями в производительности. Поскольку Vespa.ai имеет управление потоками вне среды выполнения ONNX, нам нужно указать среде выполнения ONNX использовать только один поток. Убедившись, что общее количество потоков не превышает количество физических ядер в машине, мы можем улучшить использование кеша. Vespa также поддерживает закрепление процессора.

Поскольку мы инструктируем среду выполнения ONNX работать последовательно, время вывода увеличивается, но общая пропускная способность также увеличивается. Например, мы измерили увеличение пропускной способности на 50% в приложении для ранжирования BERT. Мы еще не раскрыли настройки управления потоками ONNX Runtime для случаев, когда пользователи хотели бы настроить их самостоятельно. Это вариант, который мы могли бы рассмотреть в будущем. В этом случае каждый сеанс, имеющий собственный набор пулов потоков, будет неэффективным. Однако среда выполнения ONNX предоставляет возможность совместного использования пулов потоков между сеансами. Это достигается с помощью CreateEnvWithGlobalThreadPools C API для настройки объекта shared_env, который в Vespa.ai используется всеми исполнителями функций.

Когда мы начали использовать ONNX Runtime, его дистрибутив C ++ был связан с OpenMP. Это было проблематично для нас, поскольку настройка потока внутри операции была переопределена OpenMP, поэтому в итоге мы скомпилировали нашу собственную среду выполнения ONNX без включенного OpenMP. Однако, начиная с версии 1.6, ONNX Runtime поставляет версию без OpenMP.

Входные и выходные тензоры

Насколько это возможно, распределение памяти и владение тензорами ввода и вывода происходит внутри Vespa.ai. Рассмотрим следующие типы:

std::vector<const char *> input_names;
std::vector<const char *> output_names;
std::vector<Ort::Value>   input_values;
std::vector<Ort::Value>   output_values;

Входные значения поступают из других функций ранжирования, использующих тензорную структуру Vespa.ai. Значения во входном векторе являются оболочками для тензоров Vespa.ai. Таким образом, ONNX Runtime принимает схему памяти от Vespa.ai без копирования во внутренние буферы. Значения в выходном векторе - это заранее выделенные векторы среды выполнения ONNX, которые обертываются при последующем использовании в других функциях ранжирования.

Мы используем их непосредственно при оценке модели:

Ort::RunOptions run_opts(nullptr);
session.Run(run_opts,
    input_names.data(), input_values.data(), input_values.size(),
    output_names.data(), output_values.data(),output_values.size());

Очевидно, что такое нулевое копирование тензоров желательно с точки зрения производительности. Это работает для выходных данных, поскольку тензоры в Vespa.ai в настоящее время имеют фиксированный размер, что означает, что размеры известны во время развертывания приложения. Таким образом, даже если модели принимают входные данные с динамическими размерами, с точки зрения Vespa.ai, в настоящее время они должны быть исправлены. Поддержка динамических размеров - это работа будущего.

Одним из ограничений здесь является то, что Vespa.ai в настоящее время поддерживает только типы значений double и float в тензорах. Если возможно, Vespa.ai позаботится о преобразовании в тип, ожидаемый средой выполнения ONNX. Например, входные данные для моделей Transformer - это последовательности токенов, как правило, типа int64. Поскольку Vespa.ai в настоящее время не поддерживает типы int, они должны быть представлены, например, как float32. Преобразование из float32 в int64 может происходить с потерями, но мы пока не встретили никаких неточностей. Модели, которые принимают строки в качестве входных данных, еще не поддерживаются в Vespa.ai.

Резюме

Интеграция с ONNX Runtime прошла относительно безболезненно. Несмотря на то, что изначально нам пришлось скомпилировать собственный дистрибутив ONNX Runtime из-за OpenMP, у нас не было серьезных проблем с зависимостями. Мы с нетерпением ждем следующего выпуска, поэтому нам не нужно этого делать.

Хотя документация по C / C ++ API в ONNX Runtime в настоящее время относительно скудна, мы сочли ее достаточной. ONNX Runtime API конкретен, чист и работает так, как и следовало ожидать. На самом деле, у нас здесь вообще не было никаких проблем.

В целом, наш опыт работы с ONNX Runtime был отличным. Он отлично зарекомендовал себя для Vespa.ai, обеспечивая превосходную производительность и простой в использовании. Одним из примеров здесь является оптимизация модели BERT в среде выполнения ONNX. Также следует особо отметить оценку квантованных моделей, которую было бы сложно реализовать в Vespa.ai, учитывая текущие типы тензоров, которые мы поддерживаем.

Забегая вперед, есть и другие особенности ONNX Runtime, с которыми мы хотели бы поэкспериментировать. Одно из них - поддержка графического процессора. Хотя в настоящее время мы не уверены, выгодно ли это с точки зрения ранжирования, есть и другие случаи, когда это интересно. Одним из примеров является вычисление векторных представлений документов во время записи данных из модели Transformer.

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

Мы очень рады видеть, для чего наши пользователи будут использовать это в будущем.