tokenizers v1: кодирование, декодирование и масштабирование — измерены

Обновить на GitHub

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

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

Именно поэтому в предстоящей версии 1 tokenizers мы уделили производительности особое внимание. Токенизация должна быть лёгкой и масштабироваться вместе с вашим рабочим процессом. Ваши GPU никогда не должны простаивать в ожидании, пока CPU завершит токенизацию.

---

В этой статье мы рассмотрим, почему v1 быстрее v0.23 — часто в десятки раз.

Эта работа стала полностью возможной благодаря остальной экосистеме. Токенизация — очень активная область разработки с открытым исходным кодом, и такие библиотеки, как gigatoken, tiktoken, kitoken, tokie, fastokens, wordchipper и ai-tokenizer, а также многие другие, каждая по-своему расширили представление о том, каким может быть быстрый токенизатор. Мы изучили эту работу, и несколько приведённых ниже идей пришли к нам благодаря другим проектам, показавшим, что их стоит попробовать.

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

Мы также благодарим IBM, NVIDIA и команду ExecuTorch за предоставленные исправления и помощь в тестировании на широком спектре оборудования, что позволило расширить поддержку платформ.

Результаты

Мы представляем результаты для кандидата на выпуск tokenizers v1 в сравнении с другими широко используемыми альтернативами. Мы рассматриваем однопоточную и многопоточную работу, масштабирование по числу потоков, сравнение моделей, сравнение языков, задержку, пропускную способность декодирования, кучу памяти, а также размер crate.

Мы запускаем измерения из репозитория tokbench и добавили команду для повторного запуска бенчмарков на вашем оборудовании.

Что такое V1

v1 будет выдавать те же идентификаторы токенов, что и v0.23. Цель заключалась в сохранении выходных данных, API, словаря и рангов слияния, а также в улучшении всего, что можно улучшить. Это касается и широты поддержки. Библиотека остаётся универсальной для разных семейств токенизаторов, а не специализируется на BPE, поэтому v1 загружает всё, что загружал v0.23.

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

Именно на этапе модели выполняется большая часть описанной здесь работы. Восемь из десяти семейств моделей, измеренных в этой статье, используют кодирование парами байтов, или BPE. BPE начинает с байтов предварительного токена и последовательно объединяет соседнюю пару с наивысшим рангом, пока не останется ни одной пары с рангом. Ранжирование изучается при обучении токенизатора и поставляется вместе с ним, поэтому один и тот же текст всегда выдаёт одни и те же идентификаторы. Слияние никогда не пересекает границу предварительного токена. Два других семейства используют WordPiece и Unigram — два других типа моделей, поддерживаемых библиотекой.

На странице конвейера токенизации описаны четыре этапа. На странице алгоритмов токенизации описаны BPE, WordPiece и Unigram.

Был переработан каждый этап. Вот изменения, которые имели значение:

Разделение: битовые потоки вместо regex

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

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

Затем написанная вручную функция может использовать SIMD-инструкции (одна инструкция — множество данных) современного CPU, которые применяют одну операцию сразу к множеству байтов и хорошо подходят для текста UTF-8. bitcannon рассматривает байты входных данных как параллельные потоки битов, поэтому границы определяются булевыми операциями над целыми регистрами, а не сканированием по одному символу. За одну операцию с регистром он обрабатывает 64 байта. Та же идея используется в Parabix для обработки текста и в simdjson для JSON.

Для этого необходимо распознавать шаблон. Большинство моделей BPE на уровне байтов покрывается несколькими грамматиками, а токенизатор, чей шаблон не входит в их число, сохраняет путь через regex и не получает никакого ускорения. Поэтому приведённые выше приросты настолько различаются.

Кэш слов

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

Естественно, по мере роста входных данных число уникальных слов может расти медленнее общего числа слов. Тогда повторяющиеся слова составляют всё большую долю входных данных. Новые слова по-прежнему появляются, что объясняет отдельные промахи на приведённой ниже анимации.

Воспроизвести результат для общих префиксов можно с помощью:

tokbench measure prefix-sharing \
  --engine pipeline \
  --engine hf-tokenizers \
  --compare-to pipeline-no-cache \
  --corpus agentic_swe

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

Цикл слияния

Следующий крупный источник затрат — цикл слияния BPE. Для каждого предварительного токена цикл снова и снова находит соседнюю пару с наивысшим приоритетом и объединяет её. Предыдущая реализация выделяла новую память при каждом вызове и создавала новую очередь с приоритетами для каждого предварительного токена.

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

Каждая кандидатная пара также упаковывается в одно 64-битное значение, причём ранг слияния находится в старших битах. Сравнение двух кандидатов превращается в простое сравнение двух целых чисел, а значение «здесь нет слияния» является максимально возможным, поэтому цикл находит следующее слияние без ветвления.

Методика

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

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

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

Итог

Для десяти семейств моделей, которые охватывает путь encode в v1, он кодирует текст в 3–30 раз быстрее, чем v0.23, в однопоточном режиме на Apple M4 Max. Нижняя граница приходится на t5-base, верхняя — на gpt2. При работе восьми исполнителей масштабирование достигает 76% линейного. На протяжении всех этих изменений v1 выдаёт в точности те же идентификаторы токенов, что и выпущенная библиотека.

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

Следующий приоритет — поддержка большего числа семейств моделей. До выхода 1.0.0 мы перенесём дополнительные модели на новый цикл слияния. После стабилизации кандидатов на выпуск следующим шагом станет внедрение этих улучшений в библиотеку transformers и остальные компоненты экосистемы, зависящие от библиотеки tokenizers.

Эта публикация создана на основе результатов tokbench и будет обновляться по мере расширения поддержки.

Установка

Кандидат на выпуск v1 доступен на crates.io. Вызываемый API остаётся тем же, поэтому меняется только устанавливаемая сборка.

Обычная установка:

cargo add tokenizers --pre

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

cargo add tokenizers --pre --no-default-features --features http

Кодирование не изменилось: тот же вызов, те же идентификаторы.

use tokenizers::tokenizer::{Result, Tokenizer};

fn main() -> Result<()> {
    let tokenizer = Tokenizer::from_pretrained("deepseek-ai/DeepSeek-V4-Flash", None)?;

let encoding = tokenizer.encode("The tokenizer is no longer the bottleneck.", false)?; println!("{:?}", encoding.get_ids()); 

Ok(()) } ```

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

```rust
let encodings = tokenizer.encode_batch(documents, false)?;

Каждый показатель в этой публикации измерялся относительно этого crate. Привязки Python используют тот же код и собираются из bindings/python, но добавляют накладные расходы на каждый вызов, которые не входят ни в одно из этих измерений.

Прогресс в направлении V1

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

Кандидат на выпуск: реализовано

Эта работа доступна в предварительном выпуске Rust на crates.io:

cargo add tokenizers --pre

• разделение workspace: разделение единого crate на tk-encode, tk-serialize, tk-convert и tk-train, чтобы приложение подключало только используемые компоненты
• bitcannon: замена regex-разделения на пути кодирования операциями с битовыми потоками для GPT-2, cl100k, o200k, Tekken и DeepSeek. Это заменило конечные автоматы, поставлявшиеся в первой версии #2201 #2317
• WordCache: повторное использование идентификаторов токенов ранее обработанных предварительных токенов #2262, af5a3e3
• более быстрые структуры поиска и слияния: добавлены FlatCache, MPHF RankStore, инкрементальное слияние и BucketVocabStore #2190 #2188
• повторно используемая память модели: временное состояние модели перенесено в scratch-буферы, чтобы токенизация не выделяла память при каждом вызове #2175 #2183
• постобработка конвейера: постобработка предоставляется как этап конвейера STAGE_POST #2182
• пакетные вызовы модели: обработка нескольких диапазонов предварительных токенов одним вызовом #2304
• более быстрое декодирование: запись декодированных байтов непосредственно в повторно используемый буфер, исключение промежуточных строк и копирований, ускорение поиска токенов, поддержка буферизованной потоковой передачи и параллельное декодирование пакетов
• поддержка role_to_token #2343
• привязки Node.js #2281

1.0.0

• единая реализация кодирования: использование tk-encode при проверке обучения, чтобы обучение и инференс не могли выдавать разные результаты токенизации
• необязательные смещения и маски: вычисление этих метаданных только по запросу, чтобы не загружать путь, использующий только идентификаторы токенов
• переработка нормализаторов
• поддержка bitnorm на основе atomnorm #2209
• предкомпилированный spm
• более простые привязки Python: сокращение блокировок, типов-обёрток и написанного вручную кода диспетчеризации при сохранении наследования, сериализации, пользовательских декодеров, поведения мутаций и поддержки CPython без глобальной блокировки интерпретатора
• привязки C и C++ только для инференса для ExecuTorch и llama.cpp, за которыми, возможно, последуют привязки JVM, Swift и Go

После 1.0.0

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

Другие статьи в нашем блоге