v1.0.0-rc0: the train/encode split, repacked into 53 commits - #2384
Open
ArthurZucker wants to merge 57 commits into
Open
ArthurZucker wants to merge 57 commits into
ArthurZucker wants to merge 57 commits into
Conversation
ArthurZucker
force-pushed
the
tokenizers-rc0
branch
from
September 2, 2026 23:17
5036572 to
a0dddad
Compare
Area: core-crate-split
Area: core-pretokenization
Area: core-encode-pipeline
Part of the pipeline benchmark programme (6,694 lines over 26 commits). The harness was later extracted to the standalone tokbench repo and is not part of the rc0 tree. Squashed from 4 commits: - ci: Simple bench workflow for PipelineTokenizer (#2141) - ci: maintain pipeline-bench report in PR description + comment reactions (#2145) - ci: label-triggered pipeline bench into the PR description (#2146) - ci: run benches on more tokenizers (#2147) Area: ci-bench
Squashed from 7 commits: - refactor: extract pretok algo (#2134) - feat: `impl pipeline::PreTokenizer for FixedLength` (#2137) - feat: `impl pipeline::PreTokenizer for Digits` (#2136) - feat: `impl pipeline::PreTokenizer for CharDelimiterSplit` (#2135) - feat: `impl pipeline::PreTokenizer for UnicodeScripts` (#2138) - feat: `impl pipeline::PreTokenizer for WhitespaceSplit` (#2140) - feat: handle `Merge*` variants of the `SplitDelimiterBehavior` (#2139) Area: core-pretokenization
Area: core-normalization
472 lines. Superseded by atomsplit; utils/multi_regex.rs is not in the rc0 tree. Area: core-pretokenization
Area: core-encode-pipeline
Area: core-pretokenization
Area: core-pretokenization
Part of the pipeline benchmark programme (6,694 lines over 26 commits). The harness was later extracted to the standalone tokbench repo and is not part of the rc0 tree. Area: ci-bench
Part of the pipeline benchmark programme (6,694 lines over 26 commits). The harness was later extracted to the standalone tokbench repo and is not part of the rc0 tree. Area: ci-bench
Part of the pipeline benchmark programme (6,694 lines over 26 commits). The harness was later extracted to the standalone tokbench repo and is not part of the rc0 tree. Area: ci-bench
Part of the pipeline benchmark programme (6,694 lines over 26 commits). The harness was later extracted to the standalone tokbench repo and is not part of the rc0 tree. Squashed from 2 commits: - bench(pipeline): show per-stage split cost as share% (+ ns), not a crushed absolute bar (#2185) - bench(pipeline): multi-thread scaling sweep vs latest release (#2187) Area: ci-bench
Squashed from 2 commits: - feat: bump `ptr_hash` to `2.0.1` - fix: account for `FastPtrHash` no remap sparse layout Area: core-dependencies
Part of the pipeline benchmark programme (6,694 lines over 26 commits). The harness was later extracted to the standalone tokbench repo and is not part of the rc0 tree. Area: ci-bench
The atomsplit split engine: 12,633 lines over 6 commits. Iterated on through the tekken-style regex fast path and literal-pattern offload, then retired in favour of bitsplit -- see `feat(bitsplit): replace atomsplit with the bitsplit engine`. Area: core-splitter
Area: core-memory-allocation
Area: ci-build
Squashed from 5 commits: - post-merge fixes - post-merge format - post-merge format - ci: fix Format and README checks - chore: bump rust edition to 2024 (#2220) Area: core-toolchain
Area: core-vocab-store
Area: core-encode-pipeline
Iteration on the atomsplit engine, later retired in favour of bitsplit. Area: core-splitter
Area: core-parallelism
Iteration on the atomsplit engine, later retired in favour of bitsplit. Area: core-splitter
Area: core-pretokenization
Area: core-memory-allocation
Area: core-memory-allocation
Area: core-model-bpe
Also reinstates the scratch pool after a rebase mishap. Squashed from 2 commits: - revert: rebase mishap, reintroduce scratch pool - feat: gate parallelism (#2295) Area: core-parallelism
Squashed from 4 commits: - fmt / lint - ci: fetch gpt2.json in make test, stop masking Windows lib-test failures - ci: fail on missing fixtures instead of skipping; fetch them on Windows - make: add `data` and `models` convenience targets Area: tooling-tests
Area: core-decode
Area: core-decode
Part of the pipeline benchmark programme (6,694 lines over 26 commits). The harness was later extracted to the standalone tokbench repo and is not part of the rc0 tree. Area: ci-bench
Part of the pipeline benchmark programme (6,694 lines over 26 commits). The harness was later extracted to the standalone tokbench repo and is not part of the rc0 tree. Area: ci-bench
Removes 29,834 lines of the legacy encode engine, including atomsplit's serde layer. Squashed from 2 commits: - tk-encode: strip the legacy encode engine and its serde layer (#2352) - nit Area: core-serialization
Retires atomsplit (12,633 lines built over 6 commits) in favour of the bitsplit engine. atomsplit is removed from the tree here; its design work is recorded in the commits above. Area: core-splitter
Area: core-added-tokens
Area: bindings-node
Squashed from 2 commits: - chore: remove dead code (#2359) - lockfiles Area: core-cleanup
Area: ci-build
ArthurZucker
force-pushed
the
tokenizers-rc0
branch
from
September 2, 2026 23:26
a0dddad to
7c8576a
Compare
ArthurZucker
changed the title
…2378) * perf(pipeline): decide the single-sequence serial fallback before planning `parallel::encode` gated on `inputs.size_bytes() < PARALLEL_MIN_BYTES`, then planned, then bailed to `encode_serial` when planning produced fewer than two tasks. For one document between `PARALLEL_MIN_BYTES` and twice it that bail is guaranteed: `plan_sequence` short-circuits any sequence under `2 * PARALLEL_MIN_BYTES` into a single chunk, one chunk makes one task, and one task fails the check. So the allocations were provably discarded, and the condition is knowable from the length alone. No measurable throughput effect, and it is worth recording why rather than implying one. Because `plan_sequence` short-circuits *before* its added-token scan, the only waste in that window is about six small `Vec` allocations, against ~22 us of encoding for a 9 kB document. Measured with single-document encodes at 4 kB / 9 kB / 20 kB (the first and last as controls, since the guard cannot fire there): every cell within +-1.5%, including the 9 kB window itself. Kept because it removes work that cannot ever be used, not because it is faster. The expensive version of this waste is at 16 kB and above, where `plan_sequence` does run the full `SpecialSegmentIterator` pass over the document, finds no special token, emits one chunk and discards it -- a duplicated added-token scan. This guard does not cover that, and it cannot cheaply: whether a sequence splits is only knowable from that scan. The real fix there is the intra-sequence splitting #2331 left as follow-up, which would make the planning useful instead of merely cheap to skip. * perf(pipeline): order chunks by size class with a counting sort `plan_work` sorted the chunk list with `sort_unstable_by_key(Reverse(len))` and then folded it a second time for the per-sequence `chunk_count`. Two things are wrong with that shape: - `SequenceChunk` is 40 bytes, so a comparison sort does O(n log n) moves of 40 bytes each. The container was never the problem -- `chunks` is a `Vec`, access was already O(1) -- what it sorts is. - It runs in the serial prologue, on the calling thread, before a single worker is spawned. It is pure Amdahl fraction, and worst exactly in the many-small-inputs case tasks were introduced for in #2338: one chunk per sequence means n is the batch size, so an 80k-document batch sorts 80k structs before any encoding starts. The grouping it feeds never needs a total order -- it walks the chunks accumulating bytes and cuts a task every `PARALLEL_MIN_BYTES`. Ordering by power-of-two size class is sufficient, and that is a counting sort: histogram, exclusive prefix sums largest-class-first, scatter. O(n), moving 4-byte indices instead of 40-byte structs, with `chunk_count` coming off the same histogram pass rather than a second fold. `tasks` now index `order` rather than `chunks`, so the worker pays one extra `u32` load per chunk -- against at least `PARALLEL_MIN_BYTES` of tokenizing per task. Batched `encode` of a 5 MB corpus, min of 5 passes with the corpus repeated to fill 250 ms per pass, 3 interleaved rounds, separately built binaries: 64 B documents n_docs base MB/s this commit gpt2/chinese 80457 362.7 +17.8% llama-3/chinese 80457 344.4 +11.0% llama-3/code 32691 709.0 +6.8% gpt2/english 81983 668.4 +4.0% gpt2/code 32691 683.9 +2.9% llama-3/english 81983 663.8 +1.5% 1024 B documents n_docs base MB/s this commit llama-3/chinese 5118 600.9 +2.1% gpt2/chinese 5118 846.5 +1.8% others -1.0% .. +0.9% Positive in every cell where n is large and gone by 1 kB documents, which is the shape an O(n log n) -> O(n) change in a serial prefix should have. The noise floor for this setup is about +-4%, measured from a commit that cannot affect the batch path at all, so the two chinese cells are the solid ones and the sub-3% entries are marginal. Token ids are unaffected: chunks are encoded independently and reassembled by `(seq, idx)`, so visiting order cannot change the result. Verified identical across all four builds in every cell. Only the task boundaries move. * perf(pipeline): partition the chunks rather than counting-sort them Replaces the counting sort from the previous commit with a single partition pass, on review against #2361 which reached the same conclusion independently and expressed it in eight lines. The task grouping never needed the ordering the counting sort provided. Tasks are byte-balanced -- each grown until it holds `PARALLEL_MIN_BYTES` -- so how the chunks below that threshold sit among themselves cannot change a task's size. Exactly one property matters: a chunk big enough to be a task on its own must not be picked last and become the straggler the batch waits on. Partitioning those to the front is O(n) in place, and drops the `Vec<u32>` permutation, the `order` field, the size-class table, and the extra load per chunk in `encode_task`. Same measurement as before -- batched `encode`, 64 B documents, corpus repeated to fill 250 ms per pass, 3 interleaved rounds -- shows the two are equivalent in throughput, so the precision the counting sort bought was not worth its code: counting sort partition gpt2/chinese +17.8% +16.9% llama-3/chinese +11.0% +12.9% llama-3/code +6.8% +6.3% llama-3/english +1.5% +5.4% gpt2/code +2.9% +3.0% gpt2/english +4.0% +2.6% Net effect on this PR: 152 insertions become 88, and the worker's inner loop goes back to indexing `chunks` directly. * Apply suggestion from @McPatate Co-authored-by: Luc Georges <McPatate@users.noreply.github.com> * Apply suggestion from @McPatate Co-authored-by: Luc Georges <McPatate@users.noreply.github.com> * Apply suggestion from @ArthurZucker --------- Co-authored-by: Luc Georges <McPatate@users.noreply.github.com> (cherry picked from commit cbd94b8)
* implement padding * axe legacy pad() * padding can be a property of the tokenizer * read padding from json config * fmt * fix * to_json * allow disabling padding * test attention mask, type_ids, sentence_transformers configs * rm maybe_par_iter * test(convert): oracle the two sentence-transformers configs with truncation and padding all-MiniLM-L6-v2 and all-mpnet-base-v2 ship truncation (max_length 128) and padding (Fixed 128) in tokenizer.json, so they are where the padded attention mask and type ids are compared against the released crate. Their fixtures are pinned at hf-internal-testing/tokenizers-test-data 943303b (#11), one commit past the previous pin, so no json_oracle digest changes. (cherry picked from commit de4c7c1)
* perf(pipeline): reuse sequence A's buffer in post_process Extracted from #2361, which bundled it with the flat batch path; it is an independent change to the general path and reviews better on its own. The `batch_alloc` example the original commit added is left out: it imports `tk_encode::Tokenizer`, which #2352 removed, so it no longer compiles. #2361 deleted it two commits later for the same reason. * lol * remove bloat * Apply suggestion from @ArthurZucker * move post processor code where it belongs * more cleanup * Apply batched suggestions from code review Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * remove stuff that belongs in convert * update * nit * cleanup again * fix * fromat * nit * update * remove bloat shit * nit * Apply batched suggestions from code review Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * weave is ai * simple nit * fix(deps): bump chacha20 to 0.10.2, the yanked 0.10.1 fails cargo audit chacha20 0.10.1 and 0.10.0 are both yanked from crates.io. The audit jobs run `cargo audit -D warnings`, and a yanked crate is a denied warning, so the node audit failed. The `--ignore RUSTSEC-*` flags cannot suppress it: "yanked" is an index state, not an advisory ID. Pulled in transitively by rand 0.10.2. tokenizers/Cargo.lock was already on 0.10.2; the node and python lockfiles were missed. Verified with the exact CI command against all three lockfiles: node exit 0 python exit 0 (was 1: "1 denied warning found") tokenizers exit 0 (cherry picked from commit 0743ac0)
* perf(pipeline): give each thread its own scratch sub-pool
`ScratchPool` is one `Mutex<Vec<EncodeScratch>>` and every `encode` locks it
twice: once to pop a scratch, once to push it back when `ScratchGuard` drops.
Every encoding thread therefore serialises on one lock, which is what the note
in the source predicted:
TODO @McPatate : The Mutex can create contention, to be replaced by a
better access pattern
Measured with N threads encoding a 12 kB prompt of 30k distinct words through
one shared `PipelineTokenizer`, throughput stops improving between 16 and 32
threads and then falls back (encodes/s):
threads Intel GNR 128c NVIDIA Vera 88c AMD Turin 96c
1 52,474 67,399 56,076
8 205,768 308,619 408,580
16 298,086 327,975 640,830
32 226,496 609,208 675,155
64 191,749 475,348 668,325
full 223,578 314,715 661,471
Scaling efficiency at full width is 3.3%, 5.3% and 12.3% of linear.
This gives each thread one of 64 sub-pools, picked once per thread from a
counter and kept in a thread-local. A scratch returns to the sub-pool it came
from, so a thread reuses its own buffers and the `WordCache` inside each scratch
stays warm per thread instead of being passed around. A thread whose own
sub-pool is empty takes from another before building -- `try_lock`, so it skips
a busy sub-pool rather than queueing behind it, which would reintroduce what the
sharding removes. That keeps the property the pool exists for: it hands out what
it has rather than allocating alongside it. `keeps_at_most_one_scratch_per_
concurrent_holder` is what makes this non-optional; without the borrow it fails.
Same machines, same inputs, with the sharded pool:
threads Intel GNR 128c NVIDIA Vera 88c AMD Turin 96c
1 52,557 67,927 56,054
8 401,004 538,663 438,338
16 798,554 1,074,139 871,056
32 1,568,314 2,135,057 1,698,626
64 2,858,405 4,265,264 2,832,353
full 4,623,304 4,904,611 3,719,211
20.7x, 15.6x and 5.6x at full width, at 69%, 82% and 69% efficiency. Single
thread is unchanged on all three (1.00x), so the serial path pays nothing.
Measured with the `scratch_scaling` example added in this series, 5 s per point,
on otherwise idle machines. `cargo test -p tk-encode` goes from 131 to 132
tests passing, the new one being the borrow path above.
Reproducibility differs by arm. Three sweeps on the same idle GNR node read
4.62M, 4.80M and 4.94M encodes/s at 128 threads with the sharded pool, against
223,578, 245,339 and 539,356 with the single mutex: a contended lock is sensitive
to where the scheduler puts the threads, and 128 threads have 512 of them to
choose from there. The full-width multiplier on that machine is therefore 9x to
21x depending on which baseline sample it is divided by. The sharded arm is the
reproducible one.
One caveat on scope: an HTTP serving frontend that encodes one request at a time
per connection was measured before and after this change and is unaffected
(0.88-0.98x). Its 60 kB prompts cost about five of the 12 kB encodes tabulated
above, so at 8,500 requests/s it asks the pool for roughly 43,000 encodes/s,
against the unsharded pool's peak of 609,208 on the same machine -- 7%, short of
where this mutex begins to bind. The win here is for callers driving `encode`
from many threads at once.
* example: measure how encode scales with thread count
N threads calling `encode` on one shared `PipelineTokenizer` for a fixed wall
time, reporting encodes/s, speedup and scaling efficiency per thread count. This
is what the preceding commit's numbers come from, and what makes a claim about
pool contention checkable rather than argued.
In tk-convert because loading a released `tokenizer.json` now needs both
`tk_convert::canonicalize_file` and `tk_serialize::from_json`, and this is the
crate that has both.
Vocabulary size is a knob for a reason. The scratch holds a `WordCache`, so a
prompt cycling through a handful of words is a near-100% cache hit: absolute
throughput is inflated, and so is the share of each encode spent on the pool
lock, which is the quantity being measured. `--vocab 30000` is nearer natural
text; `--vocab 15` reproduces the artefact deliberately.
* bench: move the scaling sweep to tk-serialize
Review feedback: a bench, not an example with its own arg parser. Thread counts
from available_parallelism, prompts from big.txt. Two rows: distinct text and a
repeated line, since a WordCache hit is a real regime, not a void measurement.
Drops the Dynamo notes and the TODO this PR fixes.
---------
Co-authored-by: Arthur <arthur.zucker@gmail.com>
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
(cherry picked from commit 1f1602d)
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
4 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.
Repacked replacement for
feat/train_encode_split. Same content, readable history, credit intact.Important
Do not squash-merge. The 53 commits are the deliverable. Fast-forward (
--ff-only) or--no-ff; both keep them.What we packed
feat/train_encode_splitcarried 97 commits (91 + 6 merges). The early era predated our PR discipline, so it held 8 pure-mechanics commits (nit,fmt / lint,lockfiles,post-merge format×2) and long runs of one-PR-per-tiny-change.This branch is the same work in 53 commits, fully linear (0 merges), under three rules:
Area:trailer (core-pretokenization,core-memory-allocation,core-parallelism,core-model-bpe,core-splitter,core-serialization,ci-bench,bindings-node, …) so each contribution says what it helped.Work that was tried and later replaced — kept visible
Exploration that didn't survive into rc0 is not collapsed or hidden. It keeps its own commits, because the design work is part of the record:
atomsplit— 12,633 lines over 6 commits, the split engine that preceded bitsplit. Its arc reads end to end:feat(atomsplit): the atomsplit split engine(@ArthurZucker, +14,054)perf(atomsplit): tekken-style regex fast path(@SBrandeis)perf(atomsplit): offload literal patterns to atomsplit(@SBrandeis)refactor!: strip the legacy encode engine and its serde layer(@ArthurZucker)feat(bitsplit): replace atomsplit with the bitsplit engine(@ArthurZucker)The pipeline benchmark programme — 6,694 lines over 26 commits (@SBrandeis 5,487 / @ArthurZucker 1,177), now 9 commits rather than 4, each noting that the harness was later extracted to the standalone tokbench repo.
MultiRegex(472 lines, @McPatate) keeps its own commit, noting it was superseded by atomsplit.The pipeline pre-tokenizers themselves shipped and are not marked dropped; only the old
tokenizer/pipeline.rstrait host was refactored away (intotokenizer/pipeline/pre_tokenizer.rs).Credit
Author emails were normalised to the GitHub noreply addresses — 7 commits were previously authored under personal addresses and would not have appeared on the contribution graph at all.
Verification
git diff tokenizers-rc0 feat/train_encode_split→ only the 3 files frommain'sd5827816, which this branch contains and the old branch never merged.bindings/python/src/trainers.rswas checked line-by-line: 0 unexplained lines.cargo check --workspace --all-targets→ clean, 0 errors, 0 warnings.Area:-tagged, 0 commits containing another author's work.Pre-existing issue, not introduced here
cargo check -p tk-train --features parity-aware-bpefails withE0432:trainers/bpe/parity_trainer.rsdoesuse super::{BPE, …}buttrainers::bpedoesn't re-exportBPE. Thetk-traintrees are byte-identical between this branch andfeat/train_encode_split, so this is latent on the old branch today. Worth a separate fix before rc0 ships.