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

Что нового в версии 1

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

Результаты

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

Что такое v1

Версия 1 будет производить те же идентификаторы токенов, что и версия 0.23. Цель заключалась в том, чтобы сохранить выходные данные, API, словарь и ранги слияния, а также улучшить все, что можно улучшить. Это включает в себя широту. Библиотека остается общей для семей токенизаторов, а не специализируется на BPE, поэтому v1 загружает все, что загружала версия 0.23. Токенизатор преобразует текст в список целых чисел, которые читает модель. Токенизаторы выполняют это преобразование в четыре этапа: 1. Нормализация — применяются операции, такие как приведение к нижнему регистру или нормализация Unicode к исходному тексту. 2. Предтокенизация — текст разбивается на более мелкие части, называемые предтокенами. 3. Модель — каждая предтокен преобразуется в токены и сопоставляется с идентификаторами в словаре. 4. Постобработка — добавляются любые специальные токены, которые ожидает модель. На этапе модели происходит большая часть работы, описанной здесь. Восемь из десяти измеренных семейств моделей в этой статье используют кодировку пар байтов (BPE). BPE начинается с байтов предтокена и многократно объединяет пару с наивысшим рейтингом, пока не останется ни одной пары. Рейтинг изучается во время обучения токенизатора и поставляется с ним, так что один и тот же текст всегда производит одни и те же идентификаторы. Объединение никогда не пересекает границу предтокена.

Основные изменения

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

Разделение рабочего пространства: один пакет стал рабочим пространством.

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

Bitcannon: шаблон разделения становится логическими операциями над битовыми потоками, используя SIMD-инструкции для поиска разделений вместо регулярного выражения.

Переписывание цикла слияния: объединяемые элементы формируют инвазивный двусвязный список внутри одного предварительно выделенного буфера.

Кэширование слов: локальный кэш для потока, который сопоставляет байты предтокена с готовыми идентификаторами.

Нативный параллелизм: один общий токенизатор кодирует из многих потоков одновременно.

Разделение: битовые потоки вместо регулярных выражений

Модели BPE используют регулярное выражение для разделения входного текста на более мелкие, более удобные для обработки части, называемые предтокенами. Объединения происходят внутри предтокена и никогда не пересекают границу между двумя из них, поэтому это разделение определяет, что видит остальная часть конвейера. Регулярное выражение является фиксированным параметром модели и не меняется во время выполнения, поэтому нет необходимости в универсальном движке регулярных выражений для его интерпретации при каждом кодировании. Эквивалентную функцию разделения можно написать вручную, один раз, для шаблона, который фактически использует данная модель. Эта функция может использовать SIMD-инструкции современного центрального процессора, которые применяют одну операцию к множеству байтов одновременно и хорошо подходят для текста в кодировке UTF-8. Bitcannon рассматривает байты входных данных как параллельные потоки битов, так что границы появляются в результате логических операций по всему регистру вместо сканирования, которое продвигается по одному символу за раз.

Кэширование слов

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

Цикл слияния

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

Итог

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

Получение версии 1

Кандидат на выпуск версии 1 доступен на crates.io. API, который вы вызываете, остается прежним, поэтому единственное, что меняется, это какая сборка устанавливается. Установка обычная: cargo add tokenizers --pre. Обучение осуществляется с помощью функции по умолчанию, которая включает зависимость C++. Если вам нужно только кодировать, отключите ее, чтобы исключить реализацию обучения: cargo add tokenizers --pre --no-default-features --features http. Кодирование не изменилось: тот же вызов, те же идентификаторы.