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(); }