Preface

Today with many Large-Language-Models (LLMs) available, we mostly know how to use them (either via some chat interface or API), but knowing how they work is a different beast.

Yet this seems normal for many people, and itโ€™s totally fine! I personally donโ€™t know how the engine of my car exactly works to the screw, or how to do an open heart bypass surgery.
But yet these fascinating machines somehow can answer many questions, such as how to write code, math, general trivia, and more.

There are many blog-posts, tutorials, videos, and lectures by many talented people that cover these things in different levels of depth. But personally, maybe one thing that is missing from many of these articles is the human curiosity.
Itโ€™s not about just explaining the concept, putting little diagrams, maybe a snippet of code, and calling it a day.
It is about the blood, sweat, and tears that come from writing, optimizing, and getting burned yourself that traumatize you and cements that experience deep inside your mind (and of course the sense of accomplishment and triumph when you finally succeed).

In the following series of blog posts, I, unless procrastination wins yet again, am going to go deep-dive into LLMs. This series of blog posts will not just cover the theory, major concepts, and blindly paste code as yet another success story of โ€œpaste this and it worksโ€.
Weโ€™ll go over the code from the basic-level, and proceed with trial and error of advanced concepts and optimizations. I hope to make sure to include all the steps along the way, all the little tidbits, and cover all the missing pieces that are required for creating an LLM.

I hope to make you appreciate some of the major engineering and elegance that goes into some of the stages that go from raw texts until you get the final LLM.

Why We Need Encodings?

In this series weโ€™ll focus on the Byte-Pair-Encoding (BPE) algorithm. This algorithm is widely used by almost all major LLM models, other models use almost identical (in essence) algorithms (e.g., sentence-piece). Similar to saying that one drives Lamborghini and another a Ferrari.

Before we go into BPE, letโ€™s first try to understand why we need encoding in the first place.

When we interact with our favorite LLM, we do it either via chat or API, but one common trait in both is that we give it a text, and get a text in response. This makes sense, we humans use text as written language in books, literature, and even in this blogpost. So what really makes it special?

Interaction

To make the story short, unlike classical computer programs and dissimilar to humans, LLMs are just one big math equation (for those who are familiar thatโ€™s fine, for those that arenโ€™t, weโ€™ll [maybe] go over that in future posts). Just like the math you were taught in 2nd grade, all that LLMs do is repeatedly do things like: \(2 \times 3\) or \(0.5 \times 4\), and sometimes even \(2^5 \). Amazing isnโ€™t it?

So whatโ€™s so special about numbers? well for starters, can someone tell me exactly what is the result of \((\text{โ€œhow to boil an egg?โ€} \times 0.75)\)? Wellโ€ฆ if I ask my 6-year-old niece she might give an answer with confidence, but computers canโ€™t really do that (yet).

So how the hell does ChatGPT give me a recipe for boiling an egg, I wouldโ€™ve starved without his guidance!!
Wait! wait! There is an answer but first we need to build some suspense.

So we need a way to do math on texts, luckily better people already thought about it, and spoiler alert: in computers, text is already just numbers!

Always has been

How Text Is Represented?

This section is optional. The only thing we need later is: UTF-8 turns text into bytes, and each byte is a number from 0 to 255.

In computers one of the most commonly used standards is the Unicode format1, Unicode defines how to represent (as of today) 159,801 characters and 172 scripts in computer programs.
Unicode defines three different encoding standards (how to convert text to numbers and vice versa), the UTF-8, UTF-16, and UTF-32. With UTF-8 being the most widely used.

UTF-8 is dominant for all countries/languages on the internet, is used in most standards, often the only allowed encoding, and is supported by all modern operating systems and programming languages.

So how does Unicode work? Letโ€™s go through a running example using the one from Wikipedia:

UTF-8 encodes code points in one to four bytes, depending on the value of the code point. In the following table, the characters u to z, each representing a hexadecimal digit, are replaced by their constituent 4 bits uuuu to zzzz, from the positions U+uvwxyz:

Unicode by itself doesnโ€™t define anything special, it maps characters/glyphs to an arbitrary number, called code point, between 0 and 1,114,111. So the character โ€˜aโ€™ gets the number 61, the number 6 gets the number 36, and the Japanese Kanji ๆก gets the number 6841. When represented as Unicode values the prefix U+ is attached, so U+0061, U+0036, and U+6841 respectively.

Unicode Table
Image credit: Wikipedia

Letโ€™s do a running example with ๆก. It has the hexadecimal code point U+6841, first thing, we break the number to separate digits so โ€˜6โ€™, โ€˜8โ€™, โ€˜4โ€™, and โ€˜1โ€™. Then we convert each character to a binary representation:
6- 0110
8- 1000
4- 0100
1- 0001

So the binary representation of U+6841 is 0110 1000 0100 0001.
Since U+6841 falls between U+0800 and U+FFFF we know two things:

  1. The final Unicode representation would be 3 bytes long.
  2. The first byte is represented by 1110.

The table tells us that the first byte is 1110wwww, the 2nd byte is 10xxxxyy, and the last byte is 10yyzzzz. So how does it work?
Letโ€™s go back to our binary representation of 0110 1000 0100 0001, as we can see, itโ€™s 16 digits long, and similarly wwww xxxxyy yyzzzz is 16 digits long. So we just assign each character respectively.
wwww will become 0110
xxxxyy will become 100001
yyzzzz will become 000001

Combining all that together we get the final binary representation of:
11100110 10100001 10000001.

One thing to note, is when we look at Unicode in a computer program like Python, we might see either the value b'\xe6\xa1\x81' or [230, 161, 129].

1
2
3
4
>>> "ๆก".encode("utf-8")
b'\xe6\xa1\x81'
>>> list("ๆก".encode("utf-8"))
[230, 161, 129]

230 comes from the binary representation of the first byte 11100110, 161 comes from 10100001, and 129 from 10000001. Since each byte is represented by 8 bits, every byte can uniquely identify numbers between 0 and 255, the combination and ordering of those bytes is what determines which character is displayed.

We can also encode full sentences into UTF-8 and get a list of integers for all characters.

1
2
3
4
5
6
>>> "The quick brown fox jumped over the lazy dog".encode("utf-8")
b'The quick brown fox jumped over the lazy dog'
>>> list("The quick brown fox jumped over the lazy dog".encode("utf-8"))
[84, 104, 101, 32, 113, 117, 105, 99, 107, 32, 98, 114, 111, 119, 110,
32, 102, 111, 120, 32, 106, 117, 109, 112, 101, 100, 32, 111, 118,
101, 114, 32, 116, 104, 101, 32, 108, 97, 122, 121, 32, 100, 111, 103]

If anyone wants to expand their knowledge about Unicode and UTF-8, I recommend reading Nathan Reedโ€™s post about Unicode2.

Encodings In LLMs

Wait? how is all of that related to LLMs?
Donโ€™t worry we are getting there, right now we saw a way where we can represent texts as a list of numbers. So now all we need to do is go back to our previous example: \((\text{โ€œhow to boil an egg?โ€} \times 0.75)\)
and just change it to:
\((\text{list(โ€œhow to boil an egg?โ€.encode(โ€œutf-8โ€))} \times 0.75)\)?

Close enough but there are more steps along the way.
Remember the brown fox from above, letโ€™s do some counting.
The sentence contains 9 words, and 44 characters in total (36 without the spaces), and similarly the UTF-8 list contains 44 numbers. This is mostly due to backward compatibility with ASCII encoding so each English character is encoded into one number exactly. But what if we are in a different language? Letโ€™s say for example: Japanese:

1
2
3
4
5
6
7
>>> "็ด ๆ—ฉใ„่Œถ่‰ฒใฎ็‹ใŒๆ€ ใ‘่€…ใฎ็Šฌใ‚’้ฃ›ใณ่ถŠใˆใŸใ€‚".encode("utf-8")
b'\xe7\xb4\xa0\xe6\x97\xa9\xe3\x81\x84\xe8\x8c\xb6\xe8\x89\xb2\xe3\x81\xae\xe7\x8b\x90\xe3\x81\x8c\xe6\x80\xa0\xe3\x81\x91\xe8\x80\x85\xe3\x81\xae\xe7\x8a\xac\xe3\x82\x92\xe9\xa3\x9b\xe3\x81\xb3\xe8\xb6\x8a\xe3\x81\x88\xe3\x81\x9f\xe3\x80\x82'
>>> list("็ด ๆ—ฉใ„่Œถ่‰ฒใฎ็‹ใŒๆ€ ใ‘่€…ใฎ็Šฌใ‚’้ฃ›ใณ่ถŠใˆใŸใ€‚".encode("utf-8"))
[231, 180, 160, 230, 151, 169, 227, 129, 132, 232, 140, 182, 232, 137, 178, 227, 129, 174,
231, 139, 144, 227, 129, 140, 230, 128, 160, 227, 129, 145, 232, 128, 133, 227, 129, 174,
231, 138, 172, 227, 130, 146, 233, 163, 155, 227, 129, 179, 232, 182, 138, 227,
129, 136, 227, 129, 159, 227, 128, 130]

The same text is now represented by 20 characters, but its UTF-8 encoding needs 60 numbers! Thatโ€™s a lot of numbers.
Modern LLMs have a limit called a context window: the maximum number of tokens the model can process at once. This defines the amount of information an LLM can retain in its โ€œbrainโ€ before it needs to either give up on some information to make room for more, or lose access to relevant context and start hallucinating. This context window is fixed for a given model, and it defines the amount of numbers an LLM can do math on.
Modern LLMs have a context window that usually ranges from 128,000 to 1,000,000. If we feed our LLM UTF-8 encoding we are wasting space, we limit ourselves to ~43,000 Japanese characters (assuming context window of 128k). This is not a small number and the amount of information we can pack in 43,000 characters is large, but itโ€™s not efficient.

To show you the end result and show you what we are aiming for Iโ€™ll post here a screenshot from the GPT-4o and later modelsโ€™ tokenizer.

1
2
3
4
5
6
7
8
>>> import tiktoken
>>> enc = tiktoken.get_encoding("o200k_harmony")
>>> enc.encode("The quick brown fox jumped over the lazy dog")
[976, 4853, 19705, 68347, 48704, 1072, 290, 29082, 6446]
>>> enc.encode("็ด ๆ—ฉใ„่Œถ่‰ฒใฎ็‹ใŒๆ€ ใ‘่€…ใฎ็Šฌใ‚’้ฃ›ใณ่ถŠใˆใŸใ€‚")
[31399, 23724, 3826, 74914, 4472, 3385, 64920, 6632, 2657, 254, 15707,
8137, 3385, 129581, 7277, 169177, 58490, 30086, 18606, 5598, 788]
>>> 

Look at that! Instead of 44 numbers for our English sentence, we only use 9, and for Japanese we use only 21 numbers.
Now in our 128k context-window we can pack 128,000 English words (not characters), and in Japanese we can pack roughly ~121,000 Japanese characters. This means we can pack more information and more knowledge into our model at the same limited space we have!

The Need For Sub-Word Encoding

Word-Level Encoding

So far we talked about how text is basically represented over the entire internet, but temporarily, letโ€™s go back in time to simple eras, where we try the most obvious encoding scheme: give every word its own ID.

In the past when we wanted to represent texts in machine-learning, and later on in deep-learning models we had various techniques.
Initially we used count-based methods such as TF-IDF3, and later we learned direct word representations from texts. Yoshua Bengio published in 2003 a research about word representations4, and later on word embeddings were popularized after Mikolov published the Word2Vec paper in 20135.
TF-IDF represents a text by looking at which words appear in it, while giving more weight to words that are rare across other texts. Word2Vec is different: it learns a vector for each word based on the words that usually appear around it, so words used in similar contexts get similar representations.

In both techniques each word is associated with a unique ID that is later transformed into an embedding vector that is processed by the model. Putting aside the disadvantages in the idea of static word embeddings (i.e., a single word can have a different meaning in different contexts), the other major downside is the ability to capture all words.

Unicode Table

According to Merriam-Webster6, the English language alone can contain around 1,000,000 different words, with 500,000 being a more sane estimate.
This means that just for English we need to allocate 500,000 different IDs to capture hopefully all different words, and thatโ€™s before we include symbols, numbers, and different languages. Another major issue is typos, we all make typos, Iโ€™m ESL (English Second Language), to make sure I spell โ€˜restaurantโ€™ correctly I literally talk to myself in my head and break it into โ€˜restโ€™ โ€˜tauโ€™ โ€˜rantโ€™; I also type fast which causes me to type โ€˜recieveโ€™ instead of โ€˜receiveโ€™ many times.

So we need to not only accommodate correct spellings of words, but also accommodate entries in our tables for all possible typos we can make. I mean, ChatGPT can handle me writing โ€˜receiveโ€™ and not โ€˜receiveโ€™, so it means he has a way to handle it.

So what happens if we miss a word? Or the user wants to interact in a different language than English?

Unicode Table

Well we canโ€™t really handle it at the model level, the only thing we can do is either raise an error to the user that the word doesnโ€™t exist or we pre-allocate a unique word in our table such as: <unk> whose job is to replace any word that canโ€™t be recognized by our model.

Existing approaches already tried solving that before. For example, fastText7 is an alternative to the classical Word2Vec algorithm that learns whole-word representation based on their n-gram sub-word representation. That means that instead of allocating a unique ID per word, we break down words into overlapping sequences of characters size n (where commonly, n=3). That means that the word: eating is broken down into: <ea, eat, ati, tin, ing, ng>. (the < > is added to differentiate the start/end of a word for the middle). So when we need a representation for the word eating we extract the representations of the n-grams and sum them into a single vector.

fastText embeddings
Image credit: amitness

This way thereโ€™s a lower chance we canโ€™t represent unique, unseen words with our model, since the chance of encountering an unseen n-gram is by far lower than a whole word. The downside is that we rely on fixed size n-grams instead of giving priority to common words. e.g., the word eating above is represented by 6 3-grams, even though eating is a very common word so we wouldโ€™ve wished to assign it a unique single ID instead of spending 6 IDs on it. Additionally, this technique still suffers from the static word embedding issue mentioned above, but thatโ€™s a story for another post.

Byte-Level Encoding

So wait if the issues are unknown words why not use UTF-8? Every character/glyph can be represented as a sequence of numbers ranging from 0 to 255, so we only need to assign 256 slots in our table, and use byte encoding for our texts.

Animation description
Here is a caption for the GIF animation.

If we ignore context-window for a second the issue with byte-encoding is information compression. Think of the sequence of characters ['a', 't', 't'] which is represented by the byte sequence [97, 116, 116]. LLMs are glorious auto-completes, but they are good at it!
But how do we auto-complete that? Does the user mean: Attach? Attack? Attain? Attempt? Attract? Attune? Attest?
As you can see there are many ways to complete the sequence, and thatโ€™s before we go into different language options (although specifically in this example other languages are not possible).

So even though byte representations allow us to keep the vocabulary tiny and essentially never encounter an unknown word, we lose the ability to accurately predict the next token/word since there are way too many options to complete with.
I will say that some research about byte-level language models has been done89.

The BPE Algorithm

So we have two extremes.

Word-level encoding gives us nice, meaningful units, but it breaks on unknown words, typos, rare words, and new languages. Byte-level encoding has the opposite property: it can represent anything, but the model has to work with very small pieces of text, which wastes context and makes prediction harder.

What we want is a tokenizer that starts from bytes, so it can never fail to represent text, but then automatically learns which byte sequences are common enough to deserve their own token.

That is the idea behind BPE in modern LLM tokenizers: learned vocabulary construction.

In the byte-level tokenizer weโ€™ll build, BPE starts with a small vocabulary: the 256 possible byte values. Then it repeatedly finds the most common adjacent pair and merges that pair into a new token. Common sequences become single tokens over time, while rare sequences can still fall back to smaller byte-level pieces.

Historically, BPE was introduced as a compression algorithm, because replacing repeated sequences with shorter symbols makes text smaller. For LLM tokenizers, we use the same mechanism for a different purpose: building a useful vocabulary.

Letโ€™s look at the Wikipedia definition:

In computing, byte-pair encoding (BPE),[1][2] or digram coding,[3] is an algorithm, first described in 1994 by Philip Gage, for encoding strings of text into smaller strings by creating and using a translation table.[4] A slightly modified version of the algorithm is used in large language model tokenizers.

Wait 1994, so why havenโ€™t we used it all the time? Well for starters, we humans love to re-invent the wheel, remember the person that re-invented the integral10?

To see the original compression idea, letโ€™s go over the example from the Wikipedia page. Given the following string: aaabdaaabac we want to be able to encode it into a representation that is shorter in length and one that we can later expand back to the original without loss of information.

The method consists of two steps that repeat until we either hit a stopping criteria or stop ourselves:

  1. Find the consecutive pair with the highest number of appearances.
  2. Replace the pair with a new token that isnโ€™t used in our data.

So letโ€™s start:

Step 1: Count
{aa: 4, ab: 2, ac: 1, bd: 1, da: 1}
Step 2: Assign
aa=Z -> ZabdZabac

Step 1: Count
{Za: 2, ab: 2, ac: 1, bd: 1, da: 1}
Step 2: Assign
ab=Y -> ZYdZYac

Step 1: Count
{ZY: 2, Yd: 1, dZ: 1, Ya: 1, ac: 1}
Step 2: Assign
ZY=X -> XdXac

Weโ€™ll stop here since no pair appears more than once.

Youโ€™re probably thinking to yourself (and if you donโ€™t thatโ€™s fine), why is it better? Now we not only need to have room in memory for XdXac which contains 5 characters, we also need to handle a full dictionary of mappings: {aa=Z, ab=Y, ZY=X} that doesnโ€™t really save space, does it?
Well you are right, but that is just a small example, letโ€™s think of the English language for example. We know that verbs in English that are represented in Present-Progressive end with ing. So letโ€™s take 3 words for example: eat, walk, and drink. Their Present-Progressive forms are eating, walking, and drinking respectively.
In regular word-level representation we would allocate 6 unique IDs for these 6 words: eat->1, walk->2, drink->3, eating->4, walking->5, and drinking->6.
Letโ€™s try to do a small scale version of BPE with a single step. Letโ€™s choose the following: ing=Z.

So now instead of needing 6 unique IDs to represent our 6 words, we can just use 4 IDs: eat->1, walk->2, drink->3, ing->4. Now if we want to convert our word eating or drinking into a sequence of numbers we can represent it with 2 numbers: [1, 4] or [3, 4]. On the one hand, now every word (or sentence) will be represented by more tokens, but on the other hand we need less tokens to represent more words in our language!

Unicode Table

This seems wasteful at first but actually it allows us to learn/design โ€œsmartโ€ tokenizers, since the BPE algorithm works by merging the most common pair at each step, then common words will frequently get merged again and again, and at some point weโ€™ll get a single unique token encoding that covers, while less common words will still have a representation for that but will require slightly more tokens.

But as we see from modern LLMs, for example, the GPT-4o tokenizer (and following versions of GPT for the most part), their tokenizer only consists of 200,000 unique IDs (itโ€™s called o200k, cause the 200k is the vocabulary size), those 200,000 unique IDs are enough to capture all languages, symbols, numbers, and scripts in a way that we can never encounter unknown words but at the same time have a great trade-off between the number of unique IDs and word representations.

BPE Implementation

BPE- The Basic Version

BPE has many good implementations on the internet, Hugging Faceโ€™s famous tokenizers library is written in Rust (Blazingly Fast)11.

Weโ€™ll focus on Karpathyโ€™s minBPE repo, but we wonโ€™t only focus on explaining the code, you can watch Karpathyโ€™s amazing video on tokenizers12.

If you recall from before, we said the BPE algorithm consists of two steps: count and merge. So letโ€™s go over Karpathyโ€™s code for the basic tokenizer and understand how these steps are performed.

Step 1: Count-

1
2
3
4
5
6
7
8
9
10
def get_stats(ids, counts=None):
    """
    Given a list of integers, return a dictionary of counts of consecutive pairs
    Example: [1, 2, 3, 1, 2] -> {(1, 2): 2, (2, 3): 1, (3, 1): 1}
    Optionally allows to update an existing dictionary of counts
    """
    counts = {} if counts is None else counts
    for pair in zip(ids, ids[1:]): # iterate consecutive elements
        counts[pair] = counts.get(pair, 0) + 1
    return counts

On line 7, we first assign our dictionary to an existing one (weโ€™ll see later when itโ€™ll happen), or create a new one if it doesnโ€™t exist. Counts will hold our final pair countings.
Lines 8-9 iterate over the list, at each iteration we get a pair of consecutive items from our list, then we add 1 to their count.
This method allows us to do the counting of all the pairs in our data.

Step 2: Merge-

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def merge(ids, pair, idx):
    """
    In the list of integers (ids), replace all consecutive occurrences
    of pair with the new integer token idx
    Example: ids=[1, 2, 3, 1, 2], pair=(1, 2), idx=4 -> [4, 3, 4]
    """
    newids = []
    i = 0
    while i < len(ids):
        # if not at the very last position AND the pair matches, replace it
        if ids[i] == pair[0] and i < len(ids) - 1 and ids[i+1] == pair[1]:
            newids.append(idx)
            i += 2
        else:
            newids.append(ids[i])
            i += 1
    return newids

This method is also quite simple, given that we have our current list, and we know which pair we are going to merge, this code iterates over our list and every time it sees a mergeable pair (line 11), itโ€™ll append the new ID instead, otherwise (line 14) itโ€™ll append the current ID. The only distinction is that if we merged a pair we need to advance our index by 2 instead of 1.

Now for the entire BPE algorithm:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
    def train(self, text, vocab_size):
        assert vocab_size >= 256
        num_merges = vocab_size - 256

        # input text preprocessing
        text_bytes = text.encode("utf-8") # raw bytes
        ids = list(text_bytes) # list of integers in range 0..255

        # iteratively merge the most common pairs to create new tokens
        merges = {} # (int, int) -> int
        vocab = {idx: bytes([idx]) for idx in range(256)} # int -> bytes
        for i in range(num_merges):
            # count up the number of times every consecutive pair appears
            stats = get_stats(ids)
            # find the pair with the highest count
            pair = max(stats, key=stats.get)
            # mint a new token: assign it the next available id
            idx = 256 + i
            # replace all occurrences of pair in ids with idx
            ids = merge(ids, pair, idx)
            # save the merge
            merges[pair] = idx
            vocab[idx] = vocab[pair[0]] + vocab[pair[1]]

This is the entire algorithm LLMs use to create their tokenizers, everything above it is just optimization, so letโ€™s go over this line-by-line.

First of all, we see the method gets two parameters: text and vocab_size.
text is our training data for the tokens, this is the content that will determine which merges we do, and which tokens get their own unique ID or theyโ€™ll be represented as a sequence of IDs.
vocab_size is a hyper-parameter (decided by us) that defines the final size of our vocabulary, which means how many unique IDs weโ€™re going to handle in our model later on.

Line 3: num_merges defines the number of times weโ€™ll run the BPEโ€™s count and merge operations, it is defined as the vocab_size - 256 since we are allocating a pre-defined set of 256 IDs that our algorithm will start from. These 256 IDs are all the unique UTF-8 bytes we have! That means that at the very least in the worst case our final tokenizer will be equal to a byte-level representation of words via their UTF-8 encoding.
Which means that we guarantee that even if our data is under-represented or we encounter a new language in our model, we can always fall back to byte representation. This is what makes BPE unique, since word merges start from their byte representations and not from n-gram/character representation we can always guarantee we can represent unknowns with our tokenizers. Our model might not reply well with them but thatโ€™s a different story.

Lines 6-7: Here we convert our text to UTF-8 encoding and list representations. Remember before that we showed the example with the quick brown fox? Thatโ€™s exactly what happens in these lines, but for our entire corpus.

Unicode Table

Lines 10-11: Here we initialize our merges dictionary, this dictionary will keep track at each stage which merges happen so itโ€™ll map from a pair of IDs to a new token, e.g., (312, 56) -> 313, this means that in our process of encoding future texts, every time weโ€™ll encounter 312 followed by 56 weโ€™ll replace both of them with 313.
The other dictionary is vocab, this dictionary mainly helps us during the decode stage, where we convert a list of IDs back into textual representation to display to the user. As we said before since we allocate all 256 UTF bytes, we initialize our dictionary with 256 unique tokens. Token #1 maps to text byte #1, token #2 maps to text byte #2, and so on.

Unicode Table
Where b'x' is the representation of the x byte. e.g., b'78' = N and b'255' = รฟ

Lines 12-20: These lines run the BPE algorithm for a set number of iterations. We first count pairs in our list (line 14), then we find the pair with the highest appearance count (line 16).
Once we found the highest pair, we allocate a new token ID, one that isnโ€™t already used (remember at the first iteration we allocated 256), so in the first iteration the newest ID that isnโ€™t used is ID 257, we then apply merge on our list IDs, and every time we see a pair, weโ€™ll switch it out for 257.
You can see below an example of 3 merges:

Unicode Table
In the first iteration the pair (1,2) is the highest and is replaced with 256
Unicode Table
In the second iteration the pair (2,3) is the highest and is replaced with 257
Unicode Table
In the third iteration the pair (257, 258) is the highest and is replaced with 258

Lines 22-23: This updates the merges dictionary to keep track of all the performed merges. Vocab already contains a string representation of the entire sequence of bytes.
In our example above vocab[256] = b'1 2', vocab[257] = b'2 3', and vocab[258] = b'1 2 2 3'.
If you think about it, it makes sense. When we compress a pair we have two options, either of them is a compressed token itself, then it represents a pair, or itโ€™s a base token (0-255). So keeping on all the sequences together, it is easier for us when we have a list of IDs such as: [976, 4853, 19705, 68347, 48704, 1072, 290, 29082, 6446] to convert it directly into text rather than fully back the reverse way and find which two tokens were compressed to create token 19705.

Now that we have the full algorithm letโ€™s test it. Weโ€™ll use a small sample of 1,000 stories from the tiny-stories dataset13.
This dataset contains stories generated by GPT-4, each story in a separate line.

All experiments ran on my 8-year-old Dell XPS 7590 with i7-9750h, bless its soul.

1
2
3
4
5
6
7
time uv run --script bpe_basic.py --vocab-size 1257 --data-path part_00.txt part_01.txt --save-path basic.json
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 2/2 [00:00<00:00, 334.93it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1000/1000 [02:25<00:00,  6.85it/s]
________________________________________________________
Executed in  146.12 secs    fish           external
   usr time  143.35 secs  658.00 micros  143.35 secs
   sys time    1.11 secs  980.00 micros    1.11 secs

The script results in 211,919 tokens (final text after finishing all the BPE merges).
This is not really a great runtime :\. Modern LLMs train on corpora of billions of tokens around 300B (depending on the model size).
If we run our script on this amount of tokens itโ€™ll take only 394 days to learn 1,000 merges! Remember GPT-4o tokenizer contains 200k tokens, that means weโ€™ll only need 2,162 years! Thatโ€™s not good.
Note: number of tokens is usually tokenizer dependent, but numbers will not change greatly so 500B tokens corpus might be ยฑ 50B.

Letโ€™s see if we can improve it!

Adding A Split Pattern

Modern LLMs do not train on raw texts as a continuous stream of tokens, we want to make sure we split texts into logical chunks.
Letโ€™s see a small example:
Sentences: deep learning. deep sleep. deep blue. deep.
Weโ€™ll train on this text of 23 merges only, the result vocab contains:

1
2
3
4
5
256: b'ee', 257: b'eep', 258: b'deep', 259: b'deep ', 260: b'deep l',
261: b'deep le', 262: b'deep lea', 263: b'deep lear', 264: b'deep learn',
265: b'deep learni', 266: b'deep learnin', 267: b'deep learning', 268: b'deep learning.',
269: b'deep b', 270: b'deep bl', 271: b'deep blu', 272: b'deep blue', 273: b'deep blue.',
274: b'deep s', 275: b'deep sl', 276: b'deep sleep', 277: b'deep sleep.', 278: b'deep.'}

We can see several issues in this:

  1. No specific word boundaries, we started learning Frankenstein combinations of words. If we try to encode something like deep log later, the tokenizer will look at the prefix deep l and see it will have an entry in the merges dictionary and will replace it with the token 260 while tokenizing og into individual tokens.

  2. We donโ€™t handle punctuation and spaces efficiently, we learn both a representation for deep, deep., ..

  3. We increase our vocabulary size, by learning meaningless variations.

To do this, weโ€™ll use a regex split pattern, before starting the BPE algorithm, weโ€™ll break down each text into chunks using a regex pattern, this pattern will help us with the problems above.
This results in us not learning BPE over full sentences but rather smaller chunks that represent words, prefixes, or special characters.

For simplicity weโ€™ll use the GPT-4o regex split pattern, itโ€™s quite simple, itโ€™s just

1
2
3
4
5
6
7
8
9
10
11
pat_str = "|".join(
    [
        r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?""",
        r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?""",
        r"""\p{N}{1,3}""",
        r""" ?[^\s\p{L}\p{N}]+[\r\n/]*""",
        r"""\s*[\r\n]+""",
        r"""\s+(?!\S)""",
        r"""\s+""",
    ]
)

Iโ€™m not regex expert but our friend Gemini-3.5-flash can probably explain it:

Branch 1: Grabs lowercase or normal capitalized words, including any attached English contractions (like โ€˜s or โ€˜ll).
Branch 2: Grabs words written in ALL-CAPS, keeping acronyms or shouted text isolated from lowercase text.
Branch 3: Chops up long numbers into small chunks of at most 3 digits to keep the vocabulary size down.
Branch 4: Groups clusters of punctuation, symbols, and math operators (like !!! or ===) into their own boxes.
Branch 5: Isolates line breaks and paragraph newlines so they donโ€™t get fused to the surrounding text.
Branch 6: Isolates code indentation spaces (like Python indents) at the beginning or end of a line.
Branch 7: Catches any remaining normal spaces between words as a final cleanup step.

Letโ€™s add this to our code and see if it improves our runtime.
All we need to do is add:

1
2
3
4
compiled_pattern = re.compile(SPLIT_PATTERN)
chunks = []
for text in texts:
    chunks.extend([list(ch.encode("utf-8")) for ch in compiled_pattern.findall(text)])

And run over chunks instead of texts

1
2
3
4
5
6
7
time uv run --script bpe_pattern.py --vocab-size 1257 --data-path part_00.txt part_01.txt --save-path pattern.json
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 2/2 [00:00<00:00, 1059.03it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1000/1000 [03:28<00:00,  4.79it/s]
________________________________________________________
Executed in  209.11 secs    fish           external
   usr time  205.37 secs    0.20 millis  205.37 secs
   sys time    2.27 secs    1.92 millis    2.27 secs

Welp, our script is 60s slower, we also increase the total token count of the text 787,173, but at least we donโ€™t learn meaningless merges.

Optimizing Via Python-Native Operations

There are probably billions of clever tricks we can employ, bitwise manipulations, etc.
But one goal of this will be to optimize the code while keeping it human-readable.

One thing is to use pairwise, if we look back at get_stats we see that we iterate over a zip.\

1
for pair in zip(ids, ids[1:]): # iterate consecutive elements

This creates two slices of our lists and then zips them. Python provides a C-native version of this logic via pairwise.

1
2
3
4
5
6
7
time uv run --script bpe_pairwise.py --vocab-size 1257 --data-path part_00.txt part_01.txt --save-path pairwise.json
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 2/2 [00:00<00:00, 1213.45it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1000/1000 [03:02<00:00,  5.48it/s]
________________________________________________________
Executed in  182.95 secs    fish           external
   usr time  179.26 secs    1.58 millis  179.26 secs
   sys time    2.32 secs    0.00 millis    2.32 secs

We already got 10s improvements over our previous code. Another thing we can do is change merge.
Right now we need to handle two different lists and append to them, one thing we can do is add in-place edit of the ids list.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def merge(ids, pair: tuple[int, int], new_id: int) -> list[int]:
    """
    Performs in-place update of the list.
    if our list is [a, b, c, d],
    and our merge pair is (b, c) -> X,
    then in-place it'll create a list: [a, X, d, d].
    Then will trim the remaining part to create the final list: [a, X, d].
    """
    read, write = 0, 0
    n = len(ids)
    while read < n:
        if read < n - 1 and ids[read] == pair[0] and ids[read + 1] == pair[1]:
            ids[write] = new_id
            read += 2
        else:
            ids[write] = ids[read]
            read += 1
        write += 1
    del ids[write:]
    return ids

Now letโ€™s run it again.

1
2
3
4
5
6
7
time uv run --script bpe_inplace.py --vocab-size 1257 --data-path part_00.txt part_01.txt --save-path inplace.json
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 2/2 [00:00<00:00, 1026.88it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1000/1000 [02:53<00:00,  5.77it/s]
________________________________________________________
Executed in  173.89 secs    fish           external
   usr time  171.34 secs    0.56 millis  171.34 secs
   sys time    1.49 secs    1.03 millis    1.48 secs

Nice! We saved another 8s of runtime. Tho weโ€™re still slower (yet correct) than our initial algorithm can we do better?

Duplicate Detection

This optimization credit goes to minbpe PR #82, itโ€™s not my own idea but itโ€™s a good one nonetheless that should be implemented.
The idea behind is simple, letโ€™s do a little example:
Our corpus consists of ["the", "the", "the", "the", "the", "the", "the"]. We have 7 texts, each contains the same word "the".
When we run our BPE algorithm we need to build a global statistics of all the possible pairs and their counts. Using the current BPE algorithm weโ€™ll need to iterate 7 times over our texts counting the same things over and over again.
So why instead of iterating 7 times over the list, weโ€™ll iterate once and multiply our counting by 7?
Thatโ€™s exactly what is suggested in the linked PR.

First we change get_stats to include a frequency parameter, this indicates how many times this chunk of text appears in our entire corpus. For the example above: frequency = 7.

1
2
3
4
5
def get_stats(ids, counts=None, freq: int=1):
    counts = {} if counts is None else counts
    for pair in pairwise(ids): # iterate consecutive elements
        counts[pair] = counts.get(pair, 0) + freq
    return counts

Now we need to add an initial counting before we start our BPE to create a mapping of how many repeating chunks are there and their counts (similar to get_stats but on a sentence-level and not pair-level).
Additionally weโ€™ll add a check to only merge if both IDs are in the pair, since running a merge operation over a list that we donโ€™t merge anything is wasteful.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
compiled_pattern = re.compile(SPLIT_PATTERN)
freqs = defaultdict(int)
for text in texts:
    chunks = compiled_pattern.findall(text)
    for ch in chunks:
        freqs[ch] += 1

chunks = [list(ch.encode("utf-8")) for ch in freqs]
frequencies = list(freqs.values())

for i in trange(num_merges):
    stats: dict[tuple[int, int], int] = {}
    for ii, ids in enumerate(chunks):
        stats = get_stats(ids, stats, frequencies[ii])
    pair = max(stats, key=stats.get)
    new_id = 256 + i
    for chunk in chunks:
        if pair[0] in chunk and pair[1] in chunk:
            merge(chunk, pair, new_id)

Time to run it!

1
2
3
4
5
6
7
time uv run --script bpe_freq.py --vocab-size 1257 --data-path part_00.txt part_01.txt --save-path freq.json
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 2/2 [00:00<00:00, 1174.38it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1000/1000 [00:10<00:00, 99.10it/s]
________________________________________________________
Executed in    9.49 secs    fish           external
   usr time    9.41 secs    0.00 millis    9.41 secs
   sys time    0.02 secs    1.86 millis    0.02 secs

WOW! We went down from 170s to 9.5s! thatโ€™s almost 18x improvement over our previous version and 14x over our initial version, while being more correct!
This, like previous versions, produces the same 787,173 token count and the same BPE merges dictionary. That means for a vocabulary size of 200k and 500B our script will only take 38 years. That already gives me hope I might see this script finish some day unlike the 2,100~ years one.

But can we do better?

Multiprocessing

One thing that came to my mind while thinking about it is the following:\

Wait a second, get_stats of one chunk doesnโ€™t depend on another chunk, only the summation at the end. Additionally, merge of each chunk doesnโ€™t depend on others This means we can parallelize those two operations. For get_stats weโ€™ll use map-reduce, and for merge itโ€™s just regular multiprocessing.

True story.

Unicode Table
Map-Reduce overview: each worker gets a shard of the data, performs count and merge

The idea here is that before doing any work, we split the data into shards, then we spawn a worker process per shard which will perform the get_stats on the data, report back its statistics to the main process, then the main process will combine all the results of get_stats from all the workers, find the maximum occurring pair, and assign it a new ID.
The main process will then report that pair and new ID information back to the workers, where each worker will perform the merge on their own shard of data.

I was struggling to find a nice way to do it, Iโ€™m not a systems engineer guy and never really worked with multiprocessing by spawning my own threads and handled communication.
So I tried to like get some vague idea from Gemini on how to perform it without really giving me a solution since I wanted to experiment with that myself and sweat blood and tears (mainly tears) writing it.

Gemini then suggested me to work with pipes, this way instead of spawning and despawning processes at each iteration, each process lives in a while True loop and we just communicate information back and forth.
So I looked up the pipes documentation on python and found this small example:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from multiprocessing import Process, Pipe

def f(conn):
    data = conn.recv()
    print(f"Data from parent: {data}")
    conn.send([42, None, 'hello'])
    conn.close()

if __name__ == '__main__':
    parent_conn, child_conn = Pipe()
    p = Process(target=f, args=(child_conn,))
    p.start()
    parent_conn.send("Hello Pipe!")
    print(parent_conn.recv())   # prints "[42, None, 'hello']"
    p.join()

This example exemplifies how pipes work. A pipe object has two ends (just like real life pipes), one end we give the child process, and the other we keep to ourselves. Each end can either peek to recv data or put stuff in it to send data.
When a recv operation happens, the process hangs until it receives data from the other side.
We can use this to send either our counts from get_stats from the worker to the main process or send our (pair, new_id) from the main process to our worker.

Hereโ€™s a pseudo-code of how this will look like:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def worker_fn(shard, conn):
    chunks, freqs = build_frequencies(shard)
    while True:
        stats = get_stats(chunks, freqs) # Count local stats
        conn.send(stats) # Report back to main process your stats

        # Wait for main process data
        data = conn.recv()
        if data == "STOP": # In case we finished the BPE algorithm break the loop and end the process
            break
        pair, new_id = data
        for chunk in chunks:
            if pair[0] in chunk and pair[1] in chunk:
                merge(chunk, pair, new_id)

    conn.close()


def main_process():
    workers = {} # Technically can be a list and not a dictionary
    for shard_id in range(num_workers):
        parent_conn, worker_conn = Pipe()
        worker_process = Process(fn=worker_fn, args=(data[shard_id], worker_conn))
        workers[shard_id] = (worker_process, worker_conn) # Keep track on the workers
        worker_process.start()

        for i in range(num_merges):
            stats = {} # Our global aggregated stats
            for (_ worker_conn) in workers.values():
                local_stats = worker_conn.recv()
                stats.reduce(local_stats) # Reduce here is just a shortened form for summation

            pair = max(stats, key=stats.get)
            new_id = 256 + i
            for (_ worker_conn) in workers.values():
                worker_conn.send((pair, new_id)) # Broadcast the new id for workers to merge

    for (worker_process, worker_conn) in workers.values(): # Finish and close the processes
        worker_conn.send("STOP")
        worker_process.join()

Letโ€™s try it! The real code is not much different than this and itโ€™s an exercise for the reader to try it themselves.

1
2
3
4
5
6
7
8
time uv run --script bpe_multiprocess.py --vocab-size 1257 --data-path part_00.txt part_01.txt --save-path mp.json
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00, 660.31it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00, 482.77it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1000/1000 [00:12<00:00, 81.58it/s]
________________________________________________________
Executed in   10.38 secs    fish           external
   usr time    4.36 secs    0.61 millis    4.36 secs
   sys time    0.09 secs    1.05 millis    0.09 secs

Wait itโ€™s slower? Why? We basically count multiple in parallel so instead of one process doing X counts, we are doing \(\frac{X}{2}\) counts in parallel, we shouldโ€™ve finished faster.

Wellโ€ฆ the issue right now is the data communication between the main process and the workers. Letโ€™s try to understand what happens when we do send and recv.

Unicode Table
Each time we send data we pay a price from sending python objects.

The issue as can be seen, we need to perform pickle and serialize, and reverse those operations at each iteration per worker.
With a small amount of data itโ€™s not too bad but when we are doing counts over billions of lines our stats dictionary at each worker contains thousands of entries making it slow to send to the main process.

Diff

When we think about it, letโ€™s see exactly what happens to stats on a single list of text bytes.

Letโ€™s take the example from before with our initial list [1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 1, 1] and the merge pair (1, 2) -> 256.
Since itโ€™s the first iteration, our stats dictionary represents the initial distribution of counts over our entire texts (or text in this single line example). So our stats dictionary is:

1
2
3
4
5
6
7
8
{
  (1, 2): 3,
  (2, 2): 2,
  (2, 3): 2,
  (3, 1): 2,
  (1, 1): 1,
  (2, 1): 1,
}

After we perform the merge with (1, 2) -> 256, letโ€™s see how we count our stats at the 2nd iteration:

1
2
3
4
5
6
7
8
{
  (256, 2): 2,
  (2, 3): 2,
  (256, 256): 1,
  (3, 256): 1,
  (1, 1): 1,
  (3, 1): 1,
}

Some things look similar some things look different, letโ€™s see what happens when we investigate the differences between the dictionaries and try to subtract the latter from the former.

1
2
3
4
5
6
7
8
9
(256, 2): +2,
(256, 256): +1,
(3, 256): +1,
(1, 1): +1,
(2, 3): 0,
(3, 1): -1,
(2, 1): -1,
(2, 2): -2,
(1, 2): -3,

We can see that a pair like (2, 3) its count didnโ€™t change, while pairs containing the new ID will have a positive difference, and pairs that one of them is contributing to the new ID might have their count decrease due to not appearing anymore in the list: [256, 256, 2, 3, 257, 2, 3, 1, 1].

Of course over billions of texts the amount of pairs that arenโ€™t affected is vastly greater than the pairs that are affected.
Remember our initial vocab contains all bytes options from 0 to 255, so there are \(254 \times 253 \) potential pairs that their count doesnโ€™t change. But in our current approach we are still trying to count their occurrence. This is wasteful in terms of the final dictionary size that we need to send since we send information that is already known to us (e.g., in our small example, (2,3) doesnโ€™t need to be re-counted).

So to decrease our dictionary size and the amount of information we need to communicate between processes, a better approach would be to count initially the stats dictionary, and then have processes just communicate their local differences.

So what weโ€™ll do is the following:

First letโ€™s update how each worker counts their local stats. One simple way to do it, is to apply get_stats twice, once before the merge with a negative freq, and the second time after we apply the merge with a positive freq.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
for i, chunk in enumerate(chunks):
    if pair[0] in chunk and pair[1] in chunk:
        # We count all pairs in the chunk with negative frequency.
        # This will cause frequencies that will disappear to have negative value.
        get_stats(chunk, diffs, -frequencies[i])
        merge(chunk, pair, new_id)
        # This now updates the same dictionary.
        # If a pair is kept, e.g., (b c) -> X then a pair (a d) is unchanged.
        # Then it'll zero out its diff value.
        # New pairs, e.g., (c X) will now have new count.
        get_stats(chunk, diffs, frequencies[i])

diffs = {k: v for k, v in diffs.items() if v != 0}
conn.send(diffs)

This way we first subtract from all pairs their freq value, then after we merge we count again with positive value. If a value wasnโ€™t affected by the merge, the positive freq cancels out the negative freq which causes the count to be 0 (i.e., no difference in count from last time). If a pair was affected by the merge the second count will reflect the difference, either positive if a pair contains a new pair (which wasnโ€™t accounted for in the first count) or negative if that pair doesnโ€™t exist since we switch one of the IDs to the new ID.

Now to handle the diffs, before we send the data we need to filter all pairs that their count is 0 (didnโ€™t change), that way the dictionary weโ€™ll send will be smaller, and the inter-process communication will be faster.

Aside that we need to apply 3 more changes in our code:

The first one is: Change worker_fn to perform an initial count and send it to the main process.

Since we are handling with diffs now, a diff can only be applied if thereโ€™s an initial stats dictionary that exists. So before we start our count a merge loop, weโ€™ll perform an initial count and send that information to the main process. This way the main processโ€™s job is to hold that initial count, and apply the diffs to it when they arrive. This approach is faster than the previous one since we do not throw away our stats dictionary in the main process just to re-build it again with all the counts.

1
2
3
4
5
 stats: dict[tuple[int, int], int] = {}
    for ii, chunk in enumerate(chunks):
        stats = get_stats(chunk, stats, frequencies[ii])

conn.send(stats)

The second one: the main process needs to build the initial stats dictionary before starting the BPE algorithm:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def train_bpe():
    # shard data and initialize workers
    ... 
    stats: dict[tuple[int, int], int] = {}
    for (_ worker_conn) in workers.values():
        local_stats = worker_conn.recv()
        stats.reduce(local_stats) # Reduce here is just a shortened form for summation

    for i in range(num_merges):
        # Find the highest occurring pair and assign it a new index.
        pair = max(stats, key=stats.get)  # ty: ignore[no-matching-overload]
        new_id = 256 + i

        # Send the info to the workers to perform merge and diff count.
        for _, worker_pipe in workers.values():
            worker_pipe.send((pair, new_id))

        merges[pair] = new_id
        vocab[new_id] = vocab[pair[0]] + vocab[pair[1]]

        # Receive information from workers and update the global stats dict with the diffs.
        update_diffs(workers, stats)

The last change: The update_diffs method. This method handles updating the single stats dictionary with the diffs and removing pairs that reached 0, since if our global stats of a pair reached 0, it means that pair is no longer available in any text so we can safely remove it.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def update_diffs(workers, stats) -> None:
    """
    Updates the global stats tracking with the diffs only.
    """
    # Keep track on which pairs were changed.
    touched_pairs: set[tuple[int, int]] = set()
    # Updates the diffs for only changed values.
    for _, worker_pipe in workers.values():
        local_diffs = worker_pipe.recv()
        for pair, delta in local_diffs.items():
            stats[pair] = stats.get(pair, 0) + delta

    # Remove any pairs that have 0 value since that means they don't exist anywhere.
    for pair in touched_pairs:
        count = stats[pair]
        if count == 0:
            del stats[pair]
        elif count < 0:
            raise ValueError(f"Negative count for pair {pair}: {count}")

Time to run it:

1
2
3
4
5
6
7
8
time uv run --script bpe_multiprocess_diff.py --vocab-size 1257 --data-path part_00.txt part_01.txt --save-path mp_diff.json --num-workers 2
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00, 630.72it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00, 368.12it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1000/1000 [00:01<00:00, 550.40it/s]
________________________________________________________
Executed in    2.34 secs      fish           external
   usr time  836.54 millis    0.00 millis  836.54 millis
   sys time   40.59 millis    1.60 millis   38.99 millis

WOW WOW WOW!

We got an improvement to 2.3s thatโ€™s 4.5x improvement over the initial multiprocessing and roughly 4x improvement over the best single-process method! But can we do better?

Overcoming Serialization

One way to overcome serialization is stripping away the unnecessary information passed at every send and recv, since each operation needs to allocate Python objects for the dictionary and tuples in addition to the numbers themselves.
In practice we just need to know which two numbers are in the pair (and which precedes which), and that is the new index they are mapped to. We can do it by flattening the dictionary into 3 equal-length lists.
The first one: A list containing all the left side of the pairs.
The second one: A list containing all the right side of the pairs.
The third one: A list containing the mapping to a new index.

Unicode Table

Then we can use Pythonโ€™s Array module to pass information.

So all we need to do is change our send logic to:

1
2
3
4
5
6
records = array("Q")
for (left, right), count in stats.items():
    records.append(left)
    records.append(right)
    records.append(count)
conn.send_bytes(records)

And our receive and update logic to:

1
2
3
4
5
6
7
for _, worker_pipe in workers.values():
    records = array("Q")
    records.frombytes(worker_pipe.recv_bytes())
    for idx in range(0, len(records), 3):
        pair = (records[idx], records[idx + 1])
        count = records[idx + 2]
        stats[pair] = stats.get(pair, 0) + count

array("Q") creates a contiguous binary memory buffer, and it implements Pythonโ€™s buffer protocol. This means itโ€™s a typed array of unsigned integers, in our case we only need unsigned integers since our counts are always positive (we canโ€™t have a negative count of a pair), and send_bytes and recv_bytes work similar to send and recv respectively but over bytes-like objects.

1
2
3
4
5
6
7
8
time uv run --script bpe_multiprocess_array_diff.py --vocab-size 1257 --data-path part_00.txt part_01.txt --save-path mp_diff_array.json --num-workers 2
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00, 11.22it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00, 10.55it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1000/1000 [00:01<00:00, 596.42it/s]
________________________________________________________
Executed in    2.20 secs      fish           external
   usr time  697.38 millis    0.89 millis  696.50 millis
   sys time   53.19 millis    1.07 millis   52.12 millis

Not bad! 6% improvement, doesnโ€™t seem like much for this data but when dealing with hundreds of billions of text bytes 6% improvement can be hours of runtime.

One Last Improvement I promise!

The last improvement is not with the BPE itself, itโ€™s with the preprocessing. Right now we iterate over all the data to load it, then iterate again to split it and count the frequencies. We can combine both into a single loop to save some preprocessing time. This again helps when we talk about real-world data.

1
2
3
4
5
6
7
8
9
10
11
12
compiled_pattern = re.compile(SPLIT_PATTERN)
freqs = Counter()
if verbose:
    print(f"Worker: {worker_id}- 1. Reading shard and building freq dict.")
for file_ in tqdm(data_shard):
    with open(file_, "r") as fd:
        for line in tqdm(fd):
            split = compiled_pattern.findall(line.strip())
            freqs.update(split)

chunks = [list(chunk.encode("utf-8")) for chunk in freqs]
frequencies = list(freqs.values())
1
2
3
4
5
6
7
8
time uv run --script train_bpe_multiprocess.py --vocab-size 1257 --data-path part_00.txt part_01.txt --save-path mp_diff_array2.json --num-workers 2
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00, 11.69it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [00:00<00:00, 10.89it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1000/1000 [00:01<00:00, 621.26merges/s]
________________________________________________________
Executed in    2.09 secs      fish           external
   usr time  683.70 millis    1.62 millis  682.07 millis
   sys time   58.87 millis    0.05 millis   58.82 millis

Another 11% improvement.

Now everything begs the question, how fast is it really? So letโ€™s move on from our small data into some larger dataset. Weโ€™ll use the Tiny Stories train set, it contains 15 million lines that means 15 million stories ranging up to 5,500 characters. Weโ€™ll shard it into 8 parts and run with 8 workers.

Our approach

1
2
3
4
5
6
7
8
time uv run --script train_bpe_multiprocess.py --vocab-size 1257 --data-path test_part_00.txt test_part_01.txt test_part_02.txt test_part_03.txt test_part_04.txt test_part_05.txt test_part_06.txt test_part_07.txt --save-path freq.json --num-workers 8
1951632it [01:53, 17197.20it/s]                                                                                                                                                               | 0/1 [00:00<?, ?it/s]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1/1 [02:01<00:00, 121.33s/it]
100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 1000/1000 [00:22<00:00, 44.95merges/s]
________________________________________________________
Executed in  144.18 secs    fish           external
   usr time    3.53 secs    1.12 millis    3.53 secs
   sys time    0.13 secs    1.01 millis    0.13 secs

I wanted to compare myself to Hugging Faceโ€™s tokenizers library which is written in Rust, since it provides a BPE trainer.
But while running it my laptop ran out of RAM and almost died while it was still counting pairs.

This was the last time I managed to run the free command before my computer hung every 3 seconds.

1
2
3
4
free -h
               total        used        free      shared  buff/cache   available
Mem:            62Gi        47Gi       607Mi       604Mi        15Gi        14Gi
Swap:          126Gi        57Gi        68Gi

Unfortunate.

Closing notes

I think thereโ€™s a lot left to be desired, for example finding the max value requires O(n) search over stats, there are techniques1415 to use a max heap to reduce it to O(log n).

But I think the key takeaway for me is how simple and complex LLMs are. We saw an algorithm that was developed in 1994 being used today to fuel the biggest LLMs. But at the same time we saw the simplicity doesnโ€™t always win, thereโ€™s a lot of engineering effort that needs to be done to scale these algorithms to handle hundreds of tera-bytes, if not peta-bytes, of texts.

I think this was a really great learning experience, not just reading tutorials and implementing what was given but literally bashing my head on my desk with frustration while trying to find dumb ways to optimize it.

I also think that the best learned code and skills are learned via writing and not reading. James Clear talks about this in his Atomic Habits book16 about being in motion and taking action.
In this case being in motion is reading blogposts about BPE like this one, and taking action is writing the BPE code.

See in the next post!

References

  1. https://en.wikipedia.org/wiki/Unicodeย 

  2. https://www.reedbeta.com/blog/programmers-intro-to-unicode/ย 

  3. https://en.wikipedia.org/wiki/Tf%E2%80%93idfย 

  4. https://www.jmlr.org/papers/volume3/bengio03a/bengio03a.pdfย 

  5. https://arxiv.org/abs/1301.3781ย 

  6. https://www.merriam-webster.com/help/faq-how-many-english-wordsย 

  7. https://arxiv.org/abs/1607.04606ย 

  8. https://aclanthology.org/2025.acl-long.453/ย 

  9. https://aclanthology.org/2024.emnlp-main.1217/ย 

  10. https://diabetesjournals.org/care/article-abstract/17/2/152/17985/A-Mathematical-Model-for-the-Determination-of?redirectedFrom=fulltext/ย 

  11. https://github.com/huggingface/tokenizersย 

  12. https://www.youtube.com/watch?v=zduSFxRajkEย 

  13. https://arxiv.org/abs/2305.07759ย 

  14. https://aclanthology.org/2023.findings-acl.38/ย 

  15. https://jytan.net/blog/2025/bpe/ย 

  16. https://jamesclear.comย