Build a GPT from scratch in PyTorch

Intro

0:00 Hello and welcome to this series, where we're going to build a minimum viable GPT from scratch, a small language model that you can read, train and understand from one end to the other. Let's begin with a quick look at where we're going to be by the end of the first three episodes.

0:15 This is our model writing a story, after a first, short training run on my laptop, which is where we're going to be by episode two. And by episode three, it follows an instruction, such as which words its story should use.

0:30 My name is Adrian, and I'm a software engineer who is also a researcher. This series is going to be something of a dev log, and something I want for my own understanding as well. The goal is to understand the architecture of a GPT model end to end, along with the whole pipeline around it, from pretraining to fine-tuning and reinforcement learning. What I want to focus on is how the different blocks of parameters work together, in code that you can read line by line, without the engineering tricks a model needs at scale. Our model is going to be small enough to train in about ten minutes on a laptop, and we're going to see how far we can push it on a single machine. It's not going to do arithmetic, search the internet or be an agent. It's going to start out writing short children's stories, and while they're simple, they read like English and mostly hold together. Later on, we're going to train it on other datasets too. And it's a complete, working transformer that you can tweak, evaluate against other approaches, and play with. A quick note on how the series is made: often, the voice you're hearing is generated with AI, and it's designed to sound like me.

1:38 On the website of the series, you can find more information about the project. The code is open source on GitHub, tagged at each step of each episode, so that you can follow along. The project has docs of its own, with the diagrams we're going to go through today, and each episode gets a page here as it comes out.

1:57 I want to approach this series as an educational project. With my background as a software engineer, and my PhD in reinforcement learning for reasoning in generative models, I'm keen to explain how a GPT works from first principles. We're going to build quickly, so that we have a working model early on, but we're going to benchmark it systematically along the way. That way, each change comes with numbers that show what it did, such as the loss, and the stories from before and after. And we're going to take our time when it comes to the complexity that gets added, a piece at a time, so that you can understand each one before the next comes in. The later parts of the pipeline are quite close to my research, and I'm hoping this series helps me explain them, too.

Where to go deeper

2:40 There are some amazing resources out there, and one of the best known is nanoGPT, by Andrej Karpathy. Its README calls it the simplest, fastest repository for training medium-sized GPTs, with a model and a training loop of about three hundred lines each. It reproduces the smallest GPT-2, with a hundred and twenty-four million parameters, on a node of eight A100 GPUs in about four days. Its successor, nanochat, goes all the way to a small ChatGPT that you can talk to. And its quick start trains a character-level GPT on the works of Shakespeare, which is a good first model to train on your own machine. These go much further than we're going to on performance and scale. Here, we're going to start from a single machine, and keep each line of the model as simple as we can.

3:26 And for the mathematics, there's 3Blue1Brown, Grant Sanderson's channel, which explains maths with animations that build the intuition step by step. His series on neural networks goes from the basics of how a network learns, all the way to transformers, and how attention works inside them. I'd recommend it alongside this series, as it goes deeper into the maths than we will, while we focus on the code, and on the whole pipeline around the model.

The plan

3:54 So, here's the plan, starting with the architecture, and today, we're going to build it in PyTorch, as a blank slate that hasn't learned anything yet. Then we're going to pretrain it on a dataset of children's stories, so that it learns to continue a story. Supervised fine-tuning then teaches it to follow instructions, such as which words a story should use. After that, it's going to be direct preference optimisation. The model writes a few stories for each instruction, and our own code scores them, say on how many of the requested words they use. Then we teach the model to prefer the better ones. And then reinforcement learning, with proximal policy optimisation, and possibly others later on. Along the way, we measure each step, with metrics and plots that compare the models before and after. And we're going to touch on things such as LoRA, and a KV cache for faster generation.

The data: TinyStories

4:46 The data we're going to start with is TinyStories, from Ronen Eldan and Yuanzhi Li, at Microsoft Research. Language models of a few tens of millions of parameters struggle to write coherent English when they're trained on text from the web. So TinyStories is a synthetic dataset of short stories, written by GPT-3.5 and GPT-4, with words that a typical three or four year old would understand. The paper shows that models below ten million parameters write fluent stories with almost perfect grammar when they're trained on it. Here, they compare a model of twenty-eight million parameters, trained on TinyStories, with GPT-2 XL, which has one and a half billion, on continuing a story. And to keep the stories diverse, each one was generated around three random words, a noun, a verb and an adjective, out of a list of about fifteen hundred. Some were also asked for features, such as a dialogue, a plot twist or a bad ending. Those words and features come in handy in episode three, as they're what our instructions are going to ask for.

Attention Is All You Need

5:48 The paper that started transformers is Attention Is All You Need, from Ashish Vaswani and others at Google, in 2017. If you've been following the field at all, you've seen this diagram many, many times. It was written for translation, so it has an encoder on the left, which reads a sentence in one language into an intermediate representation. The decoder, on the right, writes it out in the other language. Ours is going to be decoder-only, which is what the most successful language models are today, from GPT onwards. So we keep the right-hand column, without the attention over the encoder's output in its middle. The input comes in at the bottom, and it's embedded, a positional encoding is added, and it goes through masked multi-head attention and a feed-forward network. Those repeat N times, six in the paper, and a linear layer and a softmax give the probabilities of the next token. Our own diagrams are going to read from left to right instead. Inside the attention blocks, there's scaled dot-product attention, on the left, which we're going to look at as a formula next. Multi-head attention, on the right, runs several of those side by side, each on its own slice of the numbers, and joins up what they find. The paper used eight heads, and we're going to use six.

7:03 In the attention formula, each token is first projected into three vectors, a query, a key and a value. The dot product of one token's query with another token's key gives a score for how much the first should take from the second. Q times K transposed does that for all the pairs of tokens at once. We divide by the square root of d k, the size of a head, as dot products grow with the number of dimensions, and large scores would push the softmax towards a single token. The softmax takes the exponential of each score and divides by their sum, so each token's scores become weights that are positive and add up to one. And the output for each token is the weighted sum of the values. Our attention is also masked, or causal, so a token attends to its own position and the ones before it, and not to the ones after it. The mask is what lets the model learn to predict the next token without seeing it. In our model, each token is three hundred and eighty-four numbers, split between six heads, so d k is sixty-four, and we divide by eight.

What we change

8:04 Our model follows the paper's decoder, with the changes that GPT-2 has. It's decoder-only, as we said, and the positions are learned, as an embedding table of their own, instead of the paper's fixed sine and cosine waves. Most recent models use rotary position embeddings, which we might look at later on. The feed-forward network uses GELU, a smoother function than ReLU, and we're going to write it out by hand. The layer norm goes before each sublayer instead of after it, which makes a deep stack of blocks easier to train. And the output layer shares its weights with the token embeddings, which saves quite a few parameters.

The architecture

8:41 So, let's look at the parameters of our model, all the way from the prompt to the next token. A prompt such as once upon a time is first split into tokens, and each token is turned into an id, a number. The tokenizer does that, and we're going to write it before we train the model. It has no parameters, as it's a fixed table from words to ids, built once by counting words, and training doesn't change it. The ids go into the token embeddings, which are a table with a row for each token in our vocabulary. The vocabulary is going to be the ten thousand most frequent words and punctuation marks in TinyStories. Each row holds three hundred and eighty-four numbers that the model learns, so that's three point eight four million parameters. Positions get a table of their own, with a row for each of the two hundred and fifty-six positions in our context window, which is the most tokens the model can look at once. The two are added together, and the sum goes into the first transformer block. There are six blocks, built alike, and each has its own weights. Inside, a layer norm before attention has seven hundred and sixty-eight parameters, a gain and a bias for each of the three hundred and eighty-four numbers. The query, key and value projections are each a matrix of three hundred and eighty-four by three hundred and eighty-four, plus a bias, and they're split into six heads of sixty-four. The output projection is one more of those, and it mixes what the six heads found. Then a second layer norm, and the feed-forward network, which we're going to call the MLP. Its up projection goes from three hundred and eighty-four numbers per token to four times that, fifteen hundred and thirty-six. The MLP is the biggest piece of each block, and it's where a lot of the computation happens. The down projection brings it back to three hundred and eighty-four, so each block's output has the shape of its input, and can be the input of the next block. That makes one point seven seven million parameters per block, and ten point six million for all six. After the sixth block, a final layer norm, and the language model head, which gives each of the ten thousand tokens a score. It shares its matrix with the token embeddings, so it doesn't add to the count. In total, the model has fourteen million, five hundred and eighty-five thousand, eight hundred and fifty-six parameters.

How data flows through it

10:54 Now let's go through the architecture again, from the point of view of the data flowing through it. Once upon a time becomes four tokens, so four ids. They're not learned: the tokenizer numbers the words by how frequent they are, so a, one of the most frequent, gets seven. And we're going to use T for how many tokens there are, four here, and up to two hundred and fifty-six. The embedding tables work as lookups, picking out rows by id. The token table holds ten thousand rows, but out of it we take the four rows for our four ids, so four by three hundred and eighty-four. And out of the position table, the rows for positions zero to three. The two are added, and that sum is how each token's vector carries both what the token is and where it is. In the block, the attention norm keeps the shape, four by three hundred and eighty-four. Then the projections give each token a query, a key and a value, each split into six heads, so each head sees four by sixty-four. Each head computes the attention formula on its own, with its own queries, keys and values, so it can learn to look for something different. The six outputs of four by sixty-four are joined back into four by three hundred and eighty-four, and the output projection mixes them. Then the residual connection: what came into the block is added to what attention produced, so attention adds to the token's vector, without replacing it. The MLP then works on each token on its own, with a norm, up to fifteen hundred and thirty-six with GELU, back down to three hundred and eighty-four, and added to its input. The block's output goes into the next block, six times over, with each block's own weights. Then the final norm, and the head gives four by ten thousand, a score for each token in the vocabulary, at each of the four positions. To generate, we take the scores at the last position, turn them into probabilities with a softmax, and sample the next token from them. It's appended to the prompt, and the whole thing runs again for the token after it. To train, we use all the positions at once. As the mask keeps each position from seeing what comes after it, the first position is asked to predict the second token, the second position the third, and so on. The cross-entropy loss measures how much probability the model gave to the token that actually came next, and training nudges the weights to give it more next time.

The model in PyTorch

13:17 All right, let's write it, starting with model dot py, which is going to hold the whole model. It opens with a docstring saying what it is: a minimum viable GPT, a decoder-only transformer you can read end to end. From dataclasses, we import dataclass, which is Python's way of declaring a class that mostly holds values. It writes the constructor for us, from the fields we list along with their types and defaults. So GPT config is going to keep all the sizes we've talked about in one place. The vocab size is ten thousand, the number of tokens the model knows. The block size is two hundred and fifty-six, and its comment says it's the context window, in tokens. N layer is the number of transformer blocks, six, and n head the number of attention heads in each of them, six as well. Then the embedding size, three hundred and eighty-four, which is how many numbers represent each token as it flows through the model, split between the heads, six times sixty-four. Dropout is ten percent: during training, it zeroes a random tenth of the values in a few places, so that the model can't lean too heavily on a few neurons, and that helps it generalise. And bias says whether our linear layers and layer norms have a bias term, which they do. Requirements dot txt lists what the code needs, which is PyTorch, from version two point five. A README says what the project is, a small GPT you can read from one end to the other, and that it's the code of this series, tagged at each step. The MIT licence lets you take the code and build on it, in your own projects too. And the git ignore file keeps Python's caches, our virtual environment and the files macOS leaves in each folder out of the repository.

Causal self-attention

14:59 Next comes attention, and for it, we import torch dot n n as n n, which is PyTorch's library of neural network building blocks, such as linear layers, along with the module class our layers are going to extend. And from torch dot n n, we import functional as F, which has those operations as plain functions, such as the softmax and scaled dot-product attention. Causal self attention extends n n module, PyTorch's base class for anything with parameters. Being a module means PyTorch finds our layers' weights, moves them to a GPU along with the model, and switches the layers between training and evaluation mode for us. In the constructor, we call the parent's init first, which sets up that bookkeeping. We keep the number of heads, as the forward pass is going to need it to split the channels. Then four linear layers, each from three hundred and eighty-four channels to three hundred and eighty-four. A linear layer multiplies its input by a matrix of weights it learns, and adds a bias. The query projection turns each token into what it's looking for, the key projection into what it offers, and the value projection into what it passes on when another token attends to it. The three could be one layer with three times the outputs, which is what GPT-2 does, as it's a bit faster. I'm keeping them apart, so that each is easy to find, in the code and when we print the model. The output projection comes after the heads, and mixes what they found. Then we keep the dropout probability, for the attention weights, and a dropout layer for the output. The forward method is what runs when the module is called on an input. X comes in with three dimensions, B, T and C, as the comment says: the batch, which is how many sequences we process at once, the tokens in each, and the channels, three hundred and eighty-four. For the queries, we project x, and view the result as B by T by six heads by sixty-four, which splits each token's channels into six slices without copying anything. Transpose then swaps the tokens and the heads, so that each head becomes a batch of its own, of T tokens by sixty-four. The keys and the values go through that too. Then scaled dot product attention, which is PyTorch's implementation of the formula from the paper. Is causal applies the causal mask we saw with the formula. Dropout on the attention weights applies while training, and not when we evaluate. We're leaning on PyTorch for attention for now, as it's fast and well tested, and we can write it out by hand later in the series. Above it, we write down what it computes: softmax of q, k transposed, over the square root of sixty-four, times v, where each token sees its own position and the tokens before it. Then we undo the split, and transpose the heads back. Contiguous copies the result into one block of memory, as view needs that, and view joins the six heads into three hundred and eighty-four channels again. Last, the output projection mixes the heads, and dropout applies to the result.

The MLP and GELU

17:57 After attention, each block has its MLP, for multi-layer perceptron, which is two linear layers with a non-linear function in between. Without that function, the two layers would collapse into a single linear one. The paper used ReLU, which keeps positive inputs as they are and turns negative ones into zero. We're going to use GELU, the Gaussian error linear unit, which GPT-2 uses as well. It's a smooth version of ReLU, where large positive inputs pass almost unchanged, and large negative ones go to about zero. Small negative inputs keep a little of their value, so their gradient isn't zero, and the curve has no sharp corner at zero. In practice, models with GELU tend to train a little more smoothly than with ReLU.

18:40 We write GELU by hand, so we import math, Python's math module, for a square root, and torch, for the error function. GELU uses the normal distribution's cumulative distribution function, which can be written with erf, the error function, as a half times one plus erf of x over the square root of two, so the function returns x times that. Above it, a comment says what that computes in words: x times the probability that a standard normal variable is at most x. PyTorch has GELU built in, along with a faster approximation through tanh, which GPT-2 used. We write it out, so that the formula is there in our code to read. Then the MLP module, where the up projection goes from three hundred and eighty-four channels to four times that, fifteen hundred and thirty-six. Four times is the ratio the paper used, and a lot of models since have kept it. The wider middle layer is where GELU applies, on four times as many channels per token. The down projection brings it back to three hundred and eighty-four, so the block's output matches its input, and then dropout. The forward pass reads from the inside out: x goes through the up projection, then GELU, then the down projection, then dropout.

The transformer block

19:55 Now the transformer block, which puts attention and the MLP together. It has a layer norm before attention, the attention, a layer norm before the MLP, and the MLP. A layer norm rescales each token's three hundred and eighty-four numbers to a mean of zero and a variance of one, and then applies a gain and a bias that it learns. That holds the values each sublayer reads in a steady range, which keeps training stable. The forward pass is where the residual connections are. X plus the attention of the normalised x: attention reads a normalised copy, and adds what it found onto x, instead of replacing it. Then x plus the MLP of the normalised x. So each sublayer makes a correction to a running representation of each token, which is called the residual stream. And the gradient has a direct path back through the additions to the first layers, which is what lets us stack blocks deep. Each norm comes before its sublayer, like in GPT-2, and the block returns x with the shape it came in with, so six of them can be chained.

Embeddings, blocks and the head

20:57 And now the GPT, which puts it all together, and keeps the config for later. The token embeddings are an embedding module, a lookup table with a row of three hundred and eighty-four learned numbers for each of the ten thousand tokens. The position embeddings are a second table, with a row for each of the two hundred and fifty-six positions in the context window. Then dropout, and the blocks, six of them in a module list, which is a Python list that PyTorch can see into, so it finds their parameters. They're all built from one class, but each has its own weights, and each learns something different. Then the final layer norm, and the language model head, a linear layer from three hundred and eighty-four channels to a score for each of the ten thousand tokens, without a bias. Then we tie the weights, so the embedding table and the head share one matrix. It's a matrix of ten thousand by three hundred and eighty-four in both, so the head scores each token by how well the final vector matches the token's own embedding. That saves three point eight four million parameters, more than a quarter of the model, and it tends to train better on small data. Above it, the comment reads: weight tying, the head scores tokens with the matrix that embeds them. And last, the constructor calls init weights, to set the starting values. Init weights goes through the modules by name, and linear layers and embeddings start from a normal distribution with a mean of zero and a standard deviation of zero point zero two, which is what GPT-2 does. The values have to be random, as weights that started equal would get equal gradients, and the neurons would stay copies of each other. The output projection and the down projection start smaller, divided by the square root of twice the number of layers, so that the stream's values don't grow as the twelve additions pile up. Above that, the comment says: these two add into the residual stream, twice per block, so they start smaller. And the biases of all the linear layers start at zero. Num params adds up the number of values in each parameter, and PyTorch lists the tied matrix once, so it's counted once. Then the forward pass, which takes idx, a batch of token ids, B sequences of T tokens, and optionally the targets. Pos holds the positions, from zero to T minus one, on the device the ids are on. We look up the token embeddings and the position embeddings, and add them. The positions are T by three hundred and eighty-four, and PyTorch broadcasts them across the batch, which means it adds those T rows to each sequence of the batch. That sum is how the model knows where each token is, as attention on its own doesn't see the order. Dropout applies to it, and then x goes through each block in turn, so the output of one is the input of the next. Then the final norm, and the head gives the logits, a score for each token in the vocabulary at each position. If we have targets, which are the tokens that actually come next, we compute the cross-entropy loss. It takes the softmax of the logits and the negative log of the probability given to the right token, and averages it over all the positions. View flattens the batch and the positions into one long list, which is the shape cross entropy expects. And forward returns the logits and the loss, which is None when there are no targets.

Generating tokens

24:12 Last in the model, generate, which writes new tokens one at a time. No grad tells PyTorch not to keep what it would need for gradients, as we're not training here, and that saves memory and time. It takes the ids so far, how many new tokens to write, a temperature, top k, and an end-of-text id to stop on. For each new token, we run the model on the ids, cropped to the last two hundred and fifty-six. The comment above says why: the model has block size positions, so it sees no more than that many tokens. We take the logits at the last position, which are the scores for the next token, and divide them by the temperature. A temperature below one sharpens the distribution towards the likeliest tokens, and above one flattens it. With top k, we find the k-th highest score, and set the scores below it to minus infinity, so the softmax gives those tokens no probability at all. That keeps the long tail of unlikely tokens out of our stories. Softmax turns the scores into probabilities, and multinomial samples a token from them. We append it to the sequence with cat, and if all the sequences in the batch have just sampled the end-of-text id, we stop early. And generate returns the ids, the prompt followed by what it wrote. Generating this way runs the whole sequence through the model for each new token, which is simple, and slower than it needs to be. A KV cache is how that gets faster, and we can come back to it later on.

The model as PyTorch sees it

25:38 To see the model as PyTorch sees it, we write summary dot py, and its docstring says what it shows: the modules, the parameters, the shapes through a forward pass, and what the model generates before any training. It imports math, torch, and our GPT and its config from model dot py. We seed PyTorch's random number generator, so that the random weights come out identical on each run, along with everything that follows from them. We build the config with its defaults, the model from it, and print the model, which lists its modules as a tree. Then we count the parameters per layer, to compare with our diagram. Named parameters gives each parameter along with its name, such as blocks dot zero dot attn dot q proj dot weight, and the second-to-last part of the name is the layer, q proj. So we add up the number of values per layer, across the six blocks, and print them, then the total. The comment says why the head isn't in the list: it shares the embedding's weight, which PyTorch counts once. Next, the shapes, and eval switches dropout off, so the forward pass is deterministic. We make a batch of random token ids, two sequences of four, and random targets. A forward hook is a function PyTorch calls with a module's output, whenever the module runs. So we register one on the embeddings, each block, the norm and the head, to print the shape of what comes out. Then we run the model once, and remove the hooks. Before training, the model is guessing, with each of the ten thousand tokens about as likely as any other. A uniform guess over ten thousand tokens has a loss of the natural log of ten thousand, so we print both, to compare. And last, we generate, although we don't have a tokenizer yet. So we pass the ids that once upon a time is going to have, sixty, sixty-one, seven and forty-seven, and ask for twelve new tokens. And the README gets a section on how to run it.

27:33 Let's run it, and first comes the model as PyTorch prints it: the embeddings, the six blocks with their attention and their MLP, the final norm and the head. The tree doesn't show how they connect, as that's in the forward methods, which is what our diagrams are for. Then come the parameters by layer, where the token embeddings are three million eight hundred and forty thousand, the up and down projections about three million five hundred thousand each, across the six blocks, and the total is fourteen million five hundred and eighty-five thousand eight hundred and fifty-six, which matches our diagram. Next, the shapes follow the data through the model. Two sequences of four ids become two by four by three hundred and eighty-four in the embeddings, while the position table is four by three hundred and eighty-four, broadcast across the batch. Each block keeps that shape, and the head gives two by four by ten thousand, a score for each token at each position. The loss is about nine point one, next to the natural log of ten thousand, nine point two, so the untrained model is close to a uniform guess, as it should be. And the twelve ids it generated after our four are random. With a tokenizer, they'd decode to English words in a random order, as the model hasn't learned anything yet.

28:45 Put together by part, the MLPs hold the most, about seven million one hundred thousand parameters, or about half the model. Attention has three million five hundred thousand, and the token embeddings three million eight hundred and forty thousand, which the head shares. The positions and the layer norms are a small fraction of the model. So the largest share of what the model learns sits in the MLPs, which is also the case in much larger models.

Docs

29:12 Last, the project gets docs, which we build with MkDocs, a static site generator for documentation written in Markdown. Its config names the site, points at the repository, and lists the Markdown extensions we need for diagrams and formulas. It uses the Material theme, and a small stylesheet gives it the colours and fonts of the series' website. The home page has the model at a glance, as a table of its sizes, and how to run it. The architecture page has today's two diagrams, written in Mermaid, which turns a few lines of text into a diagram. It goes through each piece of the model, with its code and the formula behind it, such as attention's and GELU's. A GitHub Actions workflow builds the site, and publishes it to GitHub Pages whenever the docs change on the main branch. And the README says how to serve the docs on your own machine.

30:00 The docs are going to get a page for each episode. Whenever you'd like to look at where the parameters are, or how the data flows, the architecture page has both diagrams, along with each piece of the model and the formula behind it.

Wrap-up

30:16 So, today we picked TinyStories as our data, and designed a decoder-only transformer from Attention Is All You Need, with GPT-2's changes. We looked at where its fourteen point six million parameters are, and at how a prompt flows through it, from token ids to the probabilities of the next token. And we wrote it in PyTorch, class by class. The model runs and generates tokens, but its weights are random, so its loss is close to a uniform guess. I hope that between the two diagrams and the code, you've got a good grip on how a model like this is put together.

30:51 To finish, let's read two longer stories from a bigger version of this model. In episode twelve, we make it about six times bigger and pretrain it on four rented GPUs, and then it writes stories like these, of up to three hundred tokens. They're simple, and a word goes wrong here and there, but they hold together from the first line to the last. And they come from the blocks we wrote today, with a few changes along the way, and a lot more parameters.

31:20 This is going to be an evolving project, so I'd love to hear from you. Let me know in the comments which parts of the architecture you'd like us to go deeper on, and if you'd like to follow along, subscribe to the channel. I hope you find this useful, and thanks for watching.