Токенизатор BERT в Hugging Face ускорили в 50 раз, убрав лишнюю работу с ASCII

Токенизатор превращает строку в числа, которые принимает нейросеть. В библиотеке tokenizers от Hugging Face такой этап для BERT-моделей работал неоправданно медленно. Команда профилировала код, нашла лишнюю работу и убрала её. BERT-токенизатор теперь примерно в 50 раз быстрее, чем в версии 0.23. Эта цифра складывается из двух шагов: ускорения в 8,6 раза, которое дала подготовка релиза v1.0.0, и ещё в 5,7 раза от новой правки. Правка пока лежит в пул-реквесте. Параллельно классические BPE-токенизаторы, на которых работает большинство LLM, ускорились до 30 раз.

Почему старый BERT вообще важен

BERT вышел в 2018 году, но его потомки до сих пор везде: эмбеддинги, поиск, RAG, реранкеры, классификаторы. Три самые скачиваемые модели на Hugging Face Hub используют нормализатор BERT:

• all-MiniLM-L6-v2 — 218 млн скачиваний в месяц; • кросс-энкодер ms-marco-MiniLM-L6-v2 — 79 млн; • bge-small-en-v1.5 — 59 млн.

Сам bert-base-uncased на шестом месте с 34 млн. Среди 10 000 самых популярных моделей на BPE приходится 842 млн скачиваний в месяц, на WordPiece (алгоритм BERT) — 627 млн. Больше половины BPE-скачиваний идёт на токенизаторы с нормализацией Unicode, чаще всего NFC, как в Qwen или ModernBERT. Эти улучшения ускоряют и её.

Где терялось время

Нормализация приводит разные записи одного и того же символа к единому виду. Например, «é» можно записать одной кодовой точкой или буквой «e» с отдельным знаком ударения. Без нормализации модель получит для них разные токены. У BERT без учёта регистра пять шагов: очистка управляющих и пробельных символов, пробелы вокруг китайских иероглифов, разложение по NFD, удаление диакритики и перевод в нижний регистр.

Флеймграф (график, где ширина блока показывает долю времени функции) дал однозначную картину. Нормализатор съедал 88% времени, а сам WordPiece, который выбирает токены, только 8%. Около двух третей всего времени уходило на функцию binary_search_by, которая проверяла, не является ли символ комбинируемым знаком. Это двоичный поиск по массиву из 1567 символов (6 КБ), примерно 11 сравнений с непредсказуемыми для процессора ветвлениями. Поиск выполнялся для каждого символа, включая обычные латинские «s», «l», «o» и «w».

Три ленивых трюка

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

• Пропуск ASCII через SIMD. В UTF-8 байт принадлежит ASCII, если его старший бит равен нулю. Процессор загружает 16 байт, одной инструкцией собирает старшие биты в число и сравнивает его с нулём.

Если ноль, весь блок можно не декодировать, не искать в таблицах и даже не копировать. Если нет, подсчёт нулевых бит находит первый нелатинский байт. Перевод заглавных ASCII-букв в строчные тоже делается пачкой: операция ИЛИ с 0x20.

• Кеш по блокам кодовых точек. Кириллица занимает 256 подряд идущих кодовых точек (U+0400–U+04FF). При первой встрече символа из блока алгоритм сразу считает ответ «может ли нормализация его затронуть» для всех 256 позиций.

Дальше хватает одного чтения из памяти. Блок занимает четыре строки кеша по 64 байта и остаётся в самом быстром кеше рядом с ядром. • **Ленивое копирование.

** Если текст не изменился, возвращается исходная строка без выделения памяти. При первой правке память выделяется один раз, а нетронутые участки копируются крупными кусками.

Что получилось и где не помогло

На английском тексте скорость выросла с 23 до 151 МБ/с относительно ветки main, а доля нормализатора упала с 88% до 13%. Токены совпадают с прежними на всех тестовых корпусах. Чем больше в тексте ASCII, тем сильнее выигрыш: английский и исходный код ускорились больше всего. В русском, греческом и арабском ASCII около 11%, но кеш блоков дал ускорение в 2,2–3,2 раза. Для ModernBERT с NFC английский почти не изменился (1,06 раза), потому что ASCII там и раньше пропускали, а код, арабский и русский стали примерно вдвое быстрее.

У подхода есть слабые места:

• Китайский текст не ускорился: ASCII в нём мало, а вокруг каждого иероглифа BERT всё равно ставит пробелы. • Строки до 32 байт могут обрабатываться медленнее, потому что накладные расходы выше выигрыша. В худшем случае (StripAccents на 8 байтах китайского) скорость падает до 0,35 от прежней. • Для нормализатора отдельно замедление выше погрешности заметили в 55 из 320 случаев, 19 из них на входах до 32 байт. • BERT с учётом регистра обрабатывает китайский на 8–15% медленнее при входе от 64 байт.

Для инженеров вывод простой. Если вы гоняете миллионы документов через MiniLM или bge для поиска и RAG, токенизация на CPU может быть заметной статьёй затрат, и после релиза v1.0.0 она сильно подешевеет. Тем, у кого в потоке много коротких или китайских строк, стоит замерить скорость на своих данных.