Всем хорошего дня! На связи с вами Домклик и Денис Захаров из команды чат-ботов. В этой статье я расскажу, как можно ускорить нейронки, не прибегая к закупке дорогостоящего оборудования. Статья написана по мотивам моего выступления на конференции HighLoad++ 2024.

Многие задачи машинного обучения требуют, чтобы ответ от модели пришёл как можно быстрее. Например, в голосовом боте у нас есть несколько миллисекунд на ответ, чтобы не раздражать пользователя и не ломать взаимодействие. Выход кажется простым: развернуть модель на GPU. Но мы пошли другим путём и научились ускорять модель на CPU так, что в некоторых случаях она обгоняет GPU.

Этот обзор о том, как мы используем нейронные сети в голосовых и текстовых ботах Домклик и почему решили не переезжать на GPU. Также на примере RoBERTa покажу, из чего состоит модель, как ускорить каждую её часть и как развернуть полученные артефакты в прод.

О команде и наших задачах

Я работаю в команде голосовых и текстовых ботов. Немного о том, что мы делаем. Нас можно найти:

• В чате Домклик: мы используем нейронные сети, чтобы оперативно отвечать на вопросы пользователей без подключения операторов

• При звонке на 900: если возник вопрос по ипотеке, то, скорее всего, в первую очередь пользователь столкнётся с нашим ботом

• При создании объявления: мы научились генерировать объявления о продаже с помощью LLM

• В поиске: недавно внедрили улучшения в поиск на главной странице, теперь стараемся показывать только те статьи и сервисы Домклик, которые релевантны запросу конкретного человека

В чате и голосовом боте определяем, что сказал человек, классифицируем сообщение и проводим его по соответствующему пользовательскому пути, чтобы максимально точно и понятно ответить на вопросы. С помощью LLM и векторных баз данных генерируем объявления и выдаём релевантные результаты в поиске.

Основа наших систем: RoBERTa

Ядром многих наших систем является векторизатор RoBERTa, обученный собственными аналитиками-исследователями (data scientist). С помощью этой модели мы генерируем текстовые эмбеддинги, которые затем используем для классификации, поиска по векторам в базе и создания различных решений.

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

Почему не используем GPU

Если мы нацелены на скорость, то почему бы не воспользоваться самым распространённым и, казалось бы, простым решением — развернуть модель на GPU? У нас на это несколько причин:

• Другим командам нужнее. В Домклик есть команды, работающие с более ресурсоёмкими задачами: например, распознавание текста и классификация документов. Для этого они используют довольно тяжёлые модели, иногда целые каскады, которые тоже хотят отвечать быстро. Таким образом, можно сказать, что для них GPU приоритетнее.

• Не хотелось тратить своё пространство. Пространство на GPU не бесконечно. У нас есть проекты, требующие ресурсов для работы с большими и мощными LLM. Для этого мы и планируем использовать GPU, поэтому на данный момент не хотим тратить ресурсы на что-то другое.

• Неудачные попытки оптимизации инференса. После немногочисленных подходов к оптимизации, которые не увенчались успехом, мы прекратили дальнейшие попытки.

• Высокая стоимость. Инференс на GPU — дорогое удовольствие: выше стоимость, больший объем потребляемой энергии, а значит, не самый оптимальный вариант с финансовой точки зрения.

Подходы к ускорению RoBERTa

Ранее у нас было два подхода.

1. Использование библиотеки Sentence Transformers

Самый простой подход заключался в использовании библиотеки Sentence Transformers, благодаря которой наши аналитики-исследователи обучали модель.

Достоинства:

• Простота и понятность разработки

• Необходимые метрики, потому что можем сами навесить всё, что нужно

• Большой BUS-фактор

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

Недостаток: длительное время ответа от модели (latency).

2. Использование TorchServe

Далее для ускорения инференса мы попробовали использовать TorchServe.

Достоинство: небольшой прирост производительности.

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

Недостатки:

• Сложный запуск

• Дополнительная конвертация модели

• Метрики только те, что есть

• Увеличение потребляемых ресурсов

• Дольше время запуска

• Минимальный BUS-фактор

TorchServe требует весьма сложного процесса настройки. Модель нужно конвертировать и куда-то прокинуть. Мы написали скрипты, но даже мне как разработчику потребовалось много времени, чтобы запустить всё из готового кода, который создавали другие. Те, кто придёт в команду, вряд ли смогут что-то сделать с этим монстром, который непонятно как запускается.

Кроме того, TorchServe достаточно прожорлив: он работает с JVM и потребляет много оперативной памяти.

Все эти недостатки делают TorchServe неудобным решением для наших задач.

OpenVINO как решение

Проанализировав предыдущие решения, я подумал, что можно сделать лучше, и решил использовать фреймворк, с которым уже работал ранее, — OpenVINO. Этот фреймворк от Intel предназначен для ускорения инференса на их собственных устройствах.

OpenVINO отлично работает не только на Intel: подойдут практически любые x86-процессоры. Ускорение наблюдается и на AMD, так что прирост производительности не ограничивается родным железом. Однако если речь идёт о GPU, NPU или FPGA от Intel, то здесь, конечно, преимущество у их экосистемы: они оптимизировали этот инструмент прежде всего под себя.

Ещё в OpenVINO есть поддержка всех популярных ML-фреймворков, включая PyTorch, TensorFlow и PaddlePaddle. Также он отлично работает с ONNX. Это значит, что если ваш фреймворк не входит в число самых популярных, но умеет конвертироваться в ONNX, то, скорее всего, вы тоже сможете запустить его с помощью OpenVINO.

Конвертация и ускорение с помощью OpenVINO

Для конвертации возьмём модель RoBERTa для получения эмбеддингов, которую можно найти на Hugging Face.

Для начала установим необходимые зависимости:

pip install --upgrade-strategy eager optimum[openvino,nncf]

И пробуем конвертировать модель:

from optimum.intel import OVModelForFeatureExtraction

model = OVModelForFeatureExtraction.from_pretrained(
    '{путь до неоптимизированной модели}',
    from_transformers=True,
    local_files_only=True,
)
model.save_pretrained("{куда сохранять оптимизированную модель}")

Этот процесс пройдёт без особых трудностей. Смотрим на артефакты:

Отлично! Модель сконвертирована, пробуем её запустить.

Ванильная реализация:

from sentence_transformers import SentenceTransformer


model = SentenceTransformer('ST_model', device='cpu')
outputs = model.encode("hello")
print(outputs)

>>> [0.61127436, 0.76438063, ... , 0.8329032, -0.08197051]

OpenVINO-реализация:

from openvino.runtime import Core

core = Core()
model = core.compile_model(
    model=core.read_model(model='openvino_model.xml'),
    device_name="CPU",
)
print(model("hello"))

>>> EXCEPTION

Сверху показаны способы получения эмбеддингов слова с помощью библиотеки Sentence Transformers. Вызываем метод encode у модели, отправляем «hello» и получаем эмбеддинги для этого слова. Всё просто.

Теперь попробуем сделать то же самое с OpenVINO, и у нас ничего не получится. Дело даже не в том, что мы не вызвали отсутствующий метод encode. Каким бы методом мы ни воспользовались, ошибка всё равно будет. Давайте разбираться, что пошло не так.

Анализ структуры RoBERTa

Для начала посмотрим, что такое RoBERTa и из каких компонентов она состоит. Оказывается, внутри много файлов:

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

• Сначала идёт токенизатор (на схеме выделен красным), занимающий большую часть представленных файлов

• Затем — сама модель RoBERTa (выделена зелёным)

• После этого несколько векторов с модели передаются на постпроцессинг (выделен жёлтым), который преобразует эти данные в эмбеддинги слова — это один длинный вектор

Почему возникла ошибка

Дело в том, что на предыдущем шаге мы конвертировали только саму модель (зелёную часть) и не учли два других этапа: токенизатор и постпроцессинг. Модель не может работать напрямую с текстом — она работает с токенами текста. А для их получения нужен токенизатор. Именно поэтому мы и получили ошибку.

Конвертация компонентов

Будем работать по порядку: начнём с токенизатора.

Я использовал библиотеку Hugging Face Tokenizer. OpenVINO имеет собственную библиотеку OpenVINO Tokenizers, но тестирование показало, что реализация токенайзеров от Hugging Face работает гораздо быстрее.

Запускаем токенизатор. Нужно помнить, что с его помощью мы должны получить две вещи: токены слова и attention mask. Эти данные будем подавать на вход RoBERTa.

Устанавливаем зависимости: pip install tokenizers

Запускаем токенизатор:

from tokenizers import Tokenizer

tokenizer: Tokenizer = Tokenizer.from_file("roberta/tokenizer.json")
tokens = tokenizer.encode("hello")
token_ids = tokens.ids                  # [1, 1892,24147,83,2]
attention_mask = tokens.attention_mask  # [1, 1, 1, 1, 1]

Теперь посмотрим на постпроцессинг. Заходим в соответствующую папку (в нашем случае это папка«1_Pooling») и находим config.json:

{
  "word_embedding_dimension": 1024,
  "pooling_mode_cls_token": false,
  "pooling_mode_mean_tokens": true,
  "pooling_mode_max_tokens": false,
  "pooling_mode_mean_sqrt_len_tokens": false
}

Этот JSON-файл даст полезную информацию о том, что происходит внутри. Например:
Векторы, с которыми работает постпроцессинг, имеют размерность 1024. Используется Mean Pooling: берём несколько векторов, складываем их, делим на общее количество и получаем средний вектор — именно то, что нужно.

Существует два варианта, что можно сделать с постпроцессингом.

1. Написать вручную с помощью Python.

   Преимущества: просто, быстро, понятно.

   Недостатки: долгая работа, сложно.

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

2. Создать маленькую модель, которая будет выполнять необходимые задачи.

   Преимущества: быстрая работа, универсальность.

   Недостатки: неочевидно, нужны дополнительные действия.

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

Поэтому пойдём вторым путём и создадим очень простую модель на PyTorch с одной строчкой логики. Это тот самый Mean Pooling.

import torch


class PostProcess(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, x):
        return torch.sum(x, 1) / x.size()[1].squeeze()

Суммируем все векторы и делим на их количество.

Учитывая, что OpenVINO отлично работает с ONNX, конвертируем модель в этот формат:

torch.onnx.export(
    PostProcess(),
    torch.randn(1, 1, 1024, requires_grad=True),
    'postproc.onnx',
    export_params=True,
    opset_version=11,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {1: 'batch_size', 2: 'sequence_length'},
        'output': {0: 'batch_size'},
    },
)

Конвертируем RoBERTa в ONNX. Для этого сначала установим ONNX runtime:

pip install onnxruntime

from optimum.onnxruntime import ORTModelForFeatureExtraction

model = ORTModelForFeatureExtraction.from_pretrained(
    'roberta',
    local_files_only=True,
    export=True,
)
model.save_pretrained('roberta.onnx')

После всех преобразований получаем pipeline из трёх этапов, как и планировали: токенизатор → конвертированная в ONNX RoBERTa → модель Mean Pooling тоже в ONNX.

Кажется логичным прогонять эти процессы одновременно, так как они являются частями одного pipeline.

Склейка моделей

Здесь на помощь приходит библиотека ONNX, позволяющая склеивать модели:

import onnx

model = onnx.load('roberta.onnx')
postproc = onnx.load('postproc.onnx')
combined = onnx.compose.merge_models(
    model,
    postproc,
    [('last_hidden_state', 'input')],
    prefix1='roberta',
    prefix2='pooling',
)
onnx.save_model(combined, 'final.onnx')

Нам всего лишь нужно знать выходной слой RoBERTa (в нашем случае) и входные слои постпроцессинга, на которые будем отправлять векторы. ONNX позволяет объединить эти две модели, а затем конвертировать в OpenVINO, получая модель с вшитым в неё постпроцессингом.

Конвертируем в OpenVINO:

from openvino.runtime import save_model
import openvino as ov

ov_model = ov.convert_model('final.onnx')
save_model(ov_model, 'final.xml')

Итоговый pipeline конвертации RoBERTa:

Начинаем с исходной модели transformers RoBERTa, которая делится на три части: токенизатор, Torch RoBERTa и Mean Pooling. С токенизатором мы ничего не делаем, используем Hugging Face Tokenizers.

Mean Pooling сначала реализуем как PyTorch-модель, а затем конвертируем её в ONNX.

PyTorch RoBERTa конвертируем в ONNX, после чего склеиваем её с Mean Pooling. Полученную модель конвертируем в OpenVINO, создавая двухступенчатый pipeline: вначале токенизация, а затем модель и вшитый в неё постпроцессинг.

Проверка точности

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

import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
from tokenizers import Tokenizer

from openvino.runtime import Core

st_model = SentenceTransformer('roberta_st', device='cpu')

core = Core()
ov_model = core.compile_model(
    model=core.read_model(model="roberta_ov/model.xml"),
    device_name="CPU",
)
tokenizer: Tokenizer = Tokenizer.from_file("roberta_ov/tokenizer.json")
output_layer = ov_model.output(0)


def encode_ov(text: str) -> list[float]:
    tokens = tokenizer.encode(text)
    vector = ov_model(
        [
            np.expand_dims(np.array(tokens.ids), 0),
            np.expand_dims(np.array(tokens.attention_mask), 0),
        ]
    )[output_layer]
    return np.squeeze(vector).tolist()


test_texts = [
    "some text",
    "one more text",
]

st_embeddings = [st_model.encode(text) for text in test_texts]
ov_embeddings = [encode_ov(text) for text in test_texts]
similarities = [
    cosine_similarity([st_embeddings[i]], [ov_embeddings[i]])[0][0] for i in range(len(test_texts))
]
print(similarities)

>>> [0.9999866672039361, 0.9999881433573915]

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

Метрики модели, когда запустили её в эксплуатацию, также показали, что её точность не ухудшилась, а значит, наша конвертация прошла успешно. Отличная новость!

Инференс

Итак, что надо сделать, чтобы получить эмбеддинги текста с помощью OpenVINO без лишнего кода:

import numpy as np
from tokenizers import Tokenizer
from openvino.runtime import Core


core = Core()  # инициализируем сам OpenVino
ov_model = core.compile_model(  # инициализируем модель
    model=core.read_model(model="roberta_ov/model.xml"),
    device_name="CPU",
)
tokenizer: Tokenizer = Tokenizer.from_file(  # инициализируем токенизатор
    "roberta_ov/tokenizer.json"
)
output_layer = ov_model.output(0)  # получаем адрес последнего(выходного) слоя


def encode_ov(text: str) -> list[float]:
    tokens = tokenizer.encode(text)  # получаем токены текста
    vector = ov_model(
        [
            np.expand_dims(np.array(tokens.ids), 0),  # отправляем токены на входной слой модели
            np.expand_dims(  # отправляем attention_mask на входной слой модели
                np.array(tokens.attention_mask), 0
            ),
        ]
    )[output_layer]  # получаем данные из выходного слоя
    return np.squeeze(vector).tolist()  # конвертируем в питоновский список

Замеры производительности

Измерим производительность нашей RoBERTa. Мы проделали множество сложных операций — что-то склеивали, что-то конвертировали, и теперь следует понять, оправдано ли это.

Сравнение скорости работы

Вот как выглядит график RPS:

Кроме того, я протестировал модель в нашем Dev Cluster. Продовый кластер не трогал (не самая безопасная практика). В Dev Cluster прирост составил 1,5 раза. Возможно, это связано с тем, что процессоры там не самые быстрые, или же есть ещё какие-то ограничения.

Потребление памяти

Посмотрим, как изменилось количество потребляемой памяти:

Улучшение latency

Теперь перейдём к ключевой метрике, на которую мы целились с самого начала — latency.

Вот как это выглядит в Grafana:

Приведу также замеры производительности RoBERTa, где можно сравнить работу на CPU и на GPU:

Мы переместили RoBERTa на GPU и получили скорость ответа около 20 мс, что сопоставимо с нашим худшим результатом на CPU после оптимизации. А в лучшем случае наше решение на CPU (8 мс) оказалось даже быстрее не оптимизированного варианта на GPU.

Однако стоит отметить, что инференс на GPU мы не оптимизировали: использовали обычный PyTorch в стандартном режиме и получили такие результаты. Если применить дополнительные оптимизации, то можно добиться ещё более высокой производительности.

Дополнительные методы ускорения

Если оптимизация модели не дала ожидаемых результатов, поделюсь ещё несколькими способами.

Pruning

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

Обычно модель обучается на большом количестве классов и на крупных датасетах. Но если вам нужно распознавать один-два класса (максимум три), то можно взять open source нейронную сеть, отрезать от неё половину и выиграть в производительности. При этом вряд ли вы потеряете в точности вашего достаточно небольшого распознавания. Сейчас, когда очень распространены большие LLM, именно урезание весов модели позволяет уместить языковые модели на домашние ПК. Параметры модели урезают с нескольких сотен миллиардов до нескольких десятков, а то и до нескольких миллионов. Однако это не проходит бесследно и ведёт к потере точности модели.

Квантование

С помощью квантования сжимают веса нейронной сети. Это самый простой и, наверное, самый распространённый подход.

Простейший вариант квантования: берут полноразмерную обученную модель и с помощью разных утилит конвертируют её веса в меньшую размерность. Благодаря этому сокращаются вычислительная сложность, время инференса, физический размер модели и потребление памяти. Мы заменяем веса нашей модели на веса меньшей точности и меньшего размера, скажем, int или float меньшей размерности. Но этот метод имеет и недостатки: например, уменьшается точность предсказаний модели.

Потери могут варьироваться от модели к модели, но, скорее всего, без этого не обойтись, и надо быть к этому готовым.

Когда нужно квантование

О квантовании стоит подумать, если критически важна скорость работы модели. Метод особенно эффективен в нескольких случаях:

• Вы работаете с потоковыми данными, где борьба идёт за каждую миллисекунду

• Ограничены аппаратные ресурсы и нужна оптимизация без существенной потери качества

• Применяются задачи потоковой обработки аудио и видео, LLM, одноклассовой сегментации и классификации

Не стоит использовать этот метод, если в приоритете точность модели. Работая над задачами классификации или object-detection, где важен каждый процент точности или даже каждая десятая доля процента, квантование может навредить.

Что делать, если CPU всё-таки не хватает

Что делать, если вы уже пробовали всевозможные способы оптимизации (квантование или pruning), но производительности всё равно не хватает, а модель работает медленно или не помещается в память?

В этом случае можно использовать OpenVINO Model Server — это ПО для инференса моделей, оптимизированных с помощью OpenVINO на сервере. Можно установить на компьютер всякие Intel NPU и гонять на них, либо же пойти путём применения GPU. В последнем случае тоже есть отличные инструменты для оптимизации:

• TensorRT — библиотека для оптимизации сетей под Nvidia GPU

• Nvidia Triton Inference Server — отдельный inference-сервер для прогона сетей на Nvidia GPU

В качестве дополнительного бонуса: у Nvidia Triton Inference Server есть OpenVINO backend, что позволяет использовать лучшие качества обоих миров. Вы можете запускать одни сети на CPU, а другие — на GPU. Очень классная вещь, советую с ней ознакомиться.

Итоги

Я рассказал об инструменте OpenVINO, благодаря которому нам удалось ускорить инференс нашего векторизатора RoBERTa в 3 раза. Это позволило обогнать GPU по скорости, хотя инференс на GPU мы никак не оптимизировали.

Параллельно нам удалось сэкономить ресурсы ОЗУ и ПЗУ и ускорить время сборки, потому что образ стал собираться быстрее.

Но самое главное: мы не потратили деньги на новый GPU.

Также мы узнали, из каких частей состоит RoBERTa, зачем оптимизировать каждую из них и как склеивать модели.

Зачем оптимизировать модели

Почему бы просто не запустить модели, как есть, используя PyTorch, не заморачиваясь с оптимизацией? Я имею в виду оптимизацию не только на CPU, но и на GPU.

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

• Машинное обучение находит всё больше применений.

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

• GPU становится дороже. Миграция на GPU — не всегда оптимальный путь. Иногда вместо того, чтобы инвестировать в дорогостоящие GPU, имеет смысл подумать об оптимизации на CPU. Это поможет сэкономить деньги и ресурсы.