Improve throughput with gigatoken and batching - #312
Conversation
flush_batch's fallback path appended None into a list mypy inferred as list[list[int]], which suppressed the type error but made the downstream `is None` check look unreachable. Explicitly type the list as list[list[int] | None] and narrow through a fresh binding instead of relying on a type: ignore. Also caps gigatoken to <0.11.0 since the wrapper depends on its private _hf_config() method, which isn't covered by semver.
Reintroduces the HuggingFace tokenizers Rust bindings as a selectable backend (now the default) alongside gigatoken, via --tokenizer.backend hf|huggingface|gt|gigatoken. This makes it easy to profile the two against each other and de-risks adopting gigatoken, since it's a very new library, by keeping the previously-shipped implementation one flag away. Threads a TokenizerBackend enum through Tokenizer.from_file/ from_pretrained, the tokenizer executor, tokenize_in_parallel, and the CLI config, with early validation so an unknown backend name fails fast instead of deep inside a worker process. Adds gigatoken-backend test coverage alongside the existing default-backend tests, and a --backend flag to the batch-throughput benchmark script. Also fixes a real bug this surfaced: tokenizers.Tokenizer.from_file requires a str, not a PathLike, unlike gigatoken.Tokenizer.
Two lines from the batching/backend-selection work were a few characters over the line-length limit; black wraps them differently than my editor did. No behavior change.
undfined
left a comment
undfined
left a comment
There was a problem hiding this comment.
Looks good overall. Two things to confirm before merging please.
| # Gigatoken recognizes added special tokens by default, which matches | ||
| # HuggingFace when ``encode_special_tokens`` is false. The opt-in flag | ||
| # instead encodes the special-token text as ordinary text. | ||
| if self.encode_special_tokens and any(token in text for text in inputs for token in self.special_tokens): |
There was a problem hiding this comment.
The literal scan misses normalized special-token matches
token in text does not reproduce Hugging Face’s added-token matching rules. Special tokens marked normalized=True can match only after the tokenizer normalizer runs, even when their literal content is absent from the original input.
I reproduced this with a Llama-compatible special token ▁zzzzunique and input zzzzunique: Hugging Face encoded ordinary token pieces with encode_special_tokens=True, whereas this gigatoken path emitted special-token ID 32003.
For correctness, either route every encode_special_tokens=True call through the Hugging Face fallback or implement matching equivalent to the tokenizer’s normalization and added-token rules.
| if self.encode_special_tokens and any(token in text for text in inputs for token in self.special_tokens): | ||
| fallback = self._hf_special_token_fallback | ||
| fallback.encode_special_tokens = True | ||
| return [encoding.ids for encoding in fallback.encode_batch(inputs)] |
There was a problem hiding this comment.
The fallback unexpectedly applies the tokenizer post-processor
fallback.encode_batch(inputs) defaults to add_special_tokens=True, unlike the Hugging Face path above, which explicitly disables it. Consequently, if any input contains recognized special-token text, tokenizers with a post-processor can inject BOS/EOS tokens into every document in the batch.
With the Llama fixture, I reproduced Hugging Face producing [8656, 1426, 2] while the gigatoken path produced [1, 8656, 1426, 2]. With paragraph segmentation, the post-processor may also run once per paragraph. Double check, but I think we always want this here:
fallback.encode_batch(inputs, add_special_tokens=False)
2 participants
This PR improves the speed of the
dolmatoolkit by (a) using gigatoken as a drop-in alternative for the huggingface backend and (b) adds document batching. For large datasets, this can give a substantial improvement in processing time; in my use cases, it improves from ~100kt/s to ~500Mt/s.Changes
gtsupport):Tokenizernow wraps bothgigatoken.Tokenizerand huggingface'stokenizers.Tokenizer. Selectable with a new--tokenizer.backendflag. When--tokenizer.fastisFalse, this has no effect.tokenize_filenow accumulates documents into batches (bounded by a new--tokenizer.batch_size / --tokenizer.batch_max_bytesCLI config, defaulting to 64 docs / 8 MiB) and callsencode_batchinstead of encoding per-document. Per-record error handling is preserved, so if a batch fails to encode, it falls back to encoding records individually so one bad document doesn't drop its neighbors.gigatoken>=0.10.0,<0.11.0topyproject.toml. Pinned with an upper bound since the wrapper relies on gigatoken's private_hf_config()method, and the package is still young/moving fast.Testing
On dolma's test data, using the
gtbackend over thehfbackend yields a ~10x speedup; when adding source batching, this extends to a ~18x speedup:encodeper doc)encode_batch, size 64)