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?

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!
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.
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:
- The final Unicode representation would be 3 bytes long.
- 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.
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?
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.
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.
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:
- Find the consecutive pair with the highest number of appearances.
- 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!
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.
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.
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:
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:
-
No specific word boundaries, we started learning Frankenstein combinations of words. If we try to encode something like
deep loglater, the tokenizer will look at the prefixdeep land see it will have an entry in the merges dictionary and will replace it with the token260while tokenizingoginto individual tokens. -
We donโt handle punctuation and spaces efficiently, we learn both a representation for
deep,deep.,.. -
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_statsof one chunk doesnโt depend on another chunk, only the summation at the end. Additionally,mergeof each chunk doesnโt depend on others This means we can parallelize those two operations. Forget_statsweโll use map-reduce, and for merge itโs just regular multiprocessing.
True story.
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.
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.
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
-
https://en.wikipedia.org/wiki/Unicodeย ↩
-
https://www.reedbeta.com/blog/programmers-intro-to-unicode/ย ↩
-
https://en.wikipedia.org/wiki/Tf%E2%80%93idfย ↩
-
https://www.jmlr.org/papers/volume3/bengio03a/bengio03a.pdfย ↩
-
https://arxiv.org/abs/1301.3781ย ↩
-
https://www.merriam-webster.com/help/faq-how-many-english-wordsย ↩
-
https://arxiv.org/abs/1607.04606ย ↩
-
https://aclanthology.org/2025.acl-long.453/ย ↩
-
https://aclanthology.org/2024.emnlp-main.1217/ย ↩
-
https://diabetesjournals.org/care/article-abstract/17/2/152/17985/A-Mathematical-Model-for-the-Determination-of?redirectedFrom=fulltext/ย ↩
-
https://github.com/huggingface/tokenizersย ↩
-
https://www.youtube.com/watch?v=zduSFxRajkEย ↩
-
https://arxiv.org/abs/2305.07759ย ↩
-
https://aclanthology.org/2023.findings-acl.38/ย ↩
-
https://jytan.net/blog/2025/bpe/ย ↩
-
https://jamesclear.comย ↩