This program is a dependency-free implementation of GPT-2. It loads the weight matrix and BPE file out of the original TensorFlow files, tokenizes the input with a simple byte-pair encoder, implements a basic linear algebra package with matrix math operations, defines the transformer architecture, performs transformer inference, and un-tokenizes the output with the BPE decoder. All in ~3000 bytes of C.

It's optimized efficiently enough so that GPT-2 Small takes a few seconds per reply on any modern machine. To do this I've implemented KV caching and an efficient matrix multiplication algorithm, with optional OMP parallelism.

You can then use this to create something like Chat GPT---just so long as you don't care about the quality of the output. (It's actually pretty terrible output, objectively speaking... But it does run.) There are a few quirks (especially with handling UTF-8 characters), and running the XL size model at long context length can require ~100GB of RAM. But if you're just typing with ASCII using GPT2-Small it should run just about anywhere.

I've uploaded the code to GitHub, so feel free to try and use it there.

Background: ChatGPT and transformers

In case you've been living under a rock for the past few months, ChatGPT is an application where you can talk to a type of machine learning model called a "language model" as if it was another person. It responds remarkably well, and GPT-4, the latest model that powers ChatGPT, is incredibly impressive.

This C program implements the behavior of ChatGPT using a much weaker model from 2019: GPT-2. Despite being just 2 smaller than GPT-4, it has no where near the same capabilities---but it is open source. So it has that going for it.

GPT-2 is a type of machine learning model called a "transformer". These neural networks take a fixed-size sequence of words as input, and predict the next word that will occur. By repeating the procedure over and over, you can use them to generate arbitrary-length sequences.

This post isn't meant to be an introduction to all the machine learning you'll need to know why a transformer is designed the way it is, but the rest of this post will be dedicated to describing how the above C code works.

Walkthrough the C Code

Getting started: Matrix Math (700 bytes)

Seeing as neural networks are just matrix operations. So we're going to need to get started by building a matrix library in as few bytes as possible.

My definition of a matrix is completely minimal:

We'll begin by by observing that while there are a bunch of different operations we'll need to implement, there are basically two "types" of operations"

1. Matrix-constant operations (e.g., add 7 to each entry of a matrix)
2. Matrix-matrix operations (e.g., add corresponding matrix entries)

This similarity allows us to use macros to pull out a bunch of the common logic into a meta-routine that knows how to operate on, for example, pairs of matrices and just leaves the specific operator implementation defined.

To do this in C, I'll define the function

as the following:

And so for example this lets us just write

and have it automatically expand to the full operation that perform elementwise addition or multiplication of two matrices. I define a few other easy to understand operations as well:

Now the thing about C's #defines is they're basically just glorified regexs. So when we actually run this, what's going to happen is we're going to take the line

and expand for the case of multiplication expand it to

But this replacement is almost literally just a regular expression replace. We could have put anything in place of OPERATION. This allows us to define a function like

Which at first glance looks rather confusing---what is that semi-colon doing there?---but if you just do a regular expression replace on it, you'll see that it expands to

where because the second expression doesn't do anything this is just equivalent to

(TAKE THAT LANGUAGES WITH PROPER MACROS. LISP ISN'T ALWAYS BETTER THAN C!)

Fast matrix multiplication (300 bytes)

The basic implementation of matrix multiplication is entirely straightforward: we just implement the naive cubic-time three loops: (There's nothing intelligent about my matrix multiplication. If you know how to make matrix multiplication fast you can just move along.)

Fortunately we can make it much faster with just a few bits of intelligence. Because of the way memory and caches work on most computers, it's (much!) faster to read and write to the same piece of memory over and over.

Later we're going to make one more change to the way we do inference and add a new parameter to the matrix multiply that instead allows us to only multiply part of Matrix A by Matrix B, which is useful when we've already pre-computed part of the product.

Neural network layers (300 bytes)

In order to write a transformer I'll need to define a few special neural-network specific layers. One of these is the GELU activation function, which you can just think of as magic.

I also implement a function to set the lower-diagonal of a matrix (after exponentiating the values). This is useful for what's called causal attention: we only want to attend to the past, not the future, and so we set the lower diagonal of the attention matrix to zero with this function.

And finally we need a layer normalization function. (Again another piece of magic that you can look up if you want. Basically what it does is normalize the mean and variance of each layer.)

The final piece of the model is the Linear function that just performs a matrix multiplication and adds (with tiling) a bias.

Transformer architecture (600 bytes)

With all of this out of the way, we can finally implement our transformer in just 600 bytes.

Now here's a fact that might not be completely obvious about transformer inference: once you've called the model to generate one token, you don't actually have to re-compute the entire function to generate the next token. In fact, you only need to do a small amount of additional work to generate each additional token.

This is because once you've computed the output of the transformer on the output of all the tokens up to the Nth token, you can re-use almost all of this output to compute the N+1st token (with a little bit more work.)

To actually implement this, I make all allocations in my code occur sequentially within the same block of memory, to guarantee that each matrix multiply will always use exactly the same memory. Then, at each iteration of the loop, I can just not zero-out the memory before using it for the next iteration, and the memory will already contain the results of the previous iteration. I just need to run the computation for the N+1st row.

Byte pair encoding (400 bytes)

The simplest way to build a language model is on a sequence of words. But because the total number of words is essentially unbounded, and language models need to have fixed-size inputs, it would be necessary to replace sufficiently rare words with a special [OUT OF DISTRIBUTION] token. This is no good.

While a simple “fix” for this would be to just use character-level language models that only know about individual letters, this would be a problem because it would mean that the model would have to learn the meaning of every word from scratch, and also reduces the effective context size of the language model by a factor of the average word length.

So to fix this, language models like GPT-2 work by creating tokens out of "word pieces". Some words might be tokens all by them-self, but rare words are broken up into smaller pieces. For example, the word “nicholas” might be broken up into “nich” “o” “las”.

The general algorithm for this is rather easy to implement: given a word we want to tokenize, we first split it into individual characters. Then, we look for pairs of adjacent tokens that should be merged, and merge them together. We repeat this until there are no more possible merges.

This algorithm is simple but unfortunately rather hard to implement in C because it requires a bunch of allocations, and requires keeping track of a tree-like structure of the tokens.

So instead, we'll turn the rather simple linear time algorithm into a potentially exponential time algorithm but save a bunch of code. Our basic idea will work like this in C-like pseudocode:

That is, to tokenize a word, we try each possible word in the vocabulary to see if it's a prefix of the current word. If so, we try to use this as the first token, and then recursively try to tokenize the rest of the word. We keep track of the best tokenization we've seen so far (as judged by the length, breaking ties by the index of the token in the vocab), and return that.

Weight loading (300 bytes)

We're almost done! The last thing we need to do is load the actual weights of the neural network off disk. This is actually pretty easy, because the weights are stored in a simple binary format that's easy to read in C: it's just a completely flat serialization of 32-bit floats.

The only thing we need to know is how big the various matrices are. And fortunately, this is also easy to figure out. Each of the GPT-2 model sizes have the same architecture, and the weights are saved in the same order, so all we need to do is read read the correctly-shaped matrices off of disk.

There's one final annoying thing. The layers of the neural network are not stored on disk in the order you might expect, with layer 0 first, then layer 1, then layer 2. Instead, the first layer is layer 0, then layer 1, and then layer .... TEN! (and then layer 11, and then layer 12.) This is because weights are stored when sorted lexicographically. And lexicographically, “0” comes before “1”, but “10” comes before “2”. So we have to do a bit of work to permute the weights into the correct order with the following code

Byte pair encoding loading (300 bytes)

In order to actually perform byte-pair encoding, we need to first load the byte-pair encoding vocabulary off disk. In an ideal world we'd actually just have a list of all the words in the vocabulary stored in some reasonable C-readable format, but because the original file was (a) meant for reading in Python, and (b) not meant to make it easy to parse in as few bytes as possible, we'll have to do some work here.

You might expect the file format to just be a list of words one after the other, but it's actually instead a list of the byte-pair encodings. What this means is instead of being able to read “Hello” as one token, the line is “H” “ello” which means we should be merging the tokens “H” and “ello” into a single token “Hello”.

The other challenge is that the file is encoded in smoothing-like UTF-8 (but not quite exactly that) for ... reasons. All of the printable ascii characters are encoded as themselves, but the non-printable characters from 0-31 are encoded as the value 188+the character. So for example, a space is encoded as the token “Ġ”. But now the problem is that the UTF8 encoding of “Ġ” is 0xc4 0xa0 when on disk, and so when reading it we have to do just some ugly work to convert this back to a space.

And while none of this is actually that hard to do, it still requires a fair amount of code which is annoying when you're trying to compress everything to be small.

Conclusion

It's really remarkable how you can distill so many decades of progress in machine learning to just a few thousand bytes. There is essentially nothing missing here from everything you need to run any state-of-the-art neural network (except for the actual model weights). While I mostly put this together for fun, it's a nice demonstration how simple neural networks actually are.