14 мин

169

Когда мы выпустили Auto Embeddings — функцию автоматического преобразования текстов в векторные представления — без развёртывания отдельного сервиса для работы с ML-моделью, — главный запрос пользователей касался скорости работы. Ранее для генерации эмбеддингов использовался только стек SentenceTransformers поверх Candle (Rust-рантайм Hugging Face для ML-инференса), и ресурсы CPU использовались далеко не полностью: в большинстве сценариев нагрузки показатель QPS держался на уровне нескольких десятков документов в секунду независимо от способа подачи данных, а параллельные запросы обрабатывались последовательно в рамках одной сессии модели.

Поэтому мы в течение нескольких недель оптимизировали механизм запуска ONNX-моделей в Manticore. Новый бэкенд ONNX Runtime доступен начиная с Manticore Search 27.1.5 . ONNX (Open Neural Network Exchange) — переносимый формат моделей, в котором уже публикуется большинство популярных open-source моделей для эмбеддингов: MiniLM, BGE, E5 и другие. В результате получилось решение, которое в среднем в 14 раз быстрее прежней реализации SentenceTransformers/Candle на том же оборудовании (обычный недорогой сервер с 16 ядрами / 32 потоками), с той же моделью и теми же весами, если усреднить по всей матрице замеров threads × batch, — и это преимущество сохраняется как при одном клиентском потоке, так и при тридцати двух. Предыдущая реализация во всём диапазоне нагрузок показывала 5–11 документов/с; новая реализация работает в диапазоне 70–230 документов/с.

Этот материал — инженерный отчёт / журнал разработки: что мы пробовали, что нас удивило, от чего отказались и как выглядит итоговая архитектура.

Кратко

• В среднем в 14 раз быстрее прежней реализации SentenceTransformers/Candle по всей матрице замеров threads × batch (1 / 2 / 4 / 8 / 16 / 32 потоков × размеры батча 1…128) на той же машине (16 ядер / 32 потока), с той же моделью и теми же весами.

• Начиная с Manticore Search 27.1.5 ONNX теперь используется по умолчанию как более производительное решение для любой HuggingFace-модели, которая поставляется с файлом .onnx.

• На модели all-MiniLM-L12-v2 предыдущий подход с Candle показывал 5–11 документов/с во всех протестированных конфигурациях. Новая ONNX-реализация даёт 70–230 документов/с — и примерно такое же 14-кратное преимущество сохраняется как при одном клиентском потоке, так и при тридцати двух.

• Среднее время ответа для одиночного INSERT на нашей тестовой машине: примерно 14 мс при одном клиенте, примерно 56 мс под параллельной нагрузкой в 8 потоков — оба значения заметно ниже 200+ мс, которые показывал Candle.

• Нужна максимальная пропускная способность при массовой загрузке? Используйте батч размером 32–128 документов в одном клиентском потоке. Новый бэкенд распараллеливает обработку внутри вызова и эффективнее использует ресурсы CPU, поэтому распараллеливание со стороны клиента лишь увеличивает накладные расходы на координацию — на пике на нашей машине мы зафиксировали 233 документа/с при 1 потоке и batch=64.

• Два изменения, которые дали больше всего: отключить intra_op_spinning и отказаться от батч-обработки документов внутри worker.

• Изменений в пользовательском API нет. Таблица, которая уже указывает на MODEL_NAME с поддержкой ONNX, автоматически подхватит новую реализацию. Переключить существующую таблицу на другую модель не получится одной строкой: Manticore не позволяет менять MODEL_NAME у поля FLOAT_VECTOR на месте, но и пересоздавать всю таблицу не нужно. Можно добавить рядом новый столбец с новой моделью, заново построить для него эмбеддинги и удалить старый.

Почему это важно

С автоэмбеддингами база данных сама запускает модель при каждом INSERT. Это значит, что скорость построения эмбеддингов и есть скорость INSERT: пропускная способность загрузки равна тому, что выдерживает этап эмбеддингов.

Старая реализация SentenceTransformers/Candle не раскрывала возможности CPU. Масштабирование ограничивалось конкуренцией за блокировки, а батч-обработка упиралась в padding: строки в батче дополнялись до длины самой длинной строки; между вызовами среда выполнения приостанавливала потоки так, что следующий вызов не мог эффективно продолжить работу предыдущего. Основная проблема проявлялась просто: top показывал, что сервер был загружен далеко не полностью, что бы мы ни подавали на вход. Пропускная способность — и при одиночных INSERT, и при массовых INSERT по 128 строк, и при одном клиентском потоке, и при тридцати двух клиентских потоках — оставалась на уровне 5–11 документов/с, потому что никакой способ подачи данных не позволял задействовать больше CPU.

В новой ONNX-реализации минимальная пропускная способность выросла примерно на порядок, и у пользователей появились конкретные способы влиять на производительность. Однопоточный однострочный INSERT теперь даёт 72 документа/с — уже примерно в 7 раз выше верхней границы старой реализации. Если увеличить степень параллелизма или размер батча — результат поднимается до 130–230 документов/с, а в пике — 233 документа/с в одном клиентском потоке при --batch-size=64. В среднем по всей матрице замеров threads × batch новая реализация в 14 раз быстрее старой.

Почему ONNX, а не Candle

Библиотека эмбеддингов Manticore уже некоторое время поддерживает несколько бэкендов. Реализация на Candle была хороша по качеству эмбеддингов и простоте интеграции в Manticore. Но для промышленного инференса небольших encoder-моделей вроде семейств MiniLM и BGE с ONNX Runtime трудно конкурировать:

• ONNX Runtime (или ORT — официальный, вручную оптимизированный C++-движок Microsoft для инференса ONNX-моделей) выполняет слияние графа, свёртку констант и автонастройку ядер.

• Для большинства популярных моделей эмбеддингов на HuggingFace в каталоге onnx/ уже есть заранее подготовленный model.onnx. Файл уже находится в формате, который ожидает ORT.

На той же модели all-MiniLM-L12-v2 (с теми же весами), на CPU, ONNX-реализация заметно быстрее Candle. То же качество, но значительно меньше работы на документ.

Сессия ORT настраивается с использованием следующих параметров:

let session = ort::session::Session::builder()?
    .with_optimization_level(GraphOptimizationLevel::Level3)?
    .with_intra_threads(0)?            // let ORT pick (= all cores)
    .with_intra_op_spinning(false)?    // do NOT busy-wait between calls
    .with_flush_to_zero()?             // kill denormals on attention softmax
    .with_approximate_gelu()?          // ~10% faster activation, no quality loss
    .commit_from_file(&onnx_path)?;

Большинство этих параметров выглядят очевидными. Один — нет: intra_op_spinning(false). Мы к нему ещё вернёмся: он дал самый большой выигрыш, и по сути это не столько настройка ORT, сколько решение о профиле нагрузки.

Неочевидная часть: модель параллелизма

Если дать Rust-разработчику задачу «ускорить ONNX» без других ограничений, он обычно выберет один из двух подходов. Мы попробовали оба. Оба оказались неправильными для этой нагрузки.

Подход 1: один общий Session за Mutex (Mutex — это блокировка, которая позволяет только одному потоку обращаться к сессии в каждый момент времени). Просто в реализации, но пропускная способность падает в несколько раз при увеличении параллелизма, потому что все потоки, обращающиеся к сессии, выстраиваются в очередь на блокировке. Это может подойти для CLI-инструмента, но для базы данных, обслуживающей много параллельных вставок, такой вариант не подходит.

Подход 2: пул сессий, по одной Session на CPU. Больше нет конкуренции за блокировку, но время холодного старта умножается, потребление RAM умножается, а быстрые запросы несут накладные расходы на диспетчеризацию только ради того, чтобы попасть на сессию. У нас была рабочая реализация этого подхода в отдельной ветке разработки, но она так и не дала нужного результата.

Удачное решение удалось придумать благодаря одной важной особенности, которую большинство Rust-обёрток для ONNX понимают неправильно: на Linux и macOS C API Run() в ORT является потокобезопасным. Одним Session можно пользоваться из нескольких потоков одновременно без блокировок. C++-сторона уже сериализует то, что нужно сериализовать; Rust API просто прячет это за правилами borrow checker, которые не соответствуют тому, что на самом деле разрешает базовая библиотека.

Поэтому мы оборачиваем сессию в небольшой платформозависимый тип:

#[cfg(not(target_os = "windows"))]
struct SessionWrapper {
    inner: std::cell::UnsafeCell<ort::session::Session>,
}

#[cfg(not(target_os = "windows"))]
unsafe impl Sync for SessionWrapper {}
#[cfg(not(target_os = "windows"))]
unsafe impl Send for SessionWrapper {}

impl SessionWrapper {
    fn with_session<R>(&self, f: impl FnOnce(&mut Session) -> R) -> R {
        f(unsafe { &mut *self.inner.get() })
    }
}

Да, это unsafe. Мы выводим borrow checker из цепочки, потому что базовая библиотека документирована как безопасная при используемом нами паттерне доступа. Это осознанный unsafe с однострочным обоснованием, а не ловушка.

На Windows у потоковой модели ORT есть известные проблемы, поэтому мы реализуем синхронизацию доступа через Mutex, гарантируя последовательное выполнение операций Run(). Важно, что блокировка сохраняется на протяжении всего времени выполнения замыкания, а не только вызова run() — именно это исправило гонку, которую мы видели на Windows, когда SessionOutputs одного потока ещё читались, а другой поток уже запустил новый run(). Блокировка происходит на уровне замыкания, а не отдельного вызова.

Адаптивный параллелизм — неудачные подходы

Эта часть заняла больше всего времени.

Мы токенизировали пачки из 8, 16 или 32 документов за один проход, дополняли последовательности до максимальной длины (max_len) с помощью padding-токенов и запускали по одному проходу инференса на каждый worker-поток. Пропускная способность оказалась ниже, чем при обработке тех же текстов по одному через ту же сессию. Мы провели повторные замеры. Тот же результат. Некоторое время мы пытались его опровергнуть, прежде чем принять. Отменённый коммит 980b24b "Revert: perf(model): batch inference in worker threads" — это момент, когда мы изменили архитектуру на основе данных профилирования.

За неожиданным результатом стояли две причины.

Накладные расходы, связанные с padding. Батч текстов разной длины дополняет каждую строку до самой длинной строки. Вычислительная сложность пропорциональна произведению batch_size × max_len × hidden_dim, независимо от того, сколько реально содержимого есть в батче. Реальные текстовые входы очень различаются по длине: типичный батч из 8 случайных предложений может содержать один документ длиной 60 токенов (выброс) и семь строк по 8 токенов. Модель выполняет значительное число вычислительных операций над padding-токенами и весами attention. При батчах по одному документу модель выполняет работу, пропорциональную фактическому числу токенов этого документа. Обработка отдельных документов без объединения в батчи демонстрирует более высокую производительность, как только разброс длины входов становится реалистичным.

Busy-waiting (активное ожидание). Пул потоков внутриоперационной параллелизации (intra-op) в ORT по умолчанию выполняет активное ожидание (busy-wait) между выдачей задач: потоки нагружают CPU в цикле ожидания, ожидая следующий фрагмент работы. При одном большом батче на вызов сессии это незаметно: поток постоянно занят полезной работой. При множестве параллельных маленьких вызовов это превращается в проблему: каждый такой пул потребляет 100% доступных ресурсов CPU между вызовами, и для всего остального CPU не остаётся. Вывод top показывал следующую картину: все ядра на 100%, пропускная способность ниже, чем при отключённом spinning. На первый взгляд это кажется нелогичным, пока не вспомнишь, что остальной системе тоже нужно CPU-время — модулю токенизации, процессу построения Hierarchical Navigable Small World (HNSW)-индекса, другим компонентам searchd. Установка параметра with_intra_op_spinning(false) в одной строке кода сразу подняла пропускную способность и одновременно снизила загрузку CPU.

Итоговое решение оказалось противоположным типовой рекомендации:

• Одна общая сессия, без пула.

• Один документ на вызов инференса, без батч-обработки внутри worker.

• Вызовы из нескольких потоков одновременно, масштабирование по числу CPU.

• Без spinning между вызовами — эффективно разделяем ресурсы CPU с другими компонентами системы.

fn predict_pipelined(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, _> {
    let bs = batch_size();

    // Request of small volume — single tokenize + infer, no thread overhead.
    // This is the path a 1-doc INSERT takes.
    if texts.len() <= bs {
        return Self::tokenize_and_infer(&self.session, &self.tokenizer, texts, ...);
    }

    // Large input — split across workers, each running 1-doc-at-a-time
    // through the SHARED session. This deliberately mimics the
    // many-concurrent-callers pattern that ORT is happiest with.
    let num_workers = (texts.len() / bs).min(available_cpus()).max(1);
    let docs_per_worker = texts.len().div_ceil(num_workers);

    std::thread::scope(|s| {
        for worker_texts in texts.chunks(docs_per_worker) {
            s.spawn(move || {
                for text in worker_texts {
                    Self::tokenize_and_infer(&session, &tokenizer,
                                             std::slice::from_ref(text), ...)?;
                }
                Ok(())
            });
        }
    });
    // ...
}

Архитектура с двумя ветками обработки реализована целенаправленно. INSERT на 1 строку приходит с texts.len() == 1, что <= bs, поэтому идёт по быстрой ветке без создания потоков, без отправки в каналы и без накладных расходов на координацию. Массовый REPLACE INTO с тысячами строк идёт по параллельной ветке и получает выигрыш в пропускной способности. Операции с малой вычислительной сложностью выполняются быстро, а ресурсоёмкие задачи эффективно распараллеливаются.

Мы также один раз при старте включаем параллельный процесс токенизации (TOKENIZERS_PARALLELISM=true) и заранее обрезаем входы по числу символов перед Byte Pair Encoding (BPE), чтобы текстовый объект размером 100 КБ не занимал CPU на токенизаторе целую секунду ещё до того, как модель его увидит.

Цифры

Все запуски выполнены на нашей стандартной машине для бенчмарков, с all-MiniLM-L12-v2-onnx, по 1000 документов на запуск. Сгенерировано с помощью manticore-load :

manticore-load --quiet --drop --batch-size=1 --threads=8 --total=1000 \
  --init="CREATE TABLE t (
    f text,
    v FLOAT_VECTOR KNN_TYPE='hnsw' HNSW_SIMILARITY='l2'
      MODEL_NAME='onnx-models/all-MiniLM-L12-v2-onnx' FROM=''
  )" \
  --load="INSERT INTO t(f) VALUES('<text/10/100>')"

Та же команда с --batch-size=2, 8, 32, 128, все при 8 потоках:

По сравнению с Candle на тех же 8 потоках, где результат был ровно 10 документов/с при любом размере батча, это даёт от 9 до 15 раз больше документов в секунду в зависимости от выбранного размера батча. Столбец «Среднее время ответа» отражает время выполнения одного полного оператора INSERT, а не одного документа; если разделить на размер батча, стоимость на документ попадает в диапазон 55–90 мс.

Если использовать конфигурацию с одним клиентским потоком (которая показала наилучшие результаты при массовой загрузке данных), показатели растут ещё сильнее: 72 / 76 / 93 / 175 / 233 / 222 документов/с при батчах 1 / 2 / 8 / 32 / 64 / 128. Пик всей матрицы замеров — 233 документа/с при 1 потоке × batch=64, со временем ответа на документ примерно 4.3 мс.

Как загружать данные для максимальной пропускной способности

Если вы загружаете много данных и хотите максимум документов в секунду, оптимальная стратегия следующая: отправляйте большие запросы INSERT ... VALUES (..), (..), ... (батч размером 32–128 документов) из одного клиентского потока, а не много маленьких вставок из множества потоков. Новый бэкенд распараллеливает обработку внутри вызова (см. код predict_pipelined выше), поэтому распараллеливание со стороны клиента лишь увеличивает накладные расходы на координацию поверх того, что ORT уже делает сам, — поэтому 1 поток × batch=64 (233 документа/с) уверенно обгоняет 8 потоков × batch=128 (147 документов/с).

Для сценариев с обработкой одиночных записей (например, веб-запросов, обработчиков очередей или MCP-серверов) достаточно использовать оператор INSERT INTO. Минимальный результат 72 документа/с для одного потока и одной строки уже примерно в 7 раз быстрее старой реализации на Candle, а время ответа достаточно низкое, чтобы этот уровень больше не требовал отдельной оптимизации.

До и после по всей матрице замеров

Чтобы сделать сравнение до/после, мы выполнили замеры для всей матрицы threads × batch против старой реализации Candle/trans на той же машине и тех же весах:

Каждая точка на оси X соответствует комбинации параметров: число потоков бэкенда / размер батча. Левая половина (trans …) — старая реализация Candle: документы/с держатся на уровне 5–11 во всех конфигурациях независимо от числа потоков и размера батча, при этом загрузка CPU достигает 100 %. Правая половина (onnx …) — новая реализация: документы/с на порядок выше по всему прогону. Внутри новой реализации: на маленьких батчах увеличение числа клиентских потоков повышает производительность (1T/batch=1 = 72 → 8T/batch=1 = 143); на больших батчах выигрывает один клиентский поток (1T/batch=64 = 233 — глобальный максимум).

Тот же прогон, но здесь эффективность (документов/с на % CPU) показана рядом с общей пропускной способностью. На стороне Candle (trans) обе линии остаются на низком уровне — загрузка CPU высокая, но производительность (число обрабатываемых документов в секунду) остаётся низкой. На стороне ONNX (onnx) максимальная эффективность достигается при использовании 1–2 потоков и батчах среднего размера, где каждый процент CPU даёт больше всего эмбеддингов, и остаётся заметно выше старой реализации даже при увеличении числа потоков до 32.

Что дальше

В планах развития — реализация следующих задач:

• GPU-реализация. Текущая настройка ONNX работает только на CPU. Параметр usegpu уже проброшен, но пока не подключён к провайдеру выполнения CUDA (CUDA execution provider) в ORT.

• Паритет производительности на Windows. Сейчас на Windows мы обеспечиваем последовательное выполнение операций из-за ошибки в потоковой модели ORT. Когда эта ошибка будет исправлена в исходном проекте, для Windows будет реализовано то же поведение с общей сессией, которое уже есть в Linux/macOS.

• Поддержка других архитектур в ONNX-реализации. Сейчас ONNX используется для encoder-моделей семейства BERT. T5, causal-LM и квантованные модели формата GGUF пока обрабатываются через Candle.

Как попробовать?

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

Если вы ещё не на ONNX-модели — или хотите перейти на более маленькую или более быструю, чтобы максимально использовать новый бэкенд, — учтите, что заменить модель у существующего поля нельзя. Manticore не поддерживает изменение MODEL_NAME у существующего поля FLOAT_VECTOR, поэтому миграция на месте невозможна. Есть два практических варианта на выбор, в зависимости от того, что проще в вашей среде:

Вариант A — экспортировать, отредактировать, импортировать заново. Даже если исходные данные недоступны, можно экспортировать существующую таблицу в SQL-файл с помощью утилиты mysqldump, отредактировать CREATE TABLE в этом дампе так, чтобы MODEL_NAME указывал на нужную ONNX-оптимизированную модель, и загрузить дамп в новую таблицу. Manticore заново построит эмбеддинги для каждой строки через новую реализацию при вставке.

Вариант B — добавить новый столбец рядом, перестроить, удалить старый. Если вы хотите остаться в SQL и не возиться с дампом, есть альтернатива: добавьте в ту же таблицу новый столбец FLOAT_VECTOR, который указывает на ONNX-модель, а затем запустите одноразовое перестроение эмбеддингов этого столбца из исходного текста:

ALTER TABLE t ADD COLUMN v_new FLOAT_VECTOR KNN_TYPE='hnsw'
  HNSW_SIMILARITY='l2'
  MODEL_NAME='Xenova/all-MiniLM-L6-v2'
  FROM='text_field';

ALTER TABLE t REBUILD EMBEDDINGS v_new;
-- once you've cut over reads to v_new, drop the old column
ALTER TABLE t DROP COLUMN v_old;

Точный синтаксис и ограничения см. в разделе документации Rebuilding embeddings .

Для совершенно новых таблиц всё это не важно — просто сразу выбирайте ONNX-оптимизированную модель.

Удобный источник готовых ONNX-моделей для эмбеддингов — коллекция Xenova на Hugging Face : они уже сконвертированы в ONNX и готовы для использования в MODEL_NAME='...'. Отфильтруйте модели по типу задачи feature-extraction (извлечение признаков), чтобы оставить модели для построения эмбеддингов. Несколько подходящих вариантов для старта:

• Xenova/all-MiniLM-L6-v2 — маленькая и быстрая, 384-мерная, отличный вариант по умолчанию.

• Xenova/all-MiniLM-L12-v2 — модель, которую мы тестировали в этом посте, 384-мерная, обеспечивает более высокую точность эмбеддингов.

• Xenova/bge-small-en-v1.5 — демонстрирует высокую точность при поиске по английским текстам, 384-мерная.

• Xenova/multilingual-e5-small — мультиязычная, 384-мерная.

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

📚 Документация по KNN-поиску
💬 Сообщество в Slack — будем рады получить обратную связь о работе нового подхода на ваших данных.