tokenizers v1: кодування, декодування та масштабування — виміряні

Оновити на GitHub

Токенізатор історично не був вузьким місцем у робочих процесах ML. З погляду обчислень токенізація є легкою порівняно з інтенсивним моделюванням, яке відбувається в решті конвеєра. Однак у деяких випадках вона швидко стала ключовою для прискорення (або уповільнення) роботи з машинним навчанням.

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

Саме тому ми вирішили зосередити значні зусилля на продуктивності майбутньої версії 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 створюватиме ті самі ID токенів, що й v0.23. Мета полягала в тому, щоб зберегти вихідні дані, API, словник і ранги злиття та вдосконалити все, що можна вдосконалити. Це також стосується широти підтримки. Бібліотека залишається універсальною для різних сімейств токенізаторів, а не спеціалізується на BPE, тому v1 завантажує все, що завантажувала v0.23.

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

Саме на етапі моделі відбувається більшість описаної тут роботи. Вісім із десяти сімейств моделей, виміряних у цій статті, використовують кодування парами байтів, або BPE. BPE починається з байтів претокена й послідовно об’єднує сусідню пару з найвищим рангом, доки не залишиться жодної пари з рангом. Ранг вивчається під час навчання токенізатора та постачається разом із ним, тому один і той самий текст завжди створює однакові ID. Злиття ніколи не перетинає межу претокена. Два інші сімейства використовують WordPiece та Unigram — два інші типи моделей, які підтримує бібліотека.

На сторінці конвеєра токенізації описано чотири етапи. На сторінці алгоритмів токенізації описано BPE, WordPiece та Unigram.

Було опрацьовано кожен етап. Ось зміни, які мали значення:

Розділення: бітові потоки замість регулярного виразу

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

Цей регулярний вираз є фіксованим параметром моделі. Він постачається разом із токенізатором і ніколи не змінюється під час виконання, тож немає потреби щоразу під час кодування інтерпретувати його за допомогою універсального рушія регулярних виразів. Еквівалентну функцію розділення можна написати вручну один раз для шаблону, який фактично використовує певна модель.

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

Це залежить від розпізнавання шаблону. Кілька граматик охоплюють більшість байтових моделей BPE, а токенізатор, шаблон якого не належить до них, зберігає шлях регулярного виразу й не отримує жодного з цього прискорення. Саме тому наведені вище прирости так сильно відрізняються.

Кеш слів

У реальному тексті багато повторюваних слів. Оскільки BPE завжди створює однакові ID токенів для заданого претокена, v1 може зберегти результат після першої обробки. Потоково-локальний кеш зіставляє байти кожного претокена з його ID токенів, завдяки чому наступні входження можуть пропустити процес злиття.

Зі збільшенням входу кількість унікальних слів, природно, може зростати повільніше, ніж загальна кількість слів. Тоді повторювані слова становлять дедалі більшу частку входу. Нові слова все одно з’являються, що пояснює періодичні промахи в наведеній нижче анімації.

Відтворіть результат із спільним префіксом за допомогою:

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

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

Цикл злиття

Наступні значні витрати пов’язані з циклом злиття BPE. Для кожного претокена цикл постійно знаходить сусідню пару з найвищим пріоритетом і об’єднує її. Попередня реалізація виділяла нову пам’ять для кожного виклику та створювала нову чергу з пріоритетами для кожного претокена.

v1 повторно використовує буфер-чернетку, яким володіє викликач, усуваючи ці повторювані алокації. Символи зберігаються у плоскому масиві, а сусідні символи зв’язуються їхніми позиціями в цьому масиві, що здешевлює оновлення під час злиття. Також пакет претокенів обробляється одним викликом моделі.

Кожна пара-кандидат також упаковується в одне 64-бітне значення, причому ранг злиття міститься у старших бітах. Порівняння двох кандидатів тоді зводиться до порівняння двох цілих чисел, а «тут немає злиття» є найбільшим можливим значенням, тому цикл знаходить наступне злиття без розгалуження.

Методика

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

Повторне кодування одного документа може бути швидшим за кодування потоку різних документів у тій самій збірці. Перший підхід вимірює продуктивність, коли весь документ уже представлений у кеші. Другий вимірює продуктивність на новому вході, дозволяючи раніше побаченим претокенам залишатися в кеші.

Обидві умови іноді описують як «теплі», хоча вони вимірюють різні навантаження. Наші основні результати використовують різні документи, а весь корпус надто великий, щоб поміститися в кеші. У тестах токенізаторів слід зазначати, яке навантаження використовується, оскільки цей вибір може визначати результат.

Підсумок

У десяти сімействах моделей, які охоплює шлях кодування v1, він кодує текст у 3–30 разів швидше за v0.23 в одному потоці на Apple M4 Max. Нижня межа припадає на t5-base, верхня — на gpt2. На восьми робочих потоках масштабування становить 76% від лінійного. Упродовж усіх цих змін v1 створює точно такі самі ID токенів, як і випущена бібліотека.

Загальне покращення є результатом спільної роботи кількох змін: написаного вручну розділювача замість рушія регулярних виразів, кешу, який повертає результат для повторюваного слова без повторного злиття, циклу злиття, що ніколи не звертається до алокатора, і одного виклику моделі для кожного пакета претокенів замість одного виклику для кожного претокена. Кожна зміна зменшує обсяг роботи в іншій точці конвеєра.

Наступним пріоритетом є підтримка більшої кількості сімейств моделей. До 1.0.0 ми перенесемо додаткові моделі на новий цикл злиття. Після стабілізації реліз-кандидатів наступним кроком буде впровадження цих покращень у бібліотеку transformers та решту екосистеми, які залежать від бібліотеки tokenizers.

Цей допис створено на основі результатів tokbench, і його буде оновлено в міру розширення підтримки.

Як отримати

Реліз-кандидат v1 доступний на crates.io. API, який ви викликаєте, залишається тим самим, тож змінюється лише збірка, яку ви встановлюєте.

Звичайне встановлення:

cargo add tokenizers --pre

Навчання працює через увімкнену за замовчуванням функцію, яка також підтягує залежність C++. Якщо вам потрібно лише кодування, вимкніть її, щоб виключити реалізацію навчання:

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

Кодування не змінилося: той самий виклик, ті самі ID.

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

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

1.0.0

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

Після 1.0.0

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

Більше статей у нашому блозі