perf(fast): batch ingestion in Rust with gigatoken's O200k pretokenizer#55
perf(fast): batch ingestion in Rust with gigatoken's O200k pretokenizer#55AmitMY wants to merge 2 commits into
Conversation
trainer_from_texts (fast/src/ingest.rs) takes the raw text list once and does pretokenization, unique-word dedup, graph building, doc assembly, and Trainer construction entirely in Rust with the GIL released. The Python shim gates it on the default configuration (utf8_clusters units, default GPT pretokenizer, no registered script handlers); anything else falls back to the unchanged per-document path. Pretokenization uses gigatoken's O200k implementation (our GPT_PATTERN is byte-for-byte the gpt-oss/O200k pattern), vendored under fast/src/pretok/ with MIT attribution pinned to rev 542367a3: a cargo git dependency is impossible on stable (gigatoken's manifest uses a nightly-only profile rustflags key), so the O200k subset (~3.3k lines) is vendored with the trimmed pieces documented. Only new crate dep: icu_properties. Pretokenize ladder at 20k rows (1.19M tokens, identical counts everywhere): Python HF pre_tokenize_str 576ms -> HF tokenizers crate parallel 72ms -> gigatoken O200k parallel 5ms. 20k rows / 100 merges, best of 3, digests byte-identical to the reference: - BPE: 1.73s -> 0.36s (4.8x) - BNE n=4: 2.13s -> 0.69s (3.1x) - Boundless: 3.36s -> 2.20s (1.5x; reference 68.2s -> 31x) - SuperBPE: 4.91s -> 2.07s (2.4x; reference 42.9s -> 21x) Setup collapsed 0.95s -> 0.04s, and loop times also dropped because ingestion Arc-shares each unique word graph across occurrences (1.19M occurrence trees -> 74k unique + Arc clones). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
On hold by maintainer decision — do not merge. The vendored |
…red word deltas Profiled on the entire wikitext-2 train split (23.8k docs, 500 merges). Five kept changes, each digest-verified in isolation: - Inverted-index values Vec<u32> -> FxHashSet<u32>: removal was a linear scan of up to ~23k doc ids per affected doc per step. - The per-step registry scan emits only words that actually merge; staying in the registry is the known-unchanged signal (kills a serial ~74k-entry map rebuild per step). - apply_premerge rewrite: changed-only map, scan-first without cloning untouched docs, in-place slot writes via Arc::make_mut (SuperBPE replay 4.65s -> 1.58s). - try_apply_incremental writes (position, replacement) in place instead of rebuilding each affected doc's children Vec (~2.4M Arc clones/step early). - Each changed word's candidate delta list is computed once per step and docs combine the pre-built lists weighted by occurrence count (delta emission 17.1s -> 8.3s thread-summed). Full scale, cold, best of 3: Boundless 10.90s -> 5.64s (1.9x), SuperBPE 8.36s -> 5.05s (1.7x), BPE/BNE flat (bottlenecks untouched). Digests unchanged at full scale and at the 2000-row default config. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Second commit on this branch: the full-scale profiling round (entire wikitext-2 train split, 500 merges). Five changes, each digest-verified in isolation — details in the commit message.
Cumulative vs the original full-scale baselines: Boundless 179s → 5.6s (32×), SuperBPE 216s → 5.1s (43×), and the reference comparison at this scale is now roughly 120× (Boundless).
Ranked remaining (evidence in the profiling report): incremental handling of cross-word top-match steps (~0.5–1s Boundless), a pair→words index for the registry scan (~1s Boundless, ~0.4s SuperBPE), u32 token interning (full-round redesign; only if sub-4s Boundless matters). BNE/BPE are spread thin — no single lever left. |
|
Status update: #65 now delivers this PR's ingestion structure using the HF |
1 participant
The last structural piece of the gigatoken playbook: one boundary crossing.
trainer_from_texts(newfast/src/ingest.rs) takes the raw text list and does pretokenization → unique-word dedup → graph building → doc assembly →Trainerconstruction entirely in Rust, GIL released. Previously text crossed Python↔Rust three times (pretokenizer round-trip, per-word units calls, PyObject→GraphV conversion).The Python shim gates the fast path on exactly the default configuration (
utf8_clusters, default GPT pretokenizer, no registered script handlers); custom pretokenizers/units/scripts take the unchanged fallback (covered bytest_custom_pretokenizer_*,test_callable_units). Pre-existing merge replay (SuperBPE phase 2,add_merges) composes through the same premerge-registry path — digests confirm.A cargo git dependency on gigatoken is impossible on stable: its manifest carries a nightly-only
[profile.profiling] rustflagskey that stable cargo refuses to parse even for dependencies, and the O200k scheme exists in no rev without it. So the O200k subset (~3.3k lines:o200k.rs,o200k_family.rs,mask.rs,unicode.rs) is vendored underfast/src/pretok/, every file headed with the repo URL, pinned rev542367a, and MIT attribution;mod.rsdocuments what was trimmed. Only new crate dep:icu_properties(slimmed from theicuumbrella). Build time unchanged,.so+0.17 MB. Follow-up: a small upstream PR making gigatoken's manifest stable-consumable would let us replace the vendored copy with a real dependency.Our
GPT_PATTERNis byte-for-byte gigatoken's O200k regex (PretokenizerType::from_split_regexmatches it verbatim) — and per the maintainer's call there's no dedicated split-parity test: the digest gates are the check, and all digests match everywhere (20k / 2000 / 500 rows, all four tokenizers, identical token counts).The pretokenize ladder (20k rows, 1.19M tokens)
pre_tokenize_str(old setup path)tokenizerscrate, serialtokenizerscrate, parallelResults (20k rows / 100 merges, best of 3, digests = reference)
Setup collapsed 0.95s → 0.04s; loop times also dropped (BPE 0.78→0.32s) because ingestion Arc-shares each unique word graph across its occurrences (1.19M occurrence trees → 74k unique + Arc clones).
Default benchmark (2,000 rows): fast Boundless 1.40→0.32s, SuperBPE 1.08→0.35s, BPE 0.25→0.10s, BNE 0.47→0.21s, peak RSS down across the board.
Verification
141 passed + 2 skipped from both suite entry points; clippy warning set identical to main (one scoped, commented allow in vendored code); digests byte-identical to the reference at every scale tested.
🤖 Generated with Claude Code