Continue Part 1 of a series on rebuilding the modern AI stack from first principles. The code and tests for this sprint can be found here https://github.com/Badjessa-git/ai-stack-sprints/tree/main/sprint-01-bpe-tokenizer

In the first part of this series, we ended up with a working BPE tokenizer. For this part in our series, we explore two questions:

  • Can we make this faster without changing the outputs?
  • How do production systems run tokenizers at scale? What are the trade-offs?

Before optimizing anything, we profiled v1 to see where the time was actually going;

In training, about 37% of the time was spent recounting every pair from scratch after each merge, and about 45% was spent on the merges themselves. Picking the most frequent pair was under 2%.

In decode, almost all of the time was spent searching backwards through the merges for every token being expanded.

We also had one rule for everything below: v2 had to produce the exact same output as v1, same merges and same encodings, or the numbers would mean nothing.

Optimizations

train(): our stats dictionary is being sorted every round. In our v1 tokenizer, we sorted the stats dictionary on every round just to find the pair with the maximum count. Python's sorted function uses the Timsort algorithm, which has a time complexity of O(n log n).
We don't really need to apply a sorting algorithm when we're looking for the maximum value. We can simply walk through the dictionary, keeping track of the maximum value seen so far.

On its own, the new selection measured cheaper. Timed on a representative pair sheet (~1,183 distinct pairs), per selection round:

- sorted(...) approach: 0.312 ms - Built-in max(): 0.101 ms - Champion loop: 0.044 ms

The operation itself got ~7x cheaper.

However, this was not where train spent most of its time. This search was about 2% of the train step. We did make some improvement, but there was still more headroom.

decode(): trading time complexity for space complexity at our pair lookup step. As part of our decode function, we need to figure out the pair of numbers that resulted in the new id. In v1, this is how we approached it:

matching_key = next((key for key, val in self.merges.items() if val == token_id), None)

At each step we search through our merges to find which pair produced that id. This is a cost of O(n) that is paid at every level of our recursion.

The fix here was simple: we built a lookup table where we store the pair of numbers that made a specific id.

self.merges[(left, right)] = id new_token_ids = merge_pair( token_ids=token_ids, pair=(left, right), new_id=id ) self.__byte_lookup_table[id] = (left, right) ## <-- Lookup table

Then in our recursive decode function, we replace the scan with a simple lookup, which has a time complexity of O(1), removing the continuous scan done previously.

left, right = self.__byte_lookup_table[token_id]

Comparing the decode run, we now see a 12.5x improvement in the decode step:

corpus: 139,380 bytes (sprint corpus x30) sanity check: merges and encodings identical ✓ v1 (frozen) : decode[30000 tokens] 0.124s v2 (current): decode[30000 tokens] 0.010s speedup (v1 / v2): decode 12.53x

train(): update the count of the pairs only affected by the merge instead of recounting everything. During every iteration of the loop in v1, we rebuilt our count map using get_stats. While that is a reasonable choice, rebuilding the count map every iteration feels wasteful, since only a subset of pairs are affected after a merge.

Here's a diagram to describe how it was done:

One merge round in six steps: v1 recounts all pairs after a merge, v2 updates only the pairs the merge touched

As a result, we saw a 1.44x improvement in our training steps. The max-scan change from the first point contributed to this as well, but not significantly. Most of this gain is the counting.

python3 part2/compare.py corpus: 139,380 bytes (sprint corpus x30) sanity check: merges and encodings identical ✓ v1 (frozen) : train 2.448s v2 (current): train 1.705s speedup (v1 / v2): train 1.44x

How this differs from production tokenizers

Production tokenizers split the text before BPE ever runs. GPT-style tokenizers use a regex to break the text into chunks (words, numbers, punctuation), and merges never cross those boundaries. Ours runs BPE over the raw byte stream, so a merge can swallow a space and half of the next word. SentencePiece, which the Llama family uses, does it differently and treats the space as just another character.

They are also not written in Python. OpenAI's tiktoken and Hugging Face's tokenizers run the same BPE algorithm, but with a Rust core behind a Python API. Our decode win came from a better data structure. Theirs comes from that, plus a much smaller constant factor on every step, and batch parallelism on top.

The vocabulary sizes are on a different scale too. We trained to 512. GPT-2 uses ~50k, GPT-4 uses ~100k, and GPT-4o uses ~200k. Every entry in the vocabulary costs real parameters. A 100k vocabulary at a dimension of 4,096 is roughly 400M parameters in the embedding table alone, and the rarer tokens get trained less. So vocabulary size is a trade-off, not just a case of bigger being better: bigger vocabularies shorten sequences, but they cost memory and training quality.

The counting is structured differently as well. Production trainers count pairs over a frequency table of distinct chunks, instead of one long token stream, so a merge only touches the chunks that contain that pair. Our incremental counting in change 3 had to track positions on the raw stream, which is exactly the bookkeeping their approach avoids.

So yes, it got faster without changing a single output: 1.44x on training, 12.5x on decode. Production is mostly the same algorithm, with the bookkeeping done up front and the hot loop written in Rust. Next in the series: teaching the model to learn, with a scalar autograd engine.

References

  • OpenAI, tiktoken — the Rust-core tokenizer library; the source for the cl100k (~100k, GPT-4) and o200k (~200k, GPT-4o) vocabulary sizes and the regex pretokenization patterns.
  • Hugging Face, tokenizers — production tokenizer library with a Rust core behind Python bindings.
  • Kudo & Richardson, SentencePiece (2018) — the other production family; treats whitespace as a character. Used with a 32k vocabulary in Llama 2.
  • Sennrich, Haddow & Birch, Neural Machine Translation of Rare Words with Subword Units (2016), and the subword-nmt implementation — BPE for production use, counting pairs over word-frequency tables rather than a raw token stream.
  • Radford et al., GPT-2 (2019) — the ~50k (50,257) byte-level BPE vocabulary.