From 809451ff4be5a0540718fce9274aa3c88239a935 Mon Sep 17 00:00:00 2001 From: AmitMY Date: Thu, 23 Jul 2026 15:04:14 +0200 Subject: [PATCH] fix(fast): training no longer mutates the global cluster cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #56. apply_merge_to_cluster_cache turned the cluster-graph memo into trained state: the second training run in a process started from mutated cluster graphs and produced different merges than the first (the reference has no such statefulness — its cluster building is a pure function). The mutation served nothing: no path can rely on merged cached graphs, since the reference would fail it too. Deleted, with its call sites and the now-empty on_merge hook. SuperBPE's full-scale digest changes 2a6baff51d -> 30929f0c51: phase 2 previously built graphs from clusters mutated by phase 1's merges. The new value is proven correct against the reference at 20k rows / 500 merges (both b308cdb768). All other digests unchanged; warm == cold at full scale. New shared test: train each tokenizer twice in one process and assert identical merges (all four failed before the fix under fast). Separately discovered while gating: 'cd fast && pytest ../tests/' never loaded fast/conftest.py under pytest 9, so the fast suite silently tested the REFERENCE. Running it properly (-p conftest, now documented in CLAUDE.md) exposed two unmirrored APIs, both added to the fast shim: GraphSettings.TRADE_MEMORY_FOR_SPEED and Tokenizer(cache_maxsize=...) with the word-graph lru_cache dedup in _build_graphs. Co-Authored-By: Claude Fable 5 --- CLAUDE.md | 4 +-- .../graphs/settings.py | 3 ++ .../complex_tokenization_fast/tokenizer.py | 32 ++++++++++++------- fast/src/trainer.rs | 13 ++------ fast/src/units.rs | 9 ------ tests/test_tokenizer_api.py | 17 ++++++++++ 6 files changed, 44 insertions(+), 34 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index dd6d773..91d1889 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -18,11 +18,11 @@ To verify: pip install -e . && pytest tests/ # Test fast (from fast/ directory) -cd fast && maturin develop --release && pytest ../tests/ +cd fast && maturin develop --release && pytest ../tests/ -p conftest ``` ## Testing - Reference tests run from the repo root: `pytest tests/` -- Fast tests run from `fast/`: `pytest ../tests/` (the `fast/conftest.py` aliases `complex_tokenization` → `complex_tokenization_fast`) +- Fast tests run from `fast/`: `pytest ../tests/ -p conftest` — the `-p conftest` is REQUIRED: it loads `fast/conftest.py`, which aliases `complex_tokenization` → `complex_tokenization_fast`. Without it, pytest 9 never loads that conftest (it is outside the tests/ subtree) and the run silently tests the reference implementation instead of fast - Both must pass the same 127+ tests diff --git a/fast/python/complex_tokenization_fast/graphs/settings.py b/fast/python/complex_tokenization_fast/graphs/settings.py index d28ae88..4421c14 100644 --- a/fast/python/complex_tokenization_fast/graphs/settings.py +++ b/fast/python/complex_tokenization_fast/graphs/settings.py @@ -4,6 +4,9 @@ class _GraphSettingsMeta(type): _MAX_MERGE_SIZE = 2 _ONLY_MINIMAL_MERGES = True + # Reference-implementation knob (gates its get_merges memoization); the + # fast implementation has no equivalent cache, so it only round-trips. + TRADE_MEMORY_FOR_SPEED = True @property def MAX_MERGE_SIZE(cls): # noqa: N802 diff --git a/fast/python/complex_tokenization_fast/tokenizer.py b/fast/python/complex_tokenization_fast/tokenizer.py index 3fd2fb1..b1047a6 100644 --- a/fast/python/complex_tokenization_fast/tokenizer.py +++ b/fast/python/complex_tokenization_fast/tokenizer.py @@ -1,4 +1,4 @@ -from functools import reduce +from functools import lru_cache, reduce from complex_tokenization_fast._rs import Node, Trainer from complex_tokenization_fast.graphs.settings import GraphSettings @@ -13,7 +13,8 @@ class Tokenizer: - def __init__(self, units="utf8_clusters", merge_size=2, connected=False, pretokenizer=GPTPretokenizer): + def __init__(self, units="utf8_clusters", merge_size=2, connected=False, pretokenizer=GPTPretokenizer, + cache_maxsize=None): if isinstance(units, str): if units not in UNIT_FUNCTIONS: raise ValueError(f"Unknown units: {units!r}. Choose from {list(UNIT_FUNCTIONS)}") @@ -23,6 +24,7 @@ def __init__(self, units="utf8_clusters", merge_size=2, connected=False, pretoke self.merge_size = merge_size self.connected = connected self.pretokenizer = pretokenizer + self.cache_maxsize = cache_maxsize self.merges = [] @staticmethod @@ -33,8 +35,14 @@ def add_merges(self, merges): self.merges.extend(merges) def _build_graphs(self, texts): + # Same build-local word-graph dedup as the reference: repeated words + # share one graph. cache_maxsize=None is unbounded; 0 disables. + if self.cache_maxsize == 0: + units = self.units + else: + units = lru_cache(maxsize=self.cache_maxsize)(self.units) return tuple( - words(text, connected=self.connected, units=self.units, pretokenizer=self.pretokenizer) + words(text, connected=self.connected, units=units, pretokenizer=self.pretokenizer) for text in texts ) @@ -92,28 +100,28 @@ def get_merges(self): class BPETokenizer(Tokenizer): - def __init__(self, units="utf8_clusters", pretokenizer=GPTPretokenizer): - super().__init__(units=units, merge_size=2, connected=False, pretokenizer=pretokenizer) + def __init__(self, **kwargs): + super().__init__(merge_size=2, connected=False, **kwargs) class BNETokenizer(Tokenizer): - def __init__(self, n=4, units="utf8_clusters", pretokenizer=GPTPretokenizer): - super().__init__(units=units, merge_size=n, connected=False, pretokenizer=pretokenizer) + def __init__(self, n=4, **kwargs): + super().__init__(merge_size=n, connected=False, **kwargs) class BoundlessBPETokenizer(Tokenizer): - def __init__(self, units="utf8_clusters", pretokenizer=GPTPretokenizer): - super().__init__(units=units, merge_size=2, connected=True, pretokenizer=pretokenizer) + def __init__(self, **kwargs): + super().__init__(merge_size=2, connected=True, **kwargs) class SuperBPETokenizer(Tokenizer): - def __init__(self, units="utf8_clusters", disconnected_merges=None, pretokenizer=GPTPretokenizer): - super().__init__(units=units, merge_size=2, connected=False, pretokenizer=pretokenizer) + def __init__(self, disconnected_merges=None, **kwargs): + super().__init__(merge_size=2, connected=False, **kwargs) self._disconnected_merges = disconnected_merges def train(self, texts, num_merges=100, progress=False): disconnected_merges = self._disconnected_merges or num_merges // 2 - phase1 = BPETokenizer(units=self.units, pretokenizer=self.pretokenizer) + phase1 = BPETokenizer(units=self.units, pretokenizer=self.pretokenizer, cache_maxsize=self.cache_maxsize) phase1.train(texts, num_merges=disconnected_merges) self.connected = True self.add_merges(phase1.merges) diff --git a/fast/src/trainer.rs b/fast/src/trainer.rs index ae85c7d..2e04314 100644 --- a/fast/src/trainer.rs +++ b/fast/src/trainer.rs @@ -5,7 +5,7 @@ use rustc_hash::{FxHashMap, FxHashSet}; use std::sync::Arc; use crate::graph::{graphv_to_pyobject, pyobject_to_graphv, GraphV}; -use crate::units::{apply_merge_to_cluster_cache, replace_word_cache, snapshot_word_cache}; +use crate::units::{replace_word_cache, snapshot_word_cache}; use std::collections::HashMap; // Merge score: merging a k-tuple with `count` occurrences removes (k-1)*count nodes. @@ -430,9 +430,6 @@ fn train_unconn( num_merges, verbose, true, - &mut |token, nodes| { - apply_merge_to_cluster_cache(token, nodes); - }, ); for (entry, member_idxs) in entries.iter().zip(&members) { @@ -462,7 +459,6 @@ fn train_single( println!("Merging {:?} count={}", nodes, counts[&nodes]); } *graph = graph.merge(&token, &nodes); - apply_merge_to_cluster_cache(&token, &nodes); pending.push((token, nodes)); } pending @@ -846,7 +842,6 @@ fn build_candidate_index(entries: &[WordEntry]) -> Vec, Fx // Persistent global counts + inverted index with delta updates. Per merge step // only the entries containing the chosen candidate are touched, and `global` is // patched by diffing each entry's old vs new candidates instead of being rebuilt. -// `on_merge` fires once per merge in order (used for cluster-cache side effects). // `use_try` selects try_merge semantics (legacy Unconn path) over merge; with // try_merge a no-op leaves the entry and counts untouched. fn train_entries_delta( @@ -856,7 +851,6 @@ fn train_entries_delta( num_merges: usize, verbose: bool, use_try: bool, - on_merge: &mut dyn FnMut(&GraphV, &[GraphV]), ) -> Vec<(GraphV, Vec)> { let mut pending = Vec::new(); let mut global = build_global_counts(entries); @@ -1043,7 +1037,6 @@ fn train_entries_delta( } } }); - on_merge(&token, &nodes); pending.push((token, nodes)); } pending @@ -1056,7 +1049,7 @@ fn train_streaming_disconnected( num_merges: usize, verbose: bool, ) -> Vec<(GraphV, Vec)> { - train_entries_delta(entries, None, range_start, num_merges, verbose, false, &mut |_, _| {}) + train_entries_delta(entries, None, range_start, num_merges, verbose, false) } // Connected: must assemble doc graphs for cross-word pairs @@ -1406,7 +1399,6 @@ impl Trainer { }; let token = make_token(&nodes); apply_merge_parallel(subs, &mut active, &token, &nodes); - apply_merge_to_cluster_cache(&token, &nodes); pending.push((token, nodes)); let merge_num = i + 1; @@ -1426,7 +1418,6 @@ impl Trainer { }; let token = make_token(&nodes); *graph = graph.merge(&token, &nodes); - apply_merge_to_cluster_cache(&token, &nodes); pending.push((token, nodes)); let merge_num = i + 1; diff --git a/fast/src/units.rs b/fast/src/units.rs index 8e0e4e6..4bd66b1 100644 --- a/fast/src/units.rs +++ b/fast/src/units.rs @@ -52,15 +52,6 @@ pub fn clear_handlers() { CLUSTER_CACHE.lock().unwrap().clear(); } -pub fn apply_merge_to_cluster_cache(token: &GraphV, merge: &[GraphV]) { - let mut cache = CLUSTER_CACHE.lock().unwrap(); - for graph in cache.values_mut() { - if let Some(new_g) = graph.try_merge(token, merge) { - *graph = new_g; - } - } -} - static WORD_CACHE: LazyLock>> = LazyLock::new(|| Mutex::new(HashMap::new())); diff --git a/tests/test_tokenizer_api.py b/tests/test_tokenizer_api.py index 8f8a1ce..e4d4d8f 100644 --- a/tests/test_tokenizer_api.py +++ b/tests/test_tokenizer_api.py @@ -91,3 +91,20 @@ def test_different_pretokenizer_different_merges(self): default_merges = BPETokenizer().train(texts, num_merges=5) whitespace_merges = BPETokenizer(pretokenizer=Whitespace()).train(texts, num_merges=5) assert default_merges != whitespace_merges + + @pytest.mark.parametrize("tokenizer_cls", [ + BPETokenizer, + lambda: BNETokenizer(n=4), + BoundlessBPETokenizer, + SuperBPETokenizer, + ]) + def test_repeated_training_gives_identical_merges(self, tokenizer_cls): + # Training must not observably mutate process-global state: a second + # train in the same process must reproduce the first. The corpus has + # multi-byte characters (é, à) whose UTF-8 byte pairs get merged + # within the first few merges — a cached cluster graph mutated by the + # first train would change what the second train sees. + texts = ["café déjà café déjà café déjà"] + first = tokenizer_cls().train(texts, num_merges=4) + second = tokenizer_cls().train(texts, num_merges=4) + assert first == second