Byte Pair Encoding in Python: From Naive to Fast
Build a BPE tokenizer trainer from scratch, then make it ~430× faster across three optimizations — each removing a different bottleneck.
Byte Pair Encoding (BPE) is the subword-tokenization algorithm used to build a vocabulary for LLM training and is behind GPT-2, GPT-4, Llama, and most modern tokenizers. The idea is simple: repeatedly find the most frequent adjacent pair of symbols and merge it into a new symbol, until you have the vocabulary size you want. After you obtain a vocabulary, you convert any given text to tokens (int or int16, depending on the size of the vocabulary) and give these to an LLM to generate the next token.
This post builds a BPE trainer and tokenizer in Python, then optimizes them to make them faster. Each optimization step removes one specific bottleneck and is verifiably faster, while producing exactly the same merges as the step before.
The algorithm, in one paragraph
Training
Start with a corpus of text and then split it into sentences on special tokens, e.g., <|begin_of_text|>.
Most LLMs have multiple special tokens such as <|im_start|>, <tool_response>, etc.
A list of special tokens can be found in the tokenizer_config.json file in HuggingFace models.
Once you have divided the corpus into sentences, we run a tokenization step, also called pre-tokenization, which converts the sentence into words/tokens. This step uses a regex expression similar to this:
"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
which splits the sentence ‘The cat is cute’ into “The”, ” cat”, ” is”, ” cute”. Notice that the spaces ended up in the token itself.
Here is a real example: the tokenizer_config.json for Qwen3.5-397B-A17B, which contains the special-token vocabulary, and where pretokenize_regex is the pre-tokenization pattern shown above:
Qwen3.5-397B-A17B / tokenizer_config.jsonJSON
{
"add_prefix_space": false,
"added_tokens_decoder": {
"248044": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248045": {
"content": "<|im_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248046": {
"content": "<|im_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248047": {
"content": "<|object_ref_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248048": {
"content": "<|object_ref_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248049": {
"content": "<|box_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248050": {
"content": "<|box_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248051": {
"content": "<|quad_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248052": {
"content": "<|quad_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248053": {
"content": "<|vision_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248054": {
"content": "<|vision_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248055": {
"content": "<|vision_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248056": {
"content": "<|image_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248057": {
"content": "<|video_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248058": {
"content": "<tool_call>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"248059": {
"content": "</tool_call>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"248060": {
"content": "<|fim_prefix|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"248061": {
"content": "<|fim_middle|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"248062": {
"content": "<|fim_suffix|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"248063": {
"content": "<|fim_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"248064": {
"content": "<|repo_name|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"248065": {
"content": "<|file_sep|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"248066": {
"content": "<tool_response>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"248067": {
"content": "</tool_response>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"248068": {
"content": "<think>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"248069": {
"content": "</think>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"248070": {
"content": "<|audio_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248071": {
"content": "<|audio_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248072": {
"content": "<tts_pad>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248073": {
"content": "<tts_text_bos>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248074": {
"content": "<tts_text_eod>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248075": {
"content": "<tts_text_bos_single>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"248076": {
"content": "<|audio_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
}
},
"additional_special_tokens": [
"<|im_start|>",
"<|im_end|>",
"<|object_ref_start|>",
"<|object_ref_end|>",
"<|box_start|>",
"<|box_end|>",
"<|quad_start|>",
"<|quad_end|>",
"<|vision_start|>",
"<|vision_end|>",
"<|vision_pad|>",
"<|image_pad|>",
"<|video_pad|>"
],
"bos_token": null,
"chat_template": "{%- set image_count = namespace(value=0) %}\n{%- set video_count = namespace(value=0) %}\n{%- macro render_content(content, do_vision_count, is_system_content=false) %}\n {%- if content is string %}\n {{- content }}\n {%- elif content is iterable and content is not mapping %}\n {%- for item in content %}\n {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}\n {%- if is_system_content %}\n {{- raise_exception('System message cannot contain images.') }}\n {%- endif %}\n {%- if do_vision_count %}\n {%- set image_count.value = image_count.value + 1 %}\n {%- endif %}\n {%- if add_vision_id %}\n {{- 'Picture ' ~ image_count.value ~ ': ' }}\n {%- endif %}\n {{- '<|vision_start|><|image_pad|><|vision_end|>' }}\n {%- elif 'video' in item or item.type == 'video' %}\n {%- if is_system_content %}\n {{- raise_exception('System message cannot contain videos.') }}\n {%- endif %}\n {%- if do_vision_count %}\n {%- set video_count.value = video_count.value + 1 %}\n {%- endif %}\n {%- if add_vision_id %}\n {{- 'Video ' ~ video_count.value ~ ': ' }}\n {%- endif %}\n {{- '<|vision_start|><|video_pad|><|vision_end|>' }}\n {%- elif 'text' in item %}\n {{- item.text }}\n {%- else %}\n {{- raise_exception('Unexpected item type in content.') }}\n {%- endif %}\n {%- endfor %}\n {%- elif content is none or content is undefined %}\n {{- '' }}\n {%- else %}\n {{- raise_exception('Unexpected content type.') }}\n {%- endif %}\n{%- endmacro %}\n{%- if not messages %}\n {{- raise_exception('No messages provided.') }}\n{%- endif %}\n{%- if tools and tools is iterable and tools is not mapping %}\n {{- '<|im_start|>system\\n' }}\n {{- \"# Tools\\n\\nYou have access to the following functions:\\n\\n<tools>\" }}\n {%- for tool in tools %}\n {{- \"\\n\" }}\n {{- tool | tojson }}\n {%- endfor %}\n {{- \"\\n</tools>\" }}\n {{- '\\n\\nIf you choose to call a function ONLY reply in the following format with NO suffix:\\n\\n<tool_call>\\n<function=example_function_name>\\n<parameter=example_parameter_1>\\nvalue_1\\n</parameter>\\n<parameter=example_parameter_2>\\nThis is the value for the second parameter\\nthat can span\\nmultiple lines\\n</parameter>\\n</function>\\n</tool_call>\\n\\n<IMPORTANT>\\nReminder:\\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\\n- Required parameters MUST be specified\\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\\n</IMPORTANT>' }}\n {%- if messages[0].role == 'system' %}\n {%- set content = render_content(messages[0].content, false, true)|trim %}\n {%- if content %}\n {{- '\\n\\n' + content }}\n {%- endif %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n{%- else %}\n {%- if messages[0].role == 'system' %}\n {%- set content = render_content(messages[0].content, false, true)|trim %}\n {{- '<|im_start|>system\\n' + content + '<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n {%- set index = (messages|length - 1) - loop.index0 %}\n {%- if ns.multi_step_tool and message.role == \"user\" %}\n {%- set content = render_content(message.content, false)|trim %}\n {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}\n {%- set ns.multi_step_tool = false %}\n {%- set ns.last_query_index = index %}\n {%- endif %}\n {%- endif %}\n{%- endfor %}\n{%- if ns.multi_step_tool %}\n {{- raise_exception('No user query found in messages.') }}\n{%- endif %}\n{%- for message in messages %}\n {%- set content = render_content(message.content, true)|trim %}\n {%- if message.role == \"system\" %}\n {%- if not loop.first %}\n {{- raise_exception('System message must be at the beginning.') }}\n {%- endif %}\n {%- elif message.role == \"user\" %}\n {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n {%- elif message.role == \"assistant\" %}\n {%- set reasoning_content = '' %}\n {%- if message.reasoning_content is string %}\n {%- set reasoning_content = message.reasoning_content %}\n {%- else %}\n {%- if '</think>' in content %}\n {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n {%- endif %}\n {%- endif %}\n {%- set reasoning_content = reasoning_content|trim %}\n {%- if loop.index0 > ns.last_query_index %}\n {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content + '\\n</think>\\n\\n' + content }}\n {%- else %}\n {{- '<|im_start|>' + message.role + '\\n' + content }}\n {%- endif %}\n {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}\n {%- for tool_call in message.tool_calls %}\n {%- if tool_call.function is defined %}\n {%- set tool_call = tool_call.function %}\n {%- endif %}\n {%- if loop.first %}\n {%- if content|trim %}\n {{- '\\n\\n<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n {%- else %}\n {{- '<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n {%- endif %}\n {%- else %}\n {{- '\\n<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n {%- endif %}\n {%- if tool_call.arguments is defined %}\n {%- for args_name, args_value in tool_call.arguments|items %}\n {{- '<parameter=' + args_name + '>\\n' }}\n {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}\n {{- args_value }}\n {{- '\\n</parameter>\\n' }}\n {%- endfor %}\n {%- endif %}\n {{- '</function>\\n</tool_call>' }}\n {%- endfor %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"tool\" %}\n {%- if loop.previtem and loop.previtem.role != \"tool\" %}\n {{- '<|im_start|>user' }}\n {%- endif %}\n {{- '\\n<tool_response>\\n' }}\n {{- content }}\n {{- '\\n</tool_response>' }}\n {%- if not loop.last and loop.nextitem.role != \"tool\" %}\n {{- '<|im_end|>\\n' }}\n {%- elif loop.last %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n {%- else %}\n {{- raise_exception('Unexpected message role.') }}\n {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n {{- '<|im_start|>assistant\\n' }}\n {%- if enable_thinking is defined and enable_thinking is false %}\n {{- '<think>\\n\\n</think>\\n\\n' }}\n {%- else %}\n {{- '<think>\\n' }}\n {%- endif %}\n{%- endif %}",
"clean_up_tokenization_spaces": false,
"eos_token": "<|im_end|>",
"errors": "replace",
"model_max_length": 262144,
"pad_token": "<|endoftext|>",
"split_special_tokens": false,
"tokenizer_class": "Qwen2Tokenizer",
"unk_token": null,
"add_bos_token": false,
"pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
"extra_special_tokens": {
"audio_bos_token": "<|audio_start|>",
"audio_eos_token": "<|audio_end|>",
"audio_token": "<|audio_pad|>",
"image_token": "<|image_pad|>",
"video_token": "<|video_pad|>",
"vision_bos_token": "<|vision_start|>",
"vision_eos_token": "<|vision_end|>"
}
}
After we have these tokens, we encode them to utf-8 to get a byte array bytes. Now each word is a sequence of bytes. We count each pair
(a,b) across all the words. We pick the most frequent pair, breaking ties by choosing the lexicographically largest of the tied pairs. We begin with
a vocabulary of 256 entries, one for each byte value (0-255), and assign a new index to the most frequent pair, then replace all its occurrences with that index. So, for example,
if (a, b) is the most frequent pair, then (97, 98) becomes 256, and all occurrences of (a, b) will be replaced by (256). Now we repeat this loop until we reach a
certain number of merges or vocabulary size. The resulting dictionary dict[int, bytes] is the vocabulary we use to tokenize the text for the LLM.
More formally, at each round we choose:
and then update the corpus so that every non-overlapping, left-to-right
occurrence of (a,b) collapses into a fresh token id (256, 257, 258, …)
Tokenizer
Once we have this vocabulary and these merges, we tokenize the text by running the pre-tokenization step and, for each token, applying the merges from
the merge list (in the order they were learned) to merge the bytes. We keep doing this until no more merges apply. We then use the
vocabulary with the final list[bytes], convert each byte to an int/int16, and this becomes the input to the LLM.
Dataset
We are going to use two datasets to measure performance.
Shared setup
Every tokenizer training starts with dividing the data into readable chunks and building the pair count. This reading can easily be done across multiple threads, since we are just adding to the counter after a chunk is processed. The code below takes a chunk of text and builds a word counter. The snippet below produces word counts over the entire dataset.
Pretokenization time
- Tiny Stories (8 cores, 128 chunks): 64.2 secs average time over 3 runs
- Open web text (8 cores, 128 chunks) : 356.44 secs average time over 3 runs
The timings are not 100% accurate, since the number of runs is small, but they are indicative of the expected time. The cores and chunks were chosen to avoid out-of-memory errors and to use the maximum number of CPU cores.
Pretokenization helperspython
def split_into_pretokens(chunk: str, special_tokens: list[str]):
"""
Yield GPT-2 pretokens from `chunk`.
First splits the chunk on any `special_tokens` (so they are kept intact and
never become part of a merge), then applies the GPT-2 pretokenization regex to
each segment between special tokens. Tokens are yielded lazily.
"""
split_pattern = "|".join(re.escape(token) for token in special_tokens)
split_re = re.compile(split_pattern)
gpt_pretoken_re = re.compile(PRETOKENIZE_PATTERN)
position = 0
for match in split_re.finditer(chunk):
segment = chunk[position:match.start()]
position = match.end()
for token_match in gpt_pretoken_re.finditer(segment):
yield token_match.group(0)
# Tail after the last special token.
for token_match in gpt_pretoken_re.finditer(chunk[position:]):
yield token_match.group(0)
def count_pretokens_in_chunk(byte_range: tuple[int, int], input_path: str, special_tokens: list[str]) -> Counter:
"""
Read a byte range of the input file and return a Counter of UTF-8 encoded
pretokens. One unit of parallel work during pretokenization.
"""
start, end = byte_range
pretoken_counts: Counter = Counter()
with open(input_path, "rb") as f:
f.seek(start)
chunk = f.read(end - start).decode("utf-8", errors="ignore")
for token in split_into_pretokens(chunk, special_tokens):
pretoken_counts[token.encode("utf-8")] += 1
del chunk # Free the decoded chunk early.
return pretoken_countsAfter this we build a pair count using word count.
Building the pair countpython
for key, frequency in pretoken_counts.items():
for pair in zip(key, key[1:]):
pair_count[pair] += frequencyAttempt 1 — Naive baseline
We find the pair with the maximum count, breaking ties lexicographically. We merge it and replace it with a new vocab index. Then we recalculate the pair counts for all the words.
This will give us O(M * num_merges) where M=N*L with N = token count and L = average word length.
train_bpe_naive (naive baseline)python
def train_bpe_naive(
input_path: str | os.PathLike,
vocab_size: int,
special_tokens: list[str],
word_frequencies: list[tuple[list[int], int]],
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
"""
Simplest BPE trainer.
Every round we recompute every adjacent-pair frequency from scratch and then
re-scan every word to apply the chosen merge.
"""
# Vocab starts as all 256 single bytes, followed by the special tokens.
vocab: dict[int, bytes] = {i: bytes([i]) for i in range(256)}
next_token_id = 256
for special_token in special_tokens:
vocab[next_token_id] = special_token.encode("utf-8")
next_token_id += 1
merges: list[tuple[bytes, bytes]] = []
progress_bar = tqdm(range(vocab_size - len(vocab)), desc="Training BPE (naive)")
for _ in progress_bar:
if len(vocab) >= vocab_size:
break
# Recompute pair frequencies from scratch each iteration.
pair_count: dict[tuple[int, int], int] = defaultdict(int)
for word, count in word_frequencies:
for pair in zip(word, word[1:]):
pair_count[pair] += count
# Pick the most frequent pair, breaking ties lexicographically by the
# bytes each token id maps to.
pair, _ = max(
pair_count.items(),
key=lambda kv: (kv[1], tuple(vocab[b] for b in kv[0])),
)
# Register the new merged token.
bytes_a = vocab[pair[0]]
bytes_b = vocab[pair[1]]
new_token = bytes_a + bytes_b
next_id = next_token_id # id assigned to this new token
vocab[next_id] = new_token
next_token_id += 1
merges.append((bytes_a, bytes_b))
# Replace the merged pair inside every word (two-pointer in-place rewrite).
for word, _count in word_frequencies:
read_idx = 0
write_idx = 0
while read_idx < len(word):
match = (
read_idx + 1 < len(word)
and word[read_idx] == pair[0]
and word[read_idx + 1] == pair[1]
)
if match:
word[write_idx] = next_id
read_idx += 2
else:
word[write_idx] = word[read_idx]
read_idx += 1
write_idx += 1
del word[write_idx:]
return vocab, mergesNote the lambda function key=lambda kv: (kv[1], tuple(vocab[b] for b in kv[0])) — we also include the bytes (as a tuple) in the key so that ties are broken lexicographically.
Training time (naive)
| Dataset | Vocab size | Time |
|---|---|---|
| Tiny Stories | 10000 | 811 secs |
| Open web text | 32000 | 119.12 hours |
I did not run the complete iteration for Open web text. The numbers mentioned here are just extrapolated from 13.54 sec per iteration.
Attempt 2
After running the first iteration, one thing is very obvious: the most frequent pair does not appear in every word/token.
So one optimization we can apply is to maintain a dictionary mapping each pair to the tokens that contain it, so that we update and modify only those words.
Even after this optimization, we are still rebuilding pair counts in every iteration. To fix this, we can keep track of pairs whose count changes. For
example, if a token has a frequency of 20 and contains (1,2,3), and we want to merge (1,2), we can subtract 20 from the counts of pairs (1,2) and (2,3) and increment (4,3) by 20.
This way, in the end we will only update the pair_count entries that changed and keep reusing the same pair_count.
Incremental trainer (Attempt 2)python
def apply_merge_and_track_deltas(
word_frequencies: list[tuple[list[int], int]],
affected_word_indices: list[int],
pair_to_merge: tuple[int, int],
new_token_id: int,
pair_to_word_indices: dict,
) -> dict[tuple[int, int], int]:
"""
Merge `pair_to_merge` into `new_token_id` inside every affected word IN PLACE,
while computing how the frequency of every adjacent pair changes.
Returns an aggregate `pair_deltas` mapping each changed pair to its net
frequency change, which the caller folds back into its global structures.
Side effects:
- Each affected word in `word_frequencies` is rewritten in place.
- `pair_to_word_indices` (the inverted index) is updated so it stays accurate
for future merges: words are added when a pair newly appears in them and
removed when a pair no longer appears in them.
"""
first_token, second_token = pair_to_merge
pair_deltas: dict[tuple[int, int], int] = defaultdict(int)
for word_idx in affected_word_indices:
word, count = word_frequencies[word_idx]
local_deltas: dict[tuple[int, int], int] = defaultdict(int)
# Two-pointer in-place rewrite: read_idx scans the original word while
# write_idx places tokens (possibly merged) back into the same list.
read_idx = 0
write_idx = 0
while read_idx < len(word):
is_match = (
read_idx + 1 < len(word)
and word[read_idx] == first_token
and word[read_idx + 1] == second_token
)
if is_match:
# The merged pair itself loses `count` occurrences.
local_deltas[(word[read_idx], word[read_idx + 1])] -= count
# Left neighbor: pair (prev, first) is replaced by (prev, new).
if write_idx > 0:
local_deltas[(word[write_idx - 1], word[read_idx])] -= count
local_deltas[(word[write_idx - 1], new_token_id)] += count
# Right neighbor: pair (second, next) is replaced by (new, next).
if read_idx + 2 < len(word):
local_deltas[(new_token_id, word[read_idx + 2])] += count
local_deltas[(word[read_idx + 1], word[read_idx + 2])] -= count
word[write_idx] = new_token_id
read_idx += 2
else:
word[write_idx] = word[read_idx]
read_idx += 1
write_idx += 1
# Truncate the tail left over by the now-shorter, merged word.
del word[write_idx:]
# Fold local deltas into the aggregate and keep the inverted index current.
for pair, delta in local_deltas.items():
if delta == 0:
continue
pair_deltas[pair] += delta
if pair_deltas[pair] == 0:
del pair_deltas[pair]
if delta < 0:
# This pair lost occurrences in this word; drop the word from the
# index only if the pair no longer occurs in it at all.
still_present = any(adjacent == pair for adjacent in zip(word, word[1:]))
if not still_present:
pair_to_word_indices[pair].discard(word_idx)
else:
# This pair gained occurrences in this word; record the word.
if pair not in pair_to_word_indices:
pair_to_word_indices[pair] = set()
pair_to_word_indices[pair].add(word_idx)
return pair_deltas
def train_bpe_incremental(
input_path: str | os.PathLike,
vocab_size: int,
special_tokens: list[str],
word_frequencies: list[tuple[list[int], int]],
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
"""
Incremental BPE trainer.
Same selection rule as the naive version, but instead of recomputing all pair
frequencies each round we maintain a persistent `pair_count` map plus an
inverted index `pair_to_word_indices`. After a merge, only the words that
actually contained the merged pair are re-scanned, and only the changed pair
frequencies are updated.
"""
vocab: dict[int, bytes] = {i: bytes([i]) for i in range(256)}
next_token_id = 256
for special_token in special_tokens:
vocab[next_token_id] = special_token.encode("utf-8")
next_token_id += 1
merges: list[tuple[bytes, bytes]] = []
# Build the initial pair frequencies and the inverted index (pair -> word indices).
pair_to_word_indices: dict = defaultdict(set)
pair_count: dict[tuple[int, int], int] = defaultdict(int)
for word_idx, (word, count) in enumerate(word_frequencies):
for pair in zip(word, word[1:]):
pair_count[pair] += count
pair_to_word_indices[pair].add(word_idx)
progress_bar = tqdm(range(vocab_size - len(vocab)), desc="Training BPE (incremental)")
for _ in progress_bar:
if len(pair_count) <= 0:
break
if len(vocab) >= vocab_size:
break
pair, _ = max(
pair_count.items(),
key=lambda kv: (kv[1], tuple(vocab[b] for b in kv[0])),
)
# Register the new merged token.
bytes_a = vocab[pair[0]]
bytes_b = vocab[pair[1]]
new_token = bytes_a + bytes_b
next_id = next_token_id
vocab[next_id] = new_token
next_token_id += 1
merges.append((bytes_a, bytes_b))
# Only touch the words that currently contain this pair.
affected_word_indices = list(pair_to_word_indices[pair])
pair_deltas = apply_merge_and_track_deltas(
word_frequencies, affected_word_indices, pair, next_id, pair_to_word_indices
)
# Fold the frequency deltas back into pair_count.
for changed_pair, delta in pair_deltas.items():
pair_count[changed_pair] += delta
if pair_count[changed_pair] == 0:
del pair_count[changed_pair]
return vocab, mergesTraining time (incremental)
| Dataset | Vocab size | Time |
|---|---|---|
| Tiny Stories | 10000 | 158.16 secs |
| Open web text | 32000 | 6.16 hours |
This is a big improvement over the previous attempt, but I don’t want to spend this much time building a vocab, so onto the next iteration.
Attempt 3
The final bottleneck is max(counts, ...) — an O(P)
scan over all the pairs in pair_counts, so maybe some sort of data structure could help.
In my search I came across sorteddict from sortedcontainers.
Final implementation
train_bpe_optimized (Attempt 3)python
def train_bpe_optimized(
input_path: str | os.PathLike,
vocab_size: int,
special_tokens: list[str],
word_frequencies: list[tuple[list[int], int]],
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
"""
Fully optimized BPE trainer.
Builds on the incremental version but replaces the O(pairs) max() scan with a
`count_to_pairs` SortedDict keyed by frequency. The largest frequency bucket is
the last key, so finding the most frequent pair (and then its lexicographic
tie-break within that bucket) no longer requires scanning every pair each round.
"""
vocab: dict[int, bytes] = {i: bytes([i]) for i in range(256)}
next_token_id = 256
for special_token in special_tokens:
vocab[next_token_id] = special_token.encode("utf-8")
next_token_id += 1
merges: list[tuple[bytes, bytes]] = []
# pair_count: pair -> count
pair_count: dict[tuple[int, int], int] = defaultdict(int)
# pair_to_words: pair -> [word indices] (built first, then converted to sets).
pair_to_words: dict = defaultdict(list)
for word_idx, (word, count) in enumerate(word_frequencies):
for pair in zip(word, word[1:]):
pair_to_words[pair].append(word_idx)
pair_count[pair] += count
# Inverted index with set-valued entries for O(1) add/discard.
pair_to_word_indices = {pair: set(indices) for pair, indices in pair_to_words.items()}
# Bucket pairs by frequency so the maximum frequency is the last key.
count_to_pairs: dict = defaultdict(set)
for pair, count in pair_count.items():
count_to_pairs[count].add(pair)
count_to_pairs = SortedDict(count_to_pairs)
progress_bar = tqdm(range(vocab_size - len(vocab)), desc="Training BPE (optimized)")
for _ in progress_bar:
if len(pair_count) <= 0:
break
if len(vocab) >= vocab_size:
break
# All pairs that share the current maximum frequency.
pairs_at_max_count = count_to_pairs.values()[-1]
# Tie-break lexicographically (by the bytes each token id decodes to).
pair = max(pairs_at_max_count, key=lambda v: tuple(vocab[b] for b in v))
# Register the new merged token.
bytes_a = vocab[pair[0]]
bytes_b = vocab[pair[1]]
new_token = bytes_a + bytes_b
next_id = next_token_id
vocab[next_id] = new_token
next_token_id += 1
merges.append((bytes_a, bytes_b))
# Only re-scan the words that currently contain the merged pair.
affected_word_indices = list(pair_to_word_indices[pair])
pair_deltas = apply_merge_and_track_deltas(
word_frequencies, affected_word_indices, pair, next_id, pair_to_word_indices
)
# Reconcile `pair_deltas` with both `pair_count` and `count_to_pairs`.
for changed_pair, delta in pair_deltas.items():
if delta > 0:
# A positive delta only happens for pairs involving the brand-new
# token id, so this pair is necessarily brand new.
pair_count[changed_pair] = delta
if delta not in count_to_pairs:
count_to_pairs[delta] = set()
count_to_pairs[delta].add(changed_pair)
continue
old_count = pair_count[changed_pair]
new_count = old_count + delta
if new_count < 0:
raise ValueError(f"Pair doesn't exist in pair_count")
# Move the pair out of its old frequency bucket.
count_to_pairs[old_count].discard(changed_pair)
# Drop the bucket entirely if it is now empty.
if not count_to_pairs[old_count]:
count_to_pairs.pop(old_count)
# Move the pair into its new frequency bucket (if it still occurs).
if new_count > 0:
if new_count not in count_to_pairs:
count_to_pairs[new_count] = set()
count_to_pairs[new_count].add(changed_pair)
# Keep pair_count in sync.
if new_count > 0:
pair_count[changed_pair] = new_count
else:
pair_count.pop(changed_pair)
return vocab, mergesTraining time (optimized)
| Dataset | Vocab size | Time |
|---|---|---|
| Tiny Stories | 10000 | 1.89 secs |
| Open web text | 32000 | 242.21 secs |
The full picture
| Attempt | Technique | Time | Speedup | Bottleneck removed |
|---|---|---|---|---|
| 1 | Naive recount + rescan every word | 811 s | 1× | — |
| 2 | Incremental counts + inverted index | 158.16 s | ~5× | Recomputing every pair count each round, and scanning every word per merge |
| 3 | Bucketed max selection (SortedDict) | 1.89 s | ~430× | Linear scan over all pairs to find the most frequent |
Times are for Tiny Stories at a 10,000-token vocabulary; Open Web Text (32,000-vocab) follows the same curve — 119 h → 6.2 h → 242 s.
Full source code
The complete, runnable trainer — pretokenization, all three train_bpe_*
variants, and the shared merge helper — lives in a single GitLab snippet:
BPE Tokenizer Trainer — full implementation
Where to go next
- Encode text with the vocab — training is only half the job; the other
half is applying the learned merges to new text. You can find a tokenizer which takes any text corpus and apply the learned
vocab and merges and outputs
list[int]. BPE Tokenizer — full implementation - Move to C/C++ — I think it would benefit immensely if we move the entire code to C/C++ and use Re-flex
for regex and using data oriented design for data structures with nanobind for python to C/C++. Python will just send the file paths for vocab building
and dataset conversion. C/C++ will just return a
list[int]or vocab/merges store path. This way we don’t have to worry about data ownership and it provides clean boundaries. Input will always be files and output either files orlist[int].
The leap from 13½ minutes to under two seconds didn’t come from one clever trick — it came from removing three distinct bottlenecks, one at a time.