`encode_batch` has suboptimal parallelization on high-core systems

[stargazerZJ]

Summary

On systems with 64+ CPU cores, encode_batch achieves significantly lower throughput than using direct par_iter() on the same tokenizer. The difference is 6x on a 128-core system.

Environment

• tokenizers: 0.22.2 (Rust crate)
• System: 128-core Intel Xeon (dual 8462Y+)
• OS: Linux 5.15

Reproduction

Update: the script below was a false positive that cannot be reproduced. please check out the script in the original report #1900 (manual sharding of tokenizer instances / worker-group approach) that works consistently.

use tokenizers::Tokenizer;
use rayon::prelude::*;
use std::time::Instant;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let tokenizer = Tokenizer::from_pretrained("Qwen/Qwen2.5-7B-Instruct", None)?;

    // Generate test data: 200 documents of 100KB each
    let data: Vec<String> = (0..200)
        .map(|_| "hello world ".repeat(8333))  // ~100KB per doc
        .collect();
    let texts: Vec<&str> = data.iter().map(|s| s.as_str()).collect();

    // Method 1: encode_batch
    let start = Instant::now();
    let _ = tokenizer.encode_batch(texts.clone(), false)?;
    println!("encode_batch: {:?}", start.elapsed());

    // Method 2: Direct par_iter (same tokenizer instance)
    let start = Instant::now();
    let _: Result<Vec<_>, _> = texts.par_iter()
        .map(|text| tokenizer.encode(*text, false))
        .collect();
    println!("par_iter:     {:?}", start.elapsed());

    Ok(())
}

Add to Cargo.toml:

[dependencies]
tokenizers = { version = "0.22", features = ["http"] }
rayon = "1.10"

Results (64 threads, 200 docs × 100KB)

Method	Time	Throughput
encode_batch	5.04s	2.6M tok/s
Direct par_iter	0.83s	15.8M tok/s
Speedup	6.1x

perf stat Analysis

Metric	encode_batch	par_iter	Ratio
Elapsed time	5.58s	2.46s	2.3x slower
Context switches	107,100	21,404	5x more
CPU migrations	11,710	683	17x more
Sys time	17.7s	3.6s	4.9x more
Effective GHz	1.09	2.03	1.9x lower
IPC	1.07	1.42	lower

The high context switch and CPU migration counts suggest scheduling overhead in the parallelization wrapper.

Root Cause

encode_batch uses into_maybe_par_iter() (via maybe_parallel feature and rayon_cond crate):

// tokenizer/mod.rs line ~1286
let mut encodings = inputs
    .into_maybe_par_iter()
    .map(|input| self.encode(input, add_special_tokens))
    .collect::<Result<Vec<Encoding>>>()?;

This wraps Rayon's parallel iterator in CondIterator to support the TOKENIZERS_PARALLELISM environment variable. The wrapper appears to introduce significant overhead on high-core systems.

Note on TOKENIZERS_PARALLELISM

The TOKENIZERS_PARALLELISM feature is valuable for avoiding thread explosion in nested parallel contexts. The issue is not with the feature itself, but with the performance overhead of the current implementation at high core counts.

Possible Non-Breaking Improvements

1. Internal optimization: Investigate and optimize the rayon_cond / CondIterator implementation for high-core systems. This would be completely transparent to users.

2. Check once and branch: Instead of wrapping each iterator operation, check TOKENIZERS_PARALLELISM once at the start of encode_batch and call either parallel or sequential code path:

   if current_parallelism() {
       inputs.into_par_iter().map(...).collect()
   } else {
       inputs.into_iter().map(...).collect()
   }

   This preserves the feature while avoiding per-operation wrapper overhead.

3. Add optional method (additive, non-breaking): Add encode_batch_direct() or similar that uses direct Rayon without the wrapper, for users who know they don't need nested parallelism protection.

Workaround

Users can bypass encode_batch and use direct Rayon:

use rayon::prelude::*;

let encodings: Result<Vec<_>, _> = texts.par_iter()
    .map(|text| tokenizer.encode(text, false))
    .collect();

Related

• Performance: batch_encode scales poorly on high-core Server CPUs compared to sharded tokenizer instances #1900 (original report)

[ArthurZucker]

this sounds great, and thanks for the repro!
Will see if I have time to fix asap, but you can also submit a PR for a fix if you want!!

[stargazerZJ]

Thanks a lot for taking a look (and for offering a PR path)!

Quick update from my side: I wasn’t able to reproduce the large encode_batch vs direct par_iter() gap from the new repro on my machine. Even after switching to longer + random strings and increasing the workload, I’m only seeing a small difference (on the order of ~1.1x, with par_iter slightly faster).

However, the “old trick” / original report #1900 (manual sharding of tokenizer instances / worker-group approach) does reproduce very clearly for me: I’m seeing ~3x+ speedup versus encode_batch on a heavy synthetic workload (I used that trick in production with larger and real datasets). I’ll focus on investigating that path further (profiling / isolating where the overhead comes from), since it seems to capture the scaling issue more reliably.

I’ll share concrete profiling findings and a minimal patch direction once I have them.

[ArthurZucker]

Sounds great 🫡 I'll keep an eye out here thanks a lot!

[wheynelau]

Hey I found this issue very interesting, did you manage to find any issue with <64c? I am trying to run on a 64 core but I cannot seem to replicate it

[stargazerZJ]

Hey I found this issue very interesting, did you manage to find any issue with <64c? I am trying to run on a 64 core but I cannot seem to replicate it

On my 13900K (32 core desktop CPU), the code in this current repo indeed does not reproduce. It (as aforementioned) is an false positive result that cannot reproduce on my 128C machine either. (Sorry for that!)

However, the reproduce script from the original repo #1900 (manual sharding of tokenizer instances / worker-group approach) can achieve a 11% speed increase on 13900K

❯ ./bench_data/target/release/tokenizer_bench -i bench_data -n 4 --compare-batch -b 1000 -w 32 --threads-per-tokenizer 16 --synthetic
[INFO] Tokenization Benchmark
[INFO] Workers: 32
[INFO] Batch size: 1000
[INFO] Loading tokenizer: Qwen/Qwen2.5-Coder-32B-Instruct
[INFO] Tokenizer loaded successfully
[INFO] Generating 10000 synthetic strings of 204800 bytes each
[00:00:05] ======================================== 10000/10000 strings                                                                                                                                                                                                   [INFO] Generated 10000 strings, total size: 1953.12 MB

=== Worker Pool Tokenization (32 workers, 16 threads per tokenizer) ===
[INFO] Creating 2 tokenizer groups                                                                                                                                                                
[00:01:07] ======================================== 10000/10000 PRs (147.9238/s)                                                                                                                                                                                          [THROUGHPUT] PRs: 10000 (147.9/s) | Tokens: 1346490130 (19913016.0/s) | Data: 1.91 GB (28.88 MB/s) | Elapsed: 67.6s

=== Batch Tokenization (batch_size=1000) ===
[00:01:15] ======================================== 10000/10000 PRs (131.8796/s)                                                                                                                                                                                          [THROUGHPUT] PRs: 10000 (131.9/s) | Tokens: 1346490130 (17757368.1/s) | Data: 1.91 GB (25.76 MB/s) | Elapsed: 75.8s

=== Performance Comparison ===
Worker Pool: 67.6s | Batch: 75.8s | Speedup: 0.89x

[INFO] Benchmark complete

Previously, the original script has a larger (3x) speed gain on a 192C machine and a 128C machine.
Thus, my guess is that the effect is more visible on machines with multiple sockets and non-uniform memory access

[stargazerZJ]

Quick update from my side:

Recently, I have tried to identify multiple ways to decode the exact reason why "tokenizer sharding" works. I have attempted

• sharding the cache: using multiple independent caches, each shared by a smaller number of threads, to reduce cache contention.
• change result collection scheme: The default Rayon collect::<Vec<Encoding>>() path appears to introduce substantial overhead at high core counts for large Encoding objects. Instead, I use try_for_each to write each (index, Encoding) into a pre-sized slot (UnsafeCell<Option<Encoding>> per index; each index written exactly once) and then gathers into a Vec<Encoding> in index order.
• **data synthesis method": generating simpler repeated patterns probably increases cache hit rate, and affects relative speed ratio of different methods. I'll checkout the facebook/xnli used in the official benchmark script.

However, none of these split tricks work on its own, and is not very reproducible (by reproducible I mean copying the ai-generated implementation description to a fresh workspace and let ai reproduce the results based on description).

Although I'm now having some trouble finding the real reason, I'll share more findings (that have at least convinced myself) once I have them. Thank again for your reproducing efforts!

[wheynelau]

@stargazerZJ I agree with you. Not just tokenizers, but from my experience with rayon in general, performance isn't linear after a number of threads.

I did do the tests with criterion so it shouldn't be a cache issue. Personally I use multiple pools if I need to do very heavy concurrency work, like an outer par_chunk and inner par_iter. But I would say this is the user's responsibility not the maintainer, else it's very difficult to have to configure this.

I am nowhere near a performance expert, but this might interest you.
https://gendignoux.com/blog/2024/11/18/rust-rayon-optimized.html

Edit: Performance slowed off, but did not reduce on a cluster. Unable to test without a bare metal cluster. Kubernetes may not be accurate due to CFS. Local PC was mix of E and P cores, so may not be accurate as well.