From 54c903e8076b1ae9ee511c645170b1f72e4b1514 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 5 Aug 2026 11:52:51 +0900 Subject: [PATCH] feat(tk-encode): make the encode pool opt-out behind a `parallel` feature MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `parallel` is added to the default feature set, so nothing changes unless you ask. With `--no-default-features` the private rayon `ThreadPool` is not built, no worker threads are spawned, and the job-splitting machinery is not compiled: `encode` takes the serial path that a small batch already took. Worth 66,272 bytes on the stripped `binsize_pipeline` example, 2,124,096 -> 2,057,824, measured with `progressbar` held on so the number is the pool and not indicatif. **This does not make rayon optional, and the title deliberately does not say it does.** rayon remains a hard dependency with the feature off, because `ptr_hash` declares it unconditionally to build the MPHF, and `rayon-cond` and `rdst` pull it too: rayon v1.12.0 |-- ptr_hash v2.0.2 -> tk-encode |-- rayon-cond v0.4.0 -> tk-encode |-- rdst v0.20.14 So `cargo tree` still shows rayon either way. What the flag buys is a build with no thread pool and no parallel encode path, which matters for embedded and single-threaded hosts; it is not a dependency reduction. Making rayon genuinely optional would need serial fallbacks at all 19 `maybe_par_*` call sites (7 here, 12 in tk-train) plus `tk-train` opting in — and it would still not remove rayon, for the reason above. That refactor also lands in `tokenizer/mod.rs`, `encoding.rs` and `padding.rs`, which the pipeline-only work replaces. Not worth doing unless `ptr_hash` changes, at which point it should be re-derived rather than resurrected. Two implementation notes. The parallel batch body moves out of `encode` into a gated `encode_parallel`, which takes the owned input storage rather than a slice because `JobCore` ends up owning it. And with the feature off, the striding and segment planning helpers have no caller; they are left compiled-but-unused under a scoped `allow(dead_code)` rather than cfg'd one by one, because they reference each other densely and LTO drops all of it from the binary regardless — cfg-ing each would be churn in files the pipeline-only refactor is replacing. Tests pass both ways: 291 lib tests with the feature, 290 without (the pool's own test is gated with the module), plus 15 parallel-oracle and 20 doc tests unchanged. Both configurations build warning-clean. --- tokenizers/tk-encode/Cargo.toml | 7 +- tokenizers/tk-encode/src/normalizers/mod.rs | 3 + tokenizers/tk-encode/src/tokenizer/mod.rs | 1 + .../tk-encode/src/tokenizer/pipeline.rs | 111 ++++++++++++------ tokenizers/tk-encode/src/utils/parallelism.rs | 7 ++ 5 files changed, 92 insertions(+), 37 deletions(-) diff --git a/tokenizers/tk-encode/Cargo.toml b/tokenizers/tk-encode/Cargo.toml index b67ebfccc..5e9f56f7b 100644 --- a/tokenizers/tk-encode/Cargo.toml +++ b/tokenizers/tk-encode/Cargo.toml @@ -80,7 +80,12 @@ libc = "0.2" # arbitrary (non-GPT) regex or the `Replace` normalizer — the atomsplit-native pre-tokenizers (GPT-2, # cl100k, deepseek, the class family, char-delimiter) need no backend. Without it a stub compiles and # those arbitrary-regex paths error at load. Enable with `--features fancy-regex`. -default = ["progressbar"] +# `parallel` is in the default set: the encode pool is on unless you opt out with +# --no-default-features. Note this does NOT make rayon optional — `ptr_hash` declares it +# unconditionally for the MPHF build, so rayon is linked either way. What the flag controls is +# whether we spin up a pool and compile the job-splitting machinery. +default = ["progressbar", "parallel"] +parallel = [] progressbar = ["indicatif"] http = ["hf-hub"] unstable_wasm = ["fancy-regex", "getrandom/wasm_js"] diff --git a/tokenizers/tk-encode/src/normalizers/mod.rs b/tokenizers/tk-encode/src/normalizers/mod.rs index 1763101c3..f066d5348 100644 --- a/tokenizers/tk-encode/src/normalizers/mod.rs +++ b/tokenizers/tk-encode/src/normalizers/mod.rs @@ -1,3 +1,6 @@ +// `preserves_stride_boundaries` is consulted only when planning a parallel stride. +#![cfg_attr(not(feature = "parallel"), allow(dead_code))] + pub mod bert; pub mod byte_level; pub mod precompiled; diff --git a/tokenizers/tk-encode/src/tokenizer/mod.rs b/tokenizers/tk-encode/src/tokenizer/mod.rs index f4eff8021..18a76fd9a 100644 --- a/tokenizers/tk-encode/src/tokenizer/mod.rs +++ b/tokenizers/tk-encode/src/tokenizer/mod.rs @@ -27,6 +27,7 @@ mod encoding; pub mod normalizer; pub mod pattern; pub mod pipeline; +#[cfg(feature = "parallel")] pub(crate) mod pool; pub mod pre_tokenizer; mod serialization; diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index eb647501b..c52e238f8 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -26,12 +26,19 @@ //! **Fork safety** lives in `pool`: a `pthread_atfork` child abandons the //! stale pool without touching it and lazily rebuilds. +// Without the `parallel` feature the job-splitting machinery in this file — `ParallelPlan`, +// `JobCore`, striding, segment planning — has no caller. It is left compiled-but-unused +// rather than gated item by item: those helpers reference each other densely, and LTO drops +// all of it from the binary regardless, so cfg-ing each one would be churn for no bytes. +#![cfg_attr(not(feature = "parallel"), allow(dead_code))] + use std::cell::{RefCell, UnsafeCell}; use std::convert::TryInto; use std::ops::Range; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::Arc; +#[cfg(feature = "parallel")] use rayon::prelude::*; use std::{borrow::Cow, convert::TryFrom}; @@ -61,7 +68,9 @@ use crate::{ ModelWrapper, PostProcessorWrapper, PreTokenizerWrapper, Tokenizer, }; -use super::{pool, Result, SplitDelimiterBehavior}; +use super::{Result, SplitDelimiterBehavior}; +#[cfg(feature = "parallel")] +use super::pool; pub use atomsplit::fsm::Span; @@ -1024,45 +1033,52 @@ impl JobCore { } let total: usize = chunks.iter().map(Vec::len).sum(); - // Small enough (or no pool): serial concat. - const PAR_COMMIT_MIN: usize = 64 * 1024; - let Some(pool) = (total >= PAR_COMMIT_MIN).then(pool::rayon).flatten() else { - let mut out = Vec::with_capacity(total); - for c in &chunks { - out.extend_from_slice(c); + // Parallel scatter when the result is big enough and a pool exists. Without the + // `parallel` feature none of this is compiled and the serial concat below is the only + // path. Nested `if let` rather than a let-chain: this crate is edition 2018. + #[cfg(feature = "parallel")] + { + const PAR_COMMIT_MIN: usize = 64 * 1024; + if total >= PAR_COMMIT_MIN { + if let Some(pool) = pool::rayon() { + // Parallel scatter. Prefix-sum offsets, then each chunk copies into its + // own disjoint window of the output. + let mut offsets = Vec::with_capacity(chunks.len()); + let mut acc = 0usize; + for c in &chunks { + offsets.push(acc); + acc += c.len(); + } + let mut out: Vec = Vec::with_capacity(total); + let base = out.as_mut_ptr() as usize; + pool.install(|| { + chunks + .par_iter() + .zip(offsets.par_iter()) + .for_each(|(chunk, &off)| { + // SAFETY: offsets are a prefix sum of the chunk lengths, so + // the windows are disjoint and all lie within `total` + // (reserved above); no two tasks touch the same element. + unsafe { + std::ptr::copy_nonoverlapping( + chunk.as_ptr(), + (base as *mut PipelineToken).add(off), + chunk.len(), + ); + } + }); + }); + // SAFETY: the scatter wrote every element of `0..total`. + unsafe { out.set_len(total) }; + return Ok(out); + } } - return Ok(out); - }; + } - // Parallel scatter. Prefix-sum offsets, then each chunk copies into its - // own disjoint window of the output. - let mut offsets = Vec::with_capacity(chunks.len()); - let mut acc = 0usize; + let mut out = Vec::with_capacity(total); for c in &chunks { - offsets.push(acc); - acc += c.len(); + out.extend_from_slice(c); } - let mut out: Vec = Vec::with_capacity(total); - let base = out.as_mut_ptr() as usize; - pool.install(|| { - chunks - .par_iter() - .zip(offsets.par_iter()) - .for_each(|(chunk, &off)| { - // SAFETY: offsets are a prefix sum of the chunk lengths, so - // the windows are disjoint and all lie within `total` - // (reserved above); no two tasks touch the same element. - unsafe { - std::ptr::copy_nonoverlapping( - chunk.as_ptr(), - (base as *mut PipelineToken).add(off), - chunk.len(), - ); - } - }); - }); - // SAFETY: the scatter wrote every element of `0..total`. - unsafe { out.set_len(total) }; Ok(out) } @@ -1306,6 +1322,29 @@ impl PipelineTokenizer { if total_bytes < PARALLEL_MIN_TOTAL_BYTES { return EncodeHandle::ready(self.encode_serial(&refs)); } + // Without `parallel` there is no pool at all, so every batch goes down the serial path + // a small batch already takes. + #[cfg(not(feature = "parallel"))] + return EncodeHandle::ready(self.encode_serial(&refs)); + + #[cfg(feature = "parallel")] + { + drop(refs); + self.encode_parallel(storage) + } + } + + /// The parallel batch path: inputs are split into `Item`s behind one atomic cursor which + /// pool workers and the consuming thread both drain. Compiled only with the `parallel` + /// feature; without it `encode` goes straight to `encode_serial`. + /// + /// Takes the owned storage rather than a slice because `JobCore` ends up owning it, which is + /// also why `refs` is rebuilt here. + #[cfg(feature = "parallel")] + fn encode_parallel(&self, storage: impl Inputs + 'static) -> EncodeHandle { + let n_inputs = storage.len(); + let refs: Vec<&str> = (0..n_inputs).map(|i| storage.get(i)).collect(); + let total_bytes: usize = refs.iter().map(|s| s.len()).sum(); let Some(pool) = pool::rayon() else { return EncodeHandle::ready(self.encode_serial(&refs)); }; diff --git a/tokenizers/tk-encode/src/utils/parallelism.rs b/tokenizers/tk-encode/src/utils/parallelism.rs index 98eabe393..6263dd73b 100644 --- a/tokenizers/tk-encode/src/utils/parallelism.rs +++ b/tokenizers/tk-encode/src/utils/parallelism.rs @@ -15,6 +15,10 @@ //! This module also defines the `Maybe*` helpers the legacy paths use for //! optional Rayon usage. +// `mark_parallelism_used` and `num_threads_override` exist for the pool, which is not +// compiled without the `parallel` feature. +#![cfg_attr(not(feature = "parallel"), allow(dead_code))] + use rayon::iter::IterBridge; use rayon::prelude::*; use rayon_cond::CondIterator; @@ -100,6 +104,9 @@ pub fn set_parallelism(val: bool) { /// leaked — bounded by the number of calls, so don't call this in a loop). pub fn set_num_threads(num_threads: usize) { NUM_THREADS.store(num_threads, Ordering::SeqCst); + // Nothing to invalidate without the `parallel` feature: there is no pool. The stored value + // is still kept so the setter stays callable and a later build with the feature sees it. + #[cfg(feature = "parallel")] crate::tokenizer::pool::invalidate(); }