From 375bd723074168e4f521f502f4c1c9a1a7524a65 Mon Sep 17 00:00:00 2001 From: AmitMY Date: Thu, 23 Jul 2026 16:13:38 +0200 Subject: [PATCH] perf(fast): batch merge replay for pre-existing merges Replaying N pre-existing merges (SuperBPE phase 2, add_merges) applied them one at a time, each with its own doc pass. Trainer.apply_merges replays the whole list in one call: word-level replay composes each original word pointer to its final graph via the WordRegistry (recording each word's exact collapse step), doc-top events simulate exactly the top-level matches sequential replay performs (both children Node-collapsed before step k, bytes match, left-to-right non-overlapping, cascades), and docs are rewritten once at the end. Byte-identical by construction; a direct probe (sequential x250 vs one batch call) differs in 0 of 1,290 docs, and all digests at all scales are unchanged. Non-Unconn graphs and non-pair merges fall back to the sequential path. SuperBPE full scale: 5000 merges 14.44s -> 8.25s (1.75x), 500 merges 6.48s -> 5.96s. Other cases unchanged. Trained-state reuse (to_connected) was built, measured, and rejected: the phase-1 -> phase-2 handoff is lossy by design (merges travel as lossy UTF-8 strings; non-UTF8-aligned tokens round-trip through U+FFFD and are silently dropped from phase-2 state, in both implementations), so reusing true trained state can never be byte-identical. Documented in issue. Co-Authored-By: Claude Fable 5 --- .../complex_tokenization_fast/tokenizer.py | 22 +- fast/src/trainer.rs | 189 +++++++++++++++++- 2 files changed, 194 insertions(+), 17 deletions(-) diff --git a/fast/python/complex_tokenization_fast/tokenizer.py b/fast/python/complex_tokenization_fast/tokenizer.py index b1047a6..d84b2bf 100644 --- a/fast/python/complex_tokenization_fast/tokenizer.py +++ b/fast/python/complex_tokenization_fast/tokenizer.py @@ -53,10 +53,13 @@ def make_trainer(self, texts): graphs = self._build_graphs(texts) trainer = Trainer(graphs=graphs) - for merge_strs in self.merges: - nodes = tuple(Node(value=s.encode("utf-8")) for s in merge_strs) - token = reduce(lambda a, b: a + b, nodes) - trainer.apply_merge(token, nodes) + if self.merges: + merge_list = [] + for merge_strs in self.merges: + nodes = tuple(Node(value=s.encode("utf-8")) for s in merge_strs) + token = reduce(lambda a, b: a + b, nodes) + merge_list.append((token, nodes)) + trainer.apply_merges(merge_list) return trainer @@ -75,10 +78,13 @@ def make_streaming_trainer(self, texts): trainer = Trainer(graph=Node(value=b"")) trainer.set_streaming(doc_words, connected=self.connected) - for merge_strs in self.merges: - nodes = tuple(Node(value=s.encode("utf-8")) for s in merge_strs) - token = reduce(lambda a, b: a + b, nodes) - trainer.apply_merge(token, nodes) + if self.merges: + merge_list = [] + for merge_strs in self.merges: + nodes = tuple(Node(value=s.encode("utf-8")) for s in merge_strs) + token = reduce(lambda a, b: a + b, nodes) + merge_list.append((token, nodes)) + trainer.apply_merges(merge_list) return trainer diff --git a/fast/src/trainer.rs b/fast/src/trainer.rs index c80ecfb..7e2d16a 100644 --- a/fast/src/trainer.rs +++ b/fast/src/trainer.rs @@ -1849,24 +1849,195 @@ impl Trainer { .iter() .map(|item| pyobject_to_graphv(&item)) .collect::>()?; + self.apply_one_merge(token_g, merge_nodes); + Ok(()) + } + + /// Apply a whole pre-merge list in one call. For an Unconn corpus and + /// pair merges this replays at the word level once per merge and applies + /// doc-top-level matches via an exact event simulation, then rewrites + /// each doc once — byte-identical to per-merge application, without the + /// per-merge doc scans. + fn apply_merges(&mut self, py: Python<'_>, merges: &Bound<'_, PyAny>) -> PyResult<()> { + let mut converted: Vec<(GraphV, Vec)> = Vec::new(); + for item in merges.try_iter()? { + let item = item?; + let token_g = pyobject_to_graphv(&item.get_item(0)?)?; + let nodes_obj = item.get_item(1)?; + let mut merge_nodes: Vec = Vec::new(); + for n in nodes_obj.try_iter()? { + merge_nodes.push(pyobject_to_graphv(&n?)?); + } + converted.push((token_g, merge_nodes)); + } + let batchable = matches!(&self.graph, GraphV::Unconn(_)) + && converted + .iter() + .all(|(_, n)| n.len() == 2 && n.iter().all(|g| matches!(g, GraphV::Node(_)))); + if !batchable { + for (t, n) in converted { + self.apply_one_merge(t, n); + } + return Ok(()); + } + py.allow_threads(|| self.apply_merges_batch(&converted)); + self.merges.extend(converted); + Ok(()) + } +} + +impl Trainer { + fn apply_one_merge(&mut self, token_g: GraphV, merge_nodes: Vec) { // Applying pre-merges to a corpus (e.g. SuperBPE phase-1 → phase-2) // walks every document. Canonicalize children once, then merge each // unique word once per call and rebuild docs by pointer lookup. if let GraphV::Unconn(subs_arc) = &mut self.graph { let subs = Arc::make_mut(subs_arc); - let registry = self - .premerge_registry - .get_or_insert_with(|| { - WordRegistry::new( - canonicalize_subs(subs), - std::sync::Arc::new(CandInterner::new()), - ) - }); + let registry = self.premerge_registry.get_or_insert_with(|| { + WordRegistry::new( + canonicalize_subs(subs), + std::sync::Arc::new(CandInterner::new()), + ) + }); apply_premerge(subs, registry, &token_g, &merge_nodes); } else { self.graph = self.graph.merge(&token_g, &merge_nodes); } self.merges.push((token_g, merge_nodes)); - Ok(()) + } + + fn apply_merges_batch(&mut self, list: &[(GraphV, Vec)]) { + let GraphV::Unconn(subs_arc) = &mut self.graph else { unreachable!() }; + let subs = Arc::make_mut(subs_arc); + let registry = self.premerge_registry.get_or_insert_with(|| { + WordRegistry::new( + canonicalize_subs(subs), + std::sync::Arc::new(CandInterner::new()), + ) + }); + + // Word-level replay, tracking each pre-call pointer's final graph and + // the exact step (if any) at which its word collapsed to a Node. + let mut cur2orig: FxHashMap> = + registry.words.keys().map(|&p| (p, vec![p])).collect(); + let mut final_of: FxHashMap = FxHashMap::default(); + let mut collapse_step: FxHashMap = FxHashMap::default(); + for (k, (token, nodes)) in list.iter().enumerate() { + let changed = registry.apply_merge_step(token, nodes); + for (oldp, newg) in changed { + let origs = cur2orig.remove(&oldp).unwrap_or_default(); + let is_node = matches!(newg, GraphV::Node(_)); + if let Some(np) = crate::graph::arc_ptr(&newg) { + cur2orig.entry(np).or_default().extend(origs.iter().copied()); + } + for &o in &origs { + final_of.insert(o, newg.clone()); + if is_node { + collapse_step.entry(o).or_insert(k); + } + } + } + } + + // Byte pairs of the list, for doc-top match detection. + let pairs: Vec<(Vec, Vec)> = list + .iter() + .map(|(_, n)| (n[0].to_bytes(), n[1].to_bytes())) + .collect(); + let mut merge_index: FxHashMap<(Vec, Vec), Vec> = FxHashMap::default(); + for (k, pair) in pairs.iter().enumerate() { + merge_index.entry(pair.clone()).or_default().push(k); + } + + // Per doc: simulate the doc-top matches sequential application would + // have performed. At step k a match needs both children to be Nodes + // (from the start, or collapsed at a step < k) with the pair's bytes; + // an event's token participates from the following step. Then rewrite + // the doc once: surviving children take their final word state, + // consumed spans become Nodes. + subs.par_iter_mut().for_each(|doc| { + let GraphV::Seq(children_arc) = doc else { + // Single-word docs have no doc-top level; replay directly. + for (token, nodes) in list { + if let Some(new_doc) = doc.try_merge(token, nodes) { + *doc = new_doc; + } + } + return; + }; + let children: &[GraphV] = children_arc; + // (bytes, first step at which this child is a Node, source child) + let mut items: Vec<(Vec, usize, Option)> = children + .iter() + .enumerate() + .map(|(j, c)| { + let avail = if matches!(c, GraphV::Node(_)) { + 0 + } else { + match crate::graph::arc_ptr(c).and_then(|p| collapse_step.get(&p)) { + Some(&k) => k + 1, + None => usize::MAX, + } + }; + (c.to_bytes(), avail, Some(j)) + }) + .collect(); + let mut any_event = false; + loop { + let mut best: Option = None; + for w in items.windows(2) { + if w[0].1 == usize::MAX || w[1].1 == usize::MAX { + continue; + } + if let Some(ks) = merge_index.get(&(w[0].0.clone(), w[1].0.clone())) { + let need = w[0].1.max(w[1].1); + if let Some(&k) = ks.iter().find(|&&k| k >= need) { + best = Some(best.map_or(k, |b: usize| b.min(k))); + } + } + } + let Some(k) = best else { break }; + any_event = true; + let (ka, kb) = (&pairs[k].0, &pairs[k].1); + let mut out: Vec<(Vec, usize, Option)> = + Vec::with_capacity(items.len()); + let mut i = 0; + while i + 1 < items.len() { + let (l, r) = (&items[i], &items[i + 1]); + if l.1 <= k && r.1 <= k && l.0 == *ka && r.0 == *kb { + let mut t = l.0.clone(); + t.extend_from_slice(&r.0); + out.push((t, k + 1, None)); + i += 2; + } else { + out.push(items[i].clone()); + i += 1; + } + } + if i < items.len() { + out.push(items[i].clone()); + } + items = out; + } + + let changed_child = |c: &GraphV| { + crate::graph::arc_ptr(c).and_then(|p| final_of.get(&p)).cloned() + }; + if !any_event && !children.iter().any(|c| changed_child(c).is_some()) { + return; + } + let mut new_children: Vec = items + .into_iter() + .map(|(bytes, _, src)| match src { + Some(j) => changed_child(&children[j]).unwrap_or_else(|| children[j].clone()), + None => GraphV::Node(bytes), + }) + .collect(); + *doc = if new_children.len() == 1 { + new_children.pop().unwrap() + } else { + GraphV::new_seq(new_children) + }; + }); } }