Why do OpenAI's GPT-2 weights beat mine? Part five: data quality

When I finished learning how to build an LLM from scratch, I was left with a mystery: my own models were not as good as OpenAI's original GPT-2 models, despite being based on the same architecture.

My models all had 163M parameters, and followed the design from Sebastian Raschka's book "Build a Large Language Model (from Scratch)". That meant that they were pretty much the same as the setup for the OpenAI GPT-2 "small" instance, except that they did not use weight-tying or bias on the QKV matrices. Weight-tying means that you re-use the initial embedding matrix as the output head at the end, and using it means that GPT-2 small saved quite a few parameters -- it was 124M rather than 163M -- at, at least in my own experiments, a cost in quality; similarly, while I found that QKV bias made a tiny improvement in loss terms, I'd felt it was likely within the noise.

But GPT-2 small consistently beat my models on an instruction fine-tuning (IFT) task -- also adapted from Raschka's book. That test fine-tunes the model on a subset of the Alpaca dataset, until validation loss starts rising, and then runs a test set through the resulting model. The responses to the test set questions are stored, and then I run all of the responses from all of the models under test past GPT 5.5 in one go to get an aggregate score; more details here. GPT-2 small always did better than any of my models on this.

Additionally, it did surprisingly well on a simpler eval -- one that just measured the cross entropy loss it got on a test set. It scored close to my own best models, and better than many of them. What made this result particularly interesting was that the test set in question was a split of my own training data; my models would not have seen it when training (at least, in theory), but it seems likely that it would be much more similar to their own training data than it was to OpenAI's.

I've checked two things while probing this mystery:

• It seems very likely that the GPT-2 models were overtrained by modern standards; would overtraining my own models get them closer? It turned out that no, it probably didn't help with the IFT eval (though there might have been some signal there). It did help quite a lot with the test loss eval, though.
• The way I was handling dropout in the IFT test might have been unduly benefiting some models while working against others. I decided to standardise on not using dropout during this eval, as (counter-intuitively for me) it seemed to harm the results of most models, even those that had been pre-trained with dropout. In particular, the OpenAI weights were harmed by using dropout, and making a change that benefited them (along with some of my own models) seemed the most conservative approach to take in investigating this.

The next thing I wanted to look into was the training data.

The exact dataset that the various GPT-2 models were trained on has never been released; all we know about it is from the paper, where they say:

[W]e created a new web scrape which emphasizes document quality. To do this we only scraped web pages which have been curated/filtered by humans. Manually filtering a full web scrape would be exceptionally expensive so as a starting point, we scraped all outbound links from Reddit, a social media platform, which received at least 3 karma. This can be thought of as a heuristic indicator for whether other users found the link interesting, educational, or just funny.

They called it "WebText". There is an OpenWebText that tries to replicate it, but although they tried to follow the same procedure as the original, there's no guarantee that it is all that similar.

By comparison, I'd normally been training against FineWeb. While this is a general web-scraping dataset, without the "curation" provided by using only stuff that was linked from upvoted Reddit posts, it has been refined to remove any obvious junk. I had felt that it was pretty much equivalent.

But what if I were wrong about that? I decided to see if I could get better models by using better data.

[ Read more ]

Extending Raschka's GPT-2: an MoE trained from scratch on an RTX 3090

Mixture-of-experts models are really nifty. You get inference speed close to a small model's, with a lot of the smarts and knowledge of a large one. While they use as much memory as an equivalently-sized dense (non-MoE) model, they're much faster. The frontier labs don't publish their architectures, but Claude and ChatGPT are widely rumoured to be MoEs these days -- and certainly many large open-weights models like DeepSeek and Kimi K3 are.

In this post, I'll show you how I added MoE support to the GPT-2-style code from Sebastian Raschka's book "Build a Large Language Model (from Scratch)", then used that to train a 446M-parameter model with 220M active parameters from scratch on my RTX 3090 -- essentially, GPT-2 small with 6 experts, 2 active per token. I wanted to understand how MoEs work, and as always, felt that the best way to do that is to build one (and then to write it up like this). Hopefully because it's all fresh in my mind as I write this, I should be able to explain things clearly for others who've finished Raschka's book.

The training run took just less than eight days, and the resulting model got a better loss on my test set than any of the other models I've trained so far -- which was a good thing, given that it took four times longer to train! It was also better than the original OpenAI GPT-2 small (124M parameters), but not as good as GPT-2 medium (345M parameters), although it was close. That second result was interesting, as my model had more total parameters than GPT-2 medium, but fewer active ones.

On an instruction fine-tuning test, it did better than any of my other models, but worse than both OpenAI models (why OpenAI's models are so good on that particular test is a mystery I'm digging into separately).

I'll go into those numbers in more depth later on. Firstly, though, I'll run through the code that I needed to add to GPT-2 to make this work -- not just the code for the experts themselves, but the additional code for training it. With MoEs you can't just train to minimise loss on your training set -- you need to add on an "auxiliary" loss to make sure they actually use all of their experts, and that was where it became most interesting.

Before we get into the weeds, though, let's start off with the basics; how do MoEs work in Transformers-based LLMs?

[ Read more ]

Writing an LLM from scratch, part 34b -- from bigrams to GPT-2, one component at a time (in JAX)

This post is the capstone of the most long-running series on my blog. In December 2024 (!), I started reading Sebastian Raschka's book "Build a Large Language Model (from Scratch)", and worked through it carefully. Being who I am, despite trying to apply a strict "no side quests" policy, I found myself zooming off and digging into all kinds of things.

It's time to wrap it up. I had decided that the endpoint would be to build and train an LLM from scratch just using my notes -- no reference to the book, no reference to the model code I'd written when following the book. After an X/Twitter poll, I decided to use JAX for that, just to make sure that I really was building it from scratch and not regurgitating bits of PyTorch code like a bad coding LLM spitting out half-digested lumps of Stack Overflow.

In my last post, I showed how I built a JAX training script that mirrored what I had built for the original PyTorch version of the model. To test it as I went along, I used it to train a really dumb "LLM", which instead of trying to predict the next token for every token in an input sequence, instead predicted the input -- that is, if you fed it

The fat cat sat on the mat

It would return the same thing. I called that an A-to-A model.

In this post, I'll show you how I turned it into a GPT-2 model, and then trained it from scratch on my RTX 3090 (using the parameter counts for the original paper's "small" size). What turned out really well with this is that I found a route that meant that almost every component I added made the model better! That's not guaranteed -- sometimes different aspects of an AI model depend on each other, so adding A without also adding B makes things worse. But (admittedly with a bit of backtracking in places) I was able to find a route that shows a nice clear progression.

The final training run took 37 hours 15 minutes -- compared to 40 hours, 38 minutes for an equivalent PyTorch model. That is despite it being full-fat 32-bit -- the PyTorch one was using Automatic Mixed Precision (AMP), which allowed it to use 16-bit calculations in places where it would be relatively harmless in terms of loss.

When asked to continue "Every effort moves you", it came back with a decent response:

Every effort moves you closer to your goals, but if you are unsure of what it takes, you don’t

The model got 3.418784 loss on my held-back test dataset, as compared to my PyTorch model's 3.538161, and even more impressively, it was better than the original GPT-2 small's result of 3.499677 on the same dataset! However, just as I found previously, the OpenAI weights still beat mine consistently in instruction fine-tuning challenges.

Let's get started.

[ Read more ]

Writing an LLM from scratch, part 34a -- building a JAX training loop for an LLM training run

For over a year, I've been using Sebastian Raschka's book "Build a Large Language Model (from Scratch)" -- and the multitude of side-projects that have branched out from reading it -- as something like a curriculum for learning about modern AI. The one final task I had set myself was to build and train an LLM from scratch just using my notes -- no reference to the book, no reference to the model code I'd written following the book.

As an output, I wanted something as good as my best PyTorch model based on Raschka's code -- a base model, trained on 3.2B tokens, that my (admittedly limited) evals ranked as being close to the original GPT-2 small's quality.

I wanted to use a different framework, just to make sure I wasn't parroting code that I'd somehow memorised, so I asked people on Twitter which one I should use, and the winner was JAX.

I took a slightly different route to Raschka's book; he takes an inside-out perspective, explaining things like attention, gradually building up a complete GPT-2-style model, and then building a training loop on top of it. I wanted to go outside-in: I'd put together a training harness to train the simplest-possible model with an API similar to a real LLM, get that working to my satisfaction, and then add features to that simple model, one by one, until it had the full architecture in place. The plan (which actually worked out nicely!) was that I'd be able to show how each change improved things.

That's all done now, and I'm posting about it in two parts; in this one, I'll explain how I built the training harness, and in the next, I'll show the actual building and training of the LLM.

So let's get started!

[ Read more ]

Writing an LLM from scratch, part 33 -- what I learned from finally getting round to the appendices

After finishing the main body of "Build a Large Language Model (from Scratch)", I set myself three follow-on goals.

The first was training a full GPT-2-small-style base model myself. That was reasonably easy to do but unlocked a bunch of irresistible side quests; having finally got to the end of those, it's time to move on to the others: reading through the book's appendices, and building my own GPT-2 style model in JAX.

This post is about the appendices. The TL;DR: there was stuff in there that could have saved me time in my side-questing, but I think that having to work those things out from scratch probably helped me learn them better.

[ Read more ]

Writing an LLM from scratch, part 32m -- Interventions: conclusion

Last November, when I finished the main body of "Build a Large Language Model (from Scratch)", I set myself a number of follow-on goals. One was "training the full GPT-2 base model myself".

I've reached the end of that journey, with a model that is almost -- if not quite -- as good as GPT-2 small, trained in 44 hours on my own machine, so I thought it would be worth summarising how it went.

[ Read more ]

Writing an LLM from scratch, part 32l -- Interventions: updated instruction fine-tuning results

I've been working on a GPT-2-small-style LLM based on Sebastian Raschka's book "Build a Large Language Model (from Scratch)", and have tried a bunch of different things to see if I could get it to approach the quality of the original OpenAI GPT-2-small, measured in terms of loss on a held-back test dataset. After working through them, in my last post, I managed to train one that was almost (if not quite) there.

Now, back before I started digging into these interventions, I was doing three evals for each model I built; a smoke test (to see if it could give a coherent completion to "Every effort moves you"), a test for that test set loss, and an instruction-following test that fine-tuned the model on the Alpaca dataset, got it to generate results for a test set of instructions, and then used an LLM as a judge to score them.

The idea behind this was that the loss on the test set was an interesting technical measure of the quality of a model, but it didn't really tell us much about how useful it might be in reality.

Unfortunately, in January, I realised that my methodology was bad; because I was asking the LLM to score a model in isolation, the LLM's natural randomness would mean that results were not really comparable, at least for models that were reasonably close in quality.

For example, if two models both replied to

Name the author of 'Pride and Prejudice'.

with:

The author of 'Pride and Prejudice' is Sarah Palin.

...then one run of the instruction-following test might "find the judge LLM in a good mood" and get, say, 5% -- after all, the model tried to answer, and actually used a real person's name, even if the answer was totally wrong. But in another run, the judge might be in a "worse mood" and score it at 0%.

My fix was to have two scripts:

• One that fine-tuned the model then got it to generate responses, then saved those responses in a file.
• One that took a bunch of files generated by the above, one for each of a set of different models, and presented them to the LLM together, so that it would (hopefully) be consistent in how it rated them relative to each other.

The details are here.

Because doing it that way was significantly more work, I've not been doing these tests as part of the interventions mini-series. I felt it would make more sense to wait until I'd tried a bunch of interventions and got a number of models to try.

Now I have those, so let's give it a go!

[ Read more ]

Writing an LLM from scratch, part 32k -- Interventions: training a better model locally with gradient accumulation

I've been working on a GPT-2-small-style LLM based on Sebastian Raschka's book "Build a Large Language Model (from Scratch)". I've trained various versions of it in the cloud to work out which interventions to the model and training code had the best effects on the loss it gets on a specific test dataset, and now I wanted to do a training run locally to match the best of those. For that, I wanted to match the batch size I was using for the cloud training runs.

When I first started learning this stuff, batching seemed like a performance thing -- with highly parallel systems like GPUs, it generally turned out that you could run a batch of (say) two inputs through a model in less than twice the time you could run one, so it made sense to batch them up.

For inference, that is exactly the advantage you get, but when training, it's become increasingly clear to me that you can also get an improvement in the quality of the model from batching. The best intuitive model I have is that if you run inputs through one-by-one, adjusting parameters after each, then it's easy for the model to "overcorrect" each time. With batches, you get an average set of gradients across all of the items -- which smooths things out and stabilises the training.

Of course, it's possible to overdo it. As an extreme example, imagine that you were somehow able to fit your whole training set into one batch -- then you could train by running that single batch through, doing a single backward pass, and then adjusting the parameters once. It's pretty clear that that would not work very well -- just one single update of the initially-random parameters.

When training on my local machine, I could fit a batch of six sequences into my RTX 3090. I'd found that when I moved to cloud machines, it had a very positive effect on the loss I got out of the models when I tested them. From a quick-and-dirty bit of curve-fitting, I estimated that the optimal batch size for this model, with that training run, was somewhere around 97. Conveniently, that was close to the maximum I could fit onto an 8x A100 40 GiB/GPU machine, so I used a batch size of 96 to test the different interventions I was trying.

And when I finally put all of the interventions that helped with training together, I found (somewhat to my surprise) that their combined effect -- an improvement in loss of 0.113765 -- was less than half of the loss improvement of 0.252474 that I had got from increasing the batch size.

What that all made clear was that if I wanted to do a local training run that matched the quality of the cloud-trained model, I'd need to not only add on the interventions that I'd been testing in detail, but I'd need to match the cloud batch size. And for that, I needed to learn about gradient accumulation.

[ Read more ]

Writing an LLM from scratch, part 32j -- Interventions: trying to train a better model in the cloud

Since early February, I've been trying various interventions on a 163M-parameter GPT-2-style model that I trained from scratch on my local RTX 3090, using code based on Sebastian Raschka's book "Build a Large Language Model (from Scratch)".

My original model got a loss of 3.944 on my test set, while the original GPT-2 weights got 3.500 on the same dataset. I wanted to see if I could close that gap, and had a list of potential changes to the training setup, and to the model itself. Which of them would help?

I found a list of solid-looking interventions, and in my last post I came to the conclusion that the improvements in loss I had seen with all of them -- with two possible exceptions -- seemed unlikely to be in the noise. What would happen if I tried to put them into a new model?

[ Read more ]

Writing an LLM from scratch, part 32i -- Interventions: what is in the noise?

Towards the end of last year, I trained a 163M-parameter GPT-2-style model from scratch on my local RTX 3090, using code based on Sebastian Raschka's book "Build a Large Language Model (from Scratch)".

The result was a pretty decent little model, but it wasn't as good as the original GPT-2-small, despite having more parameters (because it wasn't using weight-tying). Specifically: on a particular test set, my model gave a loss of 3.944 -- quite a lot more than the original GPT-2's 3.500 on the same dataset.

I wanted to see whether I could train a model on my own hardware (or on something that didn't cost too much to rent in the cloud) that got closer to the original model's performance. So over the last few months, I've done a bunch of further training runs, each one testing a specific intervention -- a stand-alone change that I expected to change the loss, either for better or for worse. Specifically:

• I trained a baseline model on an 8x A100 40 GiB per GPU machine on Lambda (which was better than my original locally-trained model, I believe due to the larger batch size that the larger machine made possible).
• I tried adding gradient clipping to see if that would help by limiting the effects of loss spikes.
• I tried removing dropout, given that these days people tend not to use it (because we're doing single-epoch training runs).
• I tried adding bias to the attention weight matrices -- something that was popular back in the GPT-2 era, and was used by the original weights, but which my code did not use.
• Instead of just using the learning rate of 0.0004 that was used in the code from the book, I looked into what values people use these days, and learned how to schedule it over the course of the training run.
• Similarly, I learned more about weight decay and tried some alternative values.
• Then I tried making my model more like the original GPT-2 one by introducing weight tying to see if that would help.
• Finally, I decided to try training in "full-fat" float32 instead of using PyTorch's AMP and TF32 matrix multiplication performance enhancements.

At the end of all of that, I had this table showing the effect of each intervention in terms of loss on the test set. They're sorted from least-effective to most-effective, and you can see the baseline in there too:

Winners and losers are reasonably clear:

• Weight tying and the number for weight decay I derived from a paper by Cerebras Research (probably without understanding it properly) were negatives.
• Full-fat float32, gradient clipping, attention biases, the GPT-2 weight decay parameter, removing dropout, and scheduling (and updating) the learning rate were positives.

So, for an optimal train, we'd just use the effective interventions, right? Well, not quite.

Full-fat float32 I decided wasn't worth the effort, as it meant that the train took more than twice as long, and (because it required a larger machine), cost more than three times as much.

The others did look like solid changes, but there was one concern. The effect of each intervention is actually pretty small. For example, gradient clipping reduced the loss by 0.014, from 3.692 to 3.678. That's a 0.3% improvement. Even the best intervention, scheduling the learning rate, only improved things by 2%.

Could it be that some or all of these improvements were not real, but just a result of the random nature of training deep neural networks? Could the differences just be in the noise? They seemed small enough for that to be possible.

I've trained seven more models over the last few days to try to get a feel as to how big an effect noise has for this kind of training run. The results appear to show that variations in the initial weights matter quite a lot, but randomness in the training loop (given the same initial weights) actually has a fairly minimal impact. That surprised me a bit!

Let's go through the details.

[ Read more ]