TARGET MINIMAL RUST CRATE + .tok FORMAT - #2293
Open
ArthurZucker wants to merge 97 commits into
Open
Conversation
`tokenizer.json` can only be read by linking `serde_json`, and once that is reachable LTO has to keep the whole JSON stack in every binary that can load a tokenizer. Measured on this workspace: 937 KB gzipped for the encode path with the parser reachable, 354 KB for the same encoder without it. `.tok` v1 is a flat container — 16-byte header, 16-byte section descriptors, sections at 64-byte alignment — read with a bounds check and a pointer cast. It stores only what a `tokenizer.json` stores: vocabulary, merges in rank order, added tokens, and which pre-tokenizer FSM to run. The derived tables (internal ids, merge grid, codepoint fold, the perfect hashes) are rebuilt at load exactly as today. Baking them too would save tens of milliseconds once per process, and would freeze tk-encode's internal layout into a file format. - tk-serialization: container and schema, no dependencies, write behind `write` - tk-encode: `PipelineTokenizer::from_tok`, and `to_tok` behind `tok-write` - tk-convert: the offline `tokenizer.json` -> `.tok` CLI - bindings/node-tok: a Node binding whose whole surface is the read path Byte-exact against the JSON path on gpt2, roberta, llama-3, deepseek-v4 and gemma-3 across 10 scripts (50/50), verified by `tk-convert --example tok_check`. Two fixes fell out: an empty normalizer `Sequence` (deepseek ships one) is no longer carried as a no-op call per segment, and `PipelineToken` is `repr(C)` so a slice of them views as a slice of ids.
`from_tok` routed through `Tokenizer` -> `PipelineTokenizer::try_from` to reuse the added-token id assignment. That was the lazy wiring, and it cost 91 KB gzipped: constructing a `Tokenizer` names `ModelWrapper`, `PreTokenizerWrapper` and `PostProcessorWrapper`, and each of those holds every variant, so all of them link. The reader now fills the pipeline's fields itself. The model is passed to `add_tokens` as a concrete `BPE` and the normalizer as a concrete `Replace`, so neither wrapper is named on the read path; `PipelineNormalizer` gains a `Replace` variant for that, and `read_pre_tokenizer` builds `PipelinePreTokenizer` rather than the config-level enum. The wrapper imports now live inside the `tok-write` module, where they belong. binsize_tok, opt-level=z, stripped, gzipped: 542,128 -> 451,377 Still byte-exact: 50/50 across gpt2, roberta, llama-3, deepseek-v4 and gemma-3.
The 189 KB left over the engine floor was not code, it was 202 KB of static
Unicode tables in `__const`. Nothing constructed the variants that need them —
but a match arm keeps its table alive whether or not the variant is ever built,
and `PipelineNormalizer::Declared`, `PipelineModel`'s three non-BPE arms and
`PipelinePreTokenizer`'s eight classifying arms are all such arms.
So they go behind `config` (default on): the layer that turns a parsed
`tokenizer.json` into a pipeline. With it off the crate loads a `.tok` and
nothing else, and the tables have nothing keeping them alive.
binsize_tok, opt-level=z, stripped, gzipped: 451,377 -> 317,545
the .node: 553,759 -> 331,777
Two bugs the stripped build found, both real on any build without a regex
backend:
- `Split::new` compiles every regex pattern through the system backend.
deepseek's three are not individually FSM-recognised (the pipeline runs
them as one native pass), so constructing them needed an engine that a
read-only build has no reason to carry. `Split::native` skips it.
- `is_deepseek` compares against patterns with literal CR/LF, the way the
shipped config spells them; `atomsplit::regexes` escapes them. The two are
not interchangeable, and rebuilding from the wrong one silently fell off
the native path. `DEEPSEEK_PATTERNS` is now the single source.
Verified by `tok_ids`, which digests every id and runs in both builds: 40/40
model/corpus pairs identical between the full build and the read-only one.
Gating the pipeline enums removed the Unicode tables but left `serde_json` (67
symbols) and `rayon` (243) in a binary that cannot parse JSON and never spawns
a thread. So `serde`, `serde_json`, `rayon` and `rayon-cond` become optional and
join `config`.
The reason `serde_json` survived turned out to be one line:
JsonError(#[from] serde_json::Error)
in the BPE error enum. `Error` is `Box<dyn std::error::Error>`, so that variant
puts a `serde_json::Error` vtable behind every `Result` on the encode path, and
the parser follows it in. One `#[cfg]` and it is gone.
The pass itself is mechanical — `use serde` imports, `derive(Serialize,
Deserialize)`, `#[serde(...)]` field attributes (inert, so they travel with the
derive), hand-written serde impls, and the `serialization` modules. The bulk of
it lives in one place: `impl_serde_type!` generates most of the component types,
so gating the macro covered ~280 of the 438 errors at once.
Three things needed more than a `#[cfg]`:
- `Display for SplitDelimiterBehavior` and `for PrependScheme` delegated to
the serializer to get their names. Spelled out instead, matching what serde
emitted: verbatim for the former (no `rename_all`), snake_case for the
latter.
- `pad_encodings` and `Encoding::pad` are on the ordinary encode path, so they
keep working and just pick a serial iterator when rayon is absent.
- `Model::save`, the `read_file` builders and the batch entry points are
legacy JSON artefacts, so they travel with `config`.
binsize_tok, opt-level=z, stripped, gzipped: 317,443 -> 303,177
the .node: 331,777 -> 317,774 (6.33x)
`serde_json` 0 symbols, `unicode_normalization` 0, `spm_precompiled` 0. `rayon`
remains only because `ptr_hash` depends on it unconditionally for parallel MPHF
construction — upstream, not ours, and worth ~45 KB of text.
Verified both ways: 50/50 byte-exact against the JSON path, and 40/40 pairs
identical between the full build and the serde-free one. 332 tests pass.
`config` leaves the default feature set. The default `tk-encode` is the v1 crate: it reads a `.tok`, encodes, and has no serde in its dependency tree at all. v0 — the legacy `tokenizer.json` reader, the wrapper enums, `Tokenizer` — is what `config` turns on, and `tk-convert` is the only thing that turns it on. `tokenizer/tok.rs` keeps only the read half, which is how a v1 build constructs a pipeline. The writer moves to `tk-convert`, where the conversion belongs: it is now the only crate that names `Tokenizer` or a wrapper enum. cargo tree -p tk-encode -i serde -> only via criterion (dev-dependency) 50/50 byte-exact, 332 tests pass with and without `config`.
The format describes a tokenizer, not one family of them. `Config` gains `model` and `model_param`, there is a `VOCAB_SCORES` section for Unigram, and the four pipeline model arms come back out from behind `config` — a v1 build has to be able to hold whatever a `.tok` names. BPE MERGE_PAIRS + the three MODEL_STRINGS + ignore_merges/byte_fallback/fuse_unk Unigram VOCAB_SCORES, `model_param` = unk id (u32::MAX for none), byte_fallback WordPiece `model_param` = max_input_chars_per_word, unk + continuing prefix from MODEL_STRINGS WordLevel vocab + unk The vocabulary decode is shared: one slab walk into `(token, id)` pairs, plus scores where the model has them. Unigram's vocabulary is positional, so the writer refuses one whose ids are sparse or reordered rather than silently shifting every piece. binsize_tok, opt-level=z, stripped, gzipped: 303,177 -> 334,651 That is what the other three models cost. Per-model features would let a build pay only for what it serves; not done here, because the interesting gate is the normalizer one below. Known gap: a Unigram tokenizer in the wild (albert, xlm-roberta) is blocked by its *normalizer*, not its model — `Sequence[Replace, NFKD, StripAccents, Lowercase, Precompiled]`. Carrying those means linking unicode-normalization and spm_precompiled, i.e. handing back the 154 KB of tables. That wants per-normalizer features rather than one switch, and it is a design call, so it is not in here. 50/50 byte-exact on the BPE families, 332 tests pass with and without `config`.
…on/serde Two things, both about paying only for what a build serves. `normalizers` gates the table-backed normalizers — NFC/NFD/NFKC/NFKD, StripAccents, Bert, and SentencePiece's precompiled charsmap — so `unicode-normalization`, `unicode-normalization-alignments` and `spm_precompiled` leave the dependency tree entirely rather than being compiled and then stripped. Off by default, implied by `config`. `Strip` (lstrip/rstrip) and `Lowercase` need no tables and stay. The last rayon and serde in the tree were never ours: both came from `ptr_hash`, which had them as mandatory dependencies. Upstream PR makes each optional — rayon behind `parallel`, and serde behind `serde`, since the only uses in that crate are two `Serialize` derives on its build statistics. Patched in until it lands: RagnarGrootKoerkamp/PtrHash#32 `cargo tree -p tk-encode -e normal -i {rayon,serde,serde_json}` is now empty, and dropping ptr_hash's parallel construction took it from 273 symbols to 67. binsize_tok, opt-level=z, stripped, gzipped: 334,651 -> 286,130 the .node: 2,010,745 -> 300,542 (6.69x) __text 405,064 -> 380,716 __const 62,896 (atomsplit's classification tables, which are the pre-tokenizer) 50/50 byte-exact, 40/40 ids identical across build configs, 332 tests pass.
Mirrors the ExecuTorch C++ driver it is compared against: best-of-5 per corpus, ids only, one thread, so the two numbers mean the same thing.
Collaborator
Author
vs ExecuTorch's C++ tokenizers
Their stack for the HF-json path: nlohmann/json, RE2 (+ 97 abseil archives), and PCRE2 as a lookahead fallback. cl100k's SizeSmallest program that loads a tokenizer and encodes. Stripped, LTO,
Against the ExecuTorch build that can actually serve llama-3: 1.82x smaller. Against their RE2-only build: 1.40x. ThroughputBest of 5, ids only, one thread, gpt2 — their RE2 path:
llama-3 — their PCRE2 lookahead path:
The llama-3 row wants a caveat: that is their lookahead fallback being pathological, not a fair reading of their BPE engine. The honest summary is 17-41x on the path they are actually optimised for, and that the most popular tokenizer family on the Hub falls off that path entirely because RE2 has no lookahead. This is where Load
Not a headline: we rebuild the derived tables at load on purpose, rather than freeze Reproducing# ExecuTorch
git clone --recurse-submodules https://github.com/pytorch-labs/tokenizers
cmake -S tokenizers -B build -DCMAKE_BUILD_TYPE=Release -DSUPPORT_REGEX_LOOKAHEAD=ON
cmake --build build -j
c++ -std=c++20 -O3 -flto -DNDEBUG -I... et_bench.cpp -Wl,-force_load,build/libregex_lookahead.a $(find build -name '*.a')
# ours
cargo run --release -p tk-convert -- data/gpt2.json
cargo run --release -p tk-encode --example tok_bench -- data/gpt2.tok data/corpora/*.txt
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Based on
perf/bpe-merge-review— review that first, this stacks on it.Target: on-device
A tokenizer that ships inside a phone app, a browser bundle, or an edge runtime is judged on what it adds to the binary. Ours added 2 MB gzipped, and almost none of it was the encoder — it was the JSON parser, the wrapper enums that exist only to be serde's deserialization target, and ~200 KB of Unicode tables kept alive by enum arms nothing ever constructs.
.tokis the fix: a container with no parser. 16-byte header, 16-byte section descriptors, sections at 64-byte alignment, read with a bounds check and a pointer cast. No pretty-printing, no keys, no escaping. v0 (tokenizer.json) stays supported — it moves totk-convert, which runs offline.Before / after
The shipped artifact, gzipped:
tokenizers.darwin-arm64.nodetoday.tok-only binding, v1, all four models6.69x.
Same crate, same flags (
opt-level=z,panic=abort,cgu=1,lto=fat, stripped):binsize_pipeline— v0 JSON load path reachablebinsize_tok— v1,.tokonlyWhere it went, step by step, each measured:
from_tokconfig) — 202 KB of__const.tokv1 gains Unigram / WordPiece / WordLevelnormalizers) + ptr_hash without rayon/serde__constis down from 221,096 to 62,896 — what remains isatomsplit's classification tables, which are the pre-tokenizer.cargo tree -p tk-encode -e normal -i {serde,serde_json,rayon,unicode-normalization,spm_precompiled}is empty. The last two of those were never ours: both came fromptr_hash, where they were mandatory. Upstream PR makes each optional, patched in until it lands — RagnarGrootKoerkamp/PtrHash#32.Crates
tk-serializationwritetk-encode.tok, encodestk-convertbindings/node-tokFeatures, all off by default:
config(the v0 layer —Tokenizer, the wrapper enums, serde),normalizers(the table-backed normalizers).What
.tokv1 carriesAll four models — the format describes a tokenizer, not one family of them.
MERGE_PAIRSin rank order, so a merge's rank is its indexVOCAB_SCORES,model_param= unk idmodel_param=max_input_chars_per_wordPre-tokenizers are named by FSM family, never by regex source, so loading a
.tokneeds no regex engine: GPT-2/cl100k/o200k/tekken/DeepSeek driveatomsplitnatively and a literal pattern is searched for directly.Verification
Converts, reloads through the read-only path, compares every id against the v0 path. 50/50 byte-exact across gpt2, roberta, llama-3, deepseek-v4 and gemma-3 x 10 scripts.
Digests every id and runs in every build config: 40/40 pairs identical between the full build and the v1-only one. 332 existing tests pass with and without
config.File sizes: gemma-3 33.4 -> 9.8 MB, llama-3 17.2 -> 4.9 MB, deepseek-v4 6.4 -> 3.9 MB, gpt2 1.36 -> 1.36 MB (its JSON is mostly merges, already near-minimal as id pairs).
Not in here
.tokcannot yet name a table-backed normalizer, so albert / xlm-roberta / bert convert only once theNORMALIZERsection grows a kind list. The models work; their normalizer chains are the blocker.gemma-2bfails on a pre-existing pipeline limit:<0x09>is absent from itsmodel.vocab, so byte-fallback cannot resolve. Unrelated to the format.Bugs found along the way
Both would bite any build without a regex backend,
.tokor not:Split::newcompiled every regex through the system backend. DeepSeek's three are not individually FSM-recognised (the pipeline runs them as one native pass), so constructing them demanded an engine. AddedSplit::native.is_deepseekcompares against patterns with literal CR/LF, the way the shipped config spells them;atomsplit::regexesescapes them. Not interchangeable — rebuilding from the wrong one silently fell off the native FSM onto the regex fallback.DEEPSEEK_PATTERNSis now the single source.Display for SplitDelimiterBehaviorandfor PrependSchemegot their names by delegating to the serializer, so both broke without serde. Spelled out now, matching what serde emitted.Sequence(deepseek ships one) was carried as a no-op call per segment.