Andrej Karpathy’s code and videos are so awesome! I love how concise his code is. His minGPT repo was an important inspiration for my “The Art of Transformer Programming” book.

Anyway, I just finished watching his BPE video and got inspired again. Karpathy’s code is really nice (as always), and his implementation is obviously not meant to be fast, but, we should be able to make the code much faster without sacrificing clarity too much. Hopefully, faster code can make experimentation (e.g. with different scores, instead of always taking the most popular pair) easier.

You can find my late night take on a minimal, clean, and fast BPE implementation here.

For his tests, Andrej used a snapshot of the Wikipedia article on Taylor Swift with ~185K chars. On my laptop, training/tokenizing this file with a vocab size of 10K takes:

So, in this setup, we get 100X faster training.

Training fast_minbpe on the same text but with a GPT-4-sized vocab of 100K tokens takes only slightly longer at 1.61 secs, but results in a single token. Training fast_minbpe on a GPT-4-sized 100K vocab on the English translation of Marcel Proust’s “Swann’s Way”, which is the first volume of the world’s longest novel, (you can find the file, which contains just over 1 million bytes here) takes just 8.15 seconds.

This was a really fun puzzle and, as usual, I recommend trying to solve this yourself!

Below I detail the main ideas behind my implementation, and you can find the code on github.

You really should watch Andrej’s video (or at least make sure you are familiar with the byte pair encoding algorithm) before reading on.

Why is minbpe slow?

If we are to perform $N$ merges, and the length of the training text is $L$, Karpathy’s original impl does (I think):

for i in range(N):
  calc_stats()    # O(L)
  find_max()      # O(L)
  do_merges()     # O(L)

For a total complexity of $O(NL)$ (maybe I’m neglecting some factors).

Why is fast_minbpe fast?

So fast_minbpe does something like this:

stats = calc_stats()            # O(L)
for i in range(N):
  find_max()                    # O(log(L))
  do_merges_and_update_stats()  # O(M_i)

Where $M_i$ denotes the actual number of merges we perform at the $i$th iteration. Note that $M_1 + M_2 + \dots + M_n \le L - 1$, so the overall complexity of everything (again neglecting logarithmic factors) is $O(L)$!

Doing the wrong thing

The “right thing” to do here would have been to profile Andrej’s code and measure which parts take time before attempting to optimize anything. In Python btw, this can be done with a few of lines and the builtin cProfile module:

import cProfile
import pstats
from pstats import SortKey
cProfile.run('func_to_profile()', 'pstats')
p = pstats.Stats('pstats')
p.strip_dirs().sort_stats(SortKey.CUMULATIVE).print_stats(10)
p.strip_dirs().sort_stats(SortKey.TIME).print_stats(10)

Since I did this just for fun, as a quick late night exercise though, I did not do that at all. Instead, I made a quick draft of the above complexity calculations and designed a couple of data structures implementing them. I wish I had run profiling before starting (for one, I would have had more cool numbers to report). In my original implementation I actually kept the cute find_max() one-liner with the $O(L)$ complexity, and got only a 10X speedup vs minbpe, but it was not until I made the find_max() operation take logarithmic time that I got the full improvement factors above. Tl;dr - when optimizing code you should really use a profiler, but, when dealing with large objects, complexity is never wrong.

Another very wrong thing that I did is that I basically didn’t write any tests… My code produces identical outputs to Karpathy’s code (when I force both to break ties in the same way). Again - this was just a late night fun exercise :)

Laziness Pays Off

Before we dive into the implementation, I want to make a general note on lazy data structures. By saying that a data structure is lazy I mostly mean that it only does work when it is actually needed, not ahead of time. Laziness is very often a good strategy for data structures, and will indeed turn out beneficial for both the data structures we’ll use in fast_minbpe. As my good friend Danny Lumen likes to say:

“When dealing with data structures, being lazy is often the best strategy”
- Danny Lumen

Ok, let’s consider the main data structures that I used.

IndexedList

When merging, we need to delete items from the middle of our list. A simple linked list makes this efficient. We also need to know what to iterate on - i.e. we need to efficiently find all nodes containing (the first element of) a given pair. To handle that, we’ll hold a simple python dict (the index) mapping each pair to a plain python list containing the nodes from our linked list that contain the first element of the pair. Unfortunately, updating this list is hard and inefficient. Luckily, it costs us the same to update the list and to just check that it’s still up to date when accessing it, so we’ll opt for the latter solution. Yep - being lazy pays off! To remind ourselves that the index could be stale, we name it stale_index.

You can find the implementation of the IndexedList here.

Multiset

I also needed a multiset class to hold the stats. The builtin collections.Counter class is exactly the API that I needed, but its most_common function running in linear time is too slow for our purposes. Instead, I wanted to hold all the elements in a version of a max-heap that supports increasing/decreasing counts of internal elements.

You can find my implementation of Multiset here.

Note that a tiny change here made a huge difference in performance - specifically, breaking count ties explicitly i.e. doing something like this:

class Multiset:
  class Node:
    def __lt__(self, other):
      return (self.count, self.val, self.pos) < (other.count, other.val, other.pos)

is drastically slower than breaking ties lazily like so:

class Multiset:
  class Node:
    def __lt__(self, other):
      return self.count < other.count

as the former requires the heap to update much more (and less importantly, creating and comparing tuples is slow). This is a different type of laziness as the one I mentioned above, but it pays off as well :)

Another important change (the right kind of lazy!) was accumulating modifications and only updating the heap once a query is performed (laziness FTW!).

Putting it all together

Armed with an IndexedList and our Multiset, it’s trivial to implement the merge function, at the center of the BPE algorithm:

def merge(pair, new_id, indexed_list: IndexedList, stats:Multiset=None):
  for node in indexed_list.stale_index[pair]:
    if node.val != pair[0] or node.next is None or node.next.val != pair[1]:
      continue  # The index was stale - continue.
    # Say we're merging "bc" to "X" in "abcd", and the node we're visiting now is "b".
    if stats is not None:  # Update the stats.
      stats.remove(pair)  # Remove "bc".
      if node.next.next is not None:
        stats.remove((node.next.val, node.next.next.val))  # Remove "cd".
        stats.add((new_id, node.next.next.val))  # Add "Xd".
      if node.prev is not None:
        stats.remove((node.prev.val, pair[0]))  # Remove "ab".
        stats.add((node.prev.val, new_id))  # Add "aX".
    node.next.delete()  # Delete "c", we now have "abd".
    node.val = new_id  # Update "b" to "X", we now have "aXd".
    indexed_list.update_index(node)  # Add "aX" and "Xd" to the index.

Note that importantly, we are careful to skip a pair from the index if it no longer holds the desired value.

You can find my implementation here, but, as always, I recommend attempting this yourself first!

Alternatives

If you’re only interested in fast_minbpe.py feel free to skip everything below this point!

Heapy, HeapHeapHoory, and HeapyKiYay

I actually implemented the Multiset 3 times: the first version (Heapy) was fast, but a bit convoluted, the second version (HeapHeapHooray) was tiny but slower, and final implementation (HeapyKiYay) is the one I ended up using.

Version 1 - Heapy

Since I needed each element of the heap to be aware of its own position (so I could maintain the heap condition when it increases/decreases), my first attempt consisted of writing a less standard pointer-based implementation (heaps are almost always implemented in an array). You can find this impl here.

Version 2 - HeapHeapHooray

Here I decided to optimize for code length. To make this super short, I wanted to reuse the standard implementation from the heapq module. The main idea was to create a custom list object that tells its elements where they are (so I called it a GPSList…).

Here’s the full implementation:

import heapq

class GPSList(list):
  # A list-like class that tells its elements where they are:
  # >>> l = GPSList(lambda i, x: print(f'{x} is in position {i}'))
  # >>> l.append(3)
  #   3 is in position 0
  # >>> l[0] = 5
  #   5 in in position 0
  # We only support the subset of list operations that are needed for our heap.
  def __init__(self, update_fn):
    super().__init__()
    self.update_fn = update_fn

  def __setitem__(self, i, x):
    super().__setitem__(i, x)
    self.update_fn(i, x)

  def append(self, x):
    super().append(x)
    self.update_fn(len(self) - 1, x)

  # Forbid other methods that can add elements - hope I didn't forget anything :)
  extend = insert = remove = __iadd__ = None

class HeapHeapHooray:
  def __init__(self):
    def setpos(i, x):
      x[-1] = i  # Our elements are lists of length 3 (count, pair, pos).
    self.l = GPSList(setpos)

  @property
  def max(self):
    return self.l[0][1]

  def increase(self, x):
    x[0] += 1
    heapq._siftdown_max(self.l, 0, x[-1])

  def decrease(self, x):
    x[0] -= 1
    heapq._siftup_max(self.l, x[-1])

  def insert(self, val):
    x = [0, val, -1]
    self.l.append(x)
    self.increase(x)
    return x

  def delete(self, x):
    last = self.l.pop()
    if x is last: return
    if self.l:
      self.l[x[-1]] = last
      if x < last:
        heapq._siftdown_max(self.l, 0, last[-1])
      else:
        heapq._siftup_max(self.l, last[-1])

This implementation is really short, but has two big downsides:

• It uses some internal functions of the heapq module, so might break at some point.
• It was around 5X slower than the lengthier version (but still MUCH faster than the built-in collections.Counter!).

Version 3 - HeapyKiYay

I suspected a big reason for version 2 being slow was the custom list. My final version was based on the previous one, but consisted of getting rid of GPSList and the use of the internal heapq functions, and instead, I wrote a custom version of an array-based position-aware heap. This version is clean and efficient, so I kept it. It’s the one on github.

The Leap

I also experimented with a variant where instead of using the more custom IndexedList, I created a more general data structure (the Leap) and used it to hold the consecutive pairs directly (instead of individual tokens). I ended up preferring the code with the IndexedList so I reverted the Leap, but it is more general and might be helpful elsewhere, so I’m explaining it here.

I needed a data structure that can represent an ordered array, supporting almost the same ops as a basic doubly-linked list in O(1):

• append(x): appends x to the end.
• delete(n): removes node n.
• start(x): returns the first node.
• end(x): returns the last node.
• next(n): returns the node following node n.
• prev(n): returns the node preceding node n.

I say almost because the leap importantly trades a list’s insert with an append operation (i.e. we can only insert at the end). In return, we get all of the following operations, in O(1) as well:

• first(x): returns the first node with value x.
• last(x): returns the last node with value x.
• leap(n): returns first node after node n with the same value as n.
• leapback(n): returns last node before node n with the same value as n.
• set_value(n, x): set the value of node n to x, keeping its position (n must be the last node with value x).

In addition to standard efficient in-order iteration on all elements, a leap also allows efficient in-order iteration on all elements with a given value.

I guess that this data structure must exist, but I haven’t heard of it. So I named it - “Leap” (although maybe “Leaped List” sounds better?). Either way, if this is a thing and you know its name, please lmk!

The implementation of a leap is straightforward - you can find it here.

My first implementation of a leap was around 2X shorter (by unifying the shared logic for positions and values), but that turned out to also be around 2X slower, so I opted for the more verbose version above. Note that the shorter impl also generalizes easily to the case where we want to leap by any one of several object properties. If you’re interested, here’s the more compact and general (but slower) version:

class LinkedList:  # A simple linked list.
  def __init__(self, val):
    self.val = val
    self.prev = self.next = None

  def delete(self):
    if self.prev is not None:
      self.prev.next = self.next
    if self.next is not None:
      self.next.prev = self.prev
    self.next = self.prev = None  # Not actually needed.

  def append(self, next):
    self.next, next.prev = next, self

POSITION_OBJECT = object()

class Leap:
  class Node:  # A leap-node is a superposition of two linked-list nodes.
    def __init__(self, val):
      self.val = val
      self.lls = [LinkedList(self) for _ in range('pos', 'val')]

    def next(self, val=POSITION_OBJECT):
      return self.ll(val).next.val if self.ll(val).next is not None else None

    def prev(self, val=POSITION_OBJECT):
      return self.ll(val).prev.val if self.ll(val).prev is not None else None

    def ll(self, val):
      return self.lls[0] if val is POSITION_OBJECT else self.lls[1]

  def __init__(self, init=None):
    self.first = {}  # Map from a value to its first occurrence.
    self.last = {}  # Map from a value to its last occurrence.
    if init is not None:
      for x in init:
        self.append(x)

  def nodes(self, val=POSITION_OBJECT):
    node = self.first.get(val, None)
    while node is not None:
      yield node
      node = node.next(val)
    return res

  def _delete(self, node, val):
    # This deletes the node from just one of its linked lists.
    if self.first[val] is node:
      self.first[val] = node.next(val)
    if self.last[val] is node:
      self.last[val] = node.prev(val)
    node.ll(val).delete()

  def delete(self, node):
    self._delete(node, POSITION_OBJECT)
    self._delete(node, node.val)

  def _append(self, node, val):
    # This appends the node to just one linked list.
    if val not in self.first:  # First item with value 'val'.
      self.first[val] = node
    else:
      self.last[val].ll(val).append(node.ll(val))
    self.last[val] = node

  def append(self, val):
    node = Leap.Node(val)
    self._append(node, POSITION_OBJECT)
    self._append(node, val)

  def set_value(self, node, new_val):
    # Replaces the value of a node in-place.
    self._delete(node, node.val)
    node.val = new_val
    self._append(node, node.val)

If you ever want to debug a leap, you could use something like this:

def debug_repr(leap):
  p = []
  for n in leap:
    s = str(n.val)
    if n is leap.first[n.val]: s += 'F'
    if n is leap.last[n.val]: s += 'L'
    p.append(s)
  return '[' + ', '.join(p) + ']'

Finally, since the impl with a Leap holding consecutive pairs is more involved, so let’s look at the canonical example from the BPE wiki page of performing 3 merges on the text “aaabdaaabac”. Since the byte pair “aa” (97, 97) is occurring most often, we’ll be merging it first. As our vocabulary starts with 256 tokens, one for each byte, we’ll be performing the merge (97, 97) -> 256.

Here we see that status of the leap right before this merge is performed:

For simplicity I’m only drawing forward edges.

The graph can be interpreted as follows: each node is a circle. Nodes that are first nodes in the leap are double circles. The current node the merge is iterating over is green, while the previous and next nodes are light blue. Finally the node that is just about to be changed (or that was just changed) is filled.

The below illustration (you might want to click on it to zoom!) depicts the status of the leap each time that it changes while performing 3 merges with the above text. The label under each line explains what is going on at that stage:

Hope you enjoyed this post and thanks again Andrej for the fun exercise!