Skip to content

feat(tokenizer): add early exit for truncation - #1990

Merged
McPatate merged 5 commits into
mainfrom
feat/early_exit_right_truncation
Apr 2, 2026
Merged

feat(tokenizer): add early exit for truncation#1990
McPatate merged 5 commits into
mainfrom
feat/early_exit_right_truncation

Conversation

@McPatate

@McPatate McPatate commented Mar 26, 2026

Copy link
Copy Markdown
Member

Start tackling #1573.

Here we skip tokenizing remaining splits after the max_length of the TruncationParams has been reached in the case TruncationDirection::Right.

This leaves further questions as to how we tackle shortening the pre-tokenization pipeline as well as the tokenization. Since here we introduce the bound check on the token count, we still do the full context's pretok before the token count skip.
In the benchmark I've left for reproduction in case that is of interest to anyone, there's a pre-truncated-input-512 bench that is orders of magnitude faster, because it cuts the actual input text before passing it to the tokenizer.

I also only handled right truncation as it's the simplest, need to figure out how to do left tokenization. I think I'll tackle both topics in subsequent PRs.

Benchmark improvement results (if unspecified, direction == right):

Benchmark Before After Change
truncate-to-512 (full corpus) 1.01 s 732 ms -28%
scaling: 10K chars 1.26 ms 1.06 ms -16%
scaling: 100K chars 13.2 ms 10.5 ms -21%
scaling: 500K chars 71.3 ms 56.6 ms -21%
max_length=128 6.63 ms 5.21 ms -22%
max_length=512 6.62 ms 5.30 ms -20%
max_length=2048 6.77 ms 5.46 ms -20%
max_length=8192 6.74 ms 6.00 ms -11%
right direction 1.27 ms 1.06 ms -16%
left direction 1.27 ms 1.27 ms no change (expected)
no truncation 975 ms 971 ms no change (expected)

@McPatate
McPatate requested a review from ArthurZucker March 26, 2026 17:52
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@McPatate
McPatate force-pushed the feat/early_exit_right_truncation branch from db190ee to a7fb4e2 Compare March 26, 2026 18:15
@McPatate

Copy link
Copy Markdown
Member Author

So actually implementing left truncation support was pretty easy, yields a perf boost consitent with the right truncation of ~17%:

Benchmark Before After Change
left direction 1.27 ms 1.06 ms -17%

@McPatate McPatate changed the title feat(tokenizer): add early exit for right truncation feat(tokenizer): add early exit for truncation Mar 26, 2026

@ArthurZucker ArthurZucker left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mega nice actually

continue;
}

let tokens = tokenize(&split.normalized)?;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you have 1 split, HUGE ass one -> slow -> chinese don't benefit AFAIK.
If you have many small splits -> normal
If you have many big splits -> huge wins (after 2 splits boom your done).

Comment thread tokenizers/Cargo.toml
Comment thread tokenizers/src/tokenizer/pre_tokenizer.rs
Comment thread tokenizers/src/tokenizer/mod.rs
Comment thread tokenizers/benches/truncation_benchmark.rs
Comment thread tokenizers/benches/truncation_benchmark.rs
@McPatate
McPatate merged commit f5fde3d into main Apr 2, 2026
36 checks passed
@McPatate
McPatate deleted the feat/early_exit_right_truncation branch April 2, 2026 12:56
ArthurZucker added a commit to huggingface/transformers that referenced this pull request Jun 3, 2026
tokenizers 0.23.1 early-exits right-direction truncation
(huggingface/tokenizers#1990), so return_overflowing_tokens may omit the
trailing overflow segment when it is shorter than the window. Validate
the overflow content only when that segment is present; truncation is
always checked. Keeps passing on 0.22.x (segment always present).
ArthurZucker added a commit to huggingface/transformers that referenced this pull request Jun 3, 2026
The trailing overflow segment is not essential to validate here, and
0.23.1's truncation early-exit (huggingface/tokenizers#1990) may omit it.
Keep the truncation checks; remove the now-dead overflow expectations.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants