A quick(ish) Chinchilla check
I recently overtrained a couple of GPT-2 style models, training them both on 40 tokens per parameter rather than the 20 per parameter that is generally regarded as "Chinchilla-optimal".
The normal heuristic is that instead of doing that, you should scale up the number of tokens and the number of parameters equally -- so I would have been better off scaling up the model by and the token count by the same amount. By doing that, I should expect to get a better model in terms of loss on my held-back test set than I did with my 40-tokens-per-parameter models.
My training machine poppy wasn't doing anything, so I decided to give that a go.
Would the Chinchilla rule-of-thumb hold up?
As you might expect, it did. But it was a surprisingly close-run thing, and could conceivably have been in the noise. Let's take a look.
How I use AI on this blog
Inspired by this LessWrong post, I thought I'd write about how I use AI here. This is less in the interest of disclosure, more to provide a snapshot of what I'm doing right now so that I can revisit it in the future and see how it changes. And hey, maybe it'll be of interest to you, dear readers.
If I were to summarise my working philosophy in fewer than ten words, it would be: AIs identify problems and I fix them myself. With a very specific kind of exception (which I always flag), the text and code on this blog are human-generated. That's not a moral stand, but more a constraint imposed by what this blog is meant to be -- a place for me to learn in public.
Why do OpenAI's GPT-2 weights beat mine? Part three: testing overtraining
The GPT-2-style models that I've been training work really well, and I've even managed to train some that perform better than the original OpenAI small model in terms of cross entropy loss on a test set. But as I wrote previously, there's a mystery: why do they perform worse on my instruction fine-tuning evaluation?
I had various theories about why that might be, and to me, the most plausible-seeming of them was the amount of data they were trained with. As best I can find out, OpenAI's models were, by modern standards, trained on much more data than they should have been, while I'd used the theoretically optimal amount of training data.
To put it in other words, OpenAI's models were overtrained. If I deliberately overtrained my own models, could I match their performance? This post is a write-up of what happened, but so as not to bury the lede -- it didn't seem to help much, if at all.
Let's see why.
Why do OpenAI's GPT-2 weights beat mine? Part two: the bugfix
I'm digging into why my GPT-2 style models score worse on an instruction-following eval than OpenAI's original weights; I gave the details in this post.
While I was writing up the results of my first experiment into possible causes, I ran the post past ChatGPT -- I always use an "editorial board" of AIs to check my posts for flow, style, and any technical errors (though all writing is always mine). It took a look at the eval code that I was running, and highlighted a bug.
Luckily, it doesn't change the important results -- OpenAI's models continue to be better than mine at instruction-following. But it was enough to change the baseline numbers, re-ordering how well my own models did. So I fixed it and regenerated the baseline so that future experiments are based on solid ground.
Why do OpenAI's GPT-2 weights beat mine?
When I finished my project training an LLM from scratch, I was left with a minor mystery. Why were my models worse at instruction-following than the original OpenAI GPT-2 small weights?
I had an evaluation that I was running, based on the instruction fine-tuning code in chapter 7 of "Build a Large Language Model (from Scratch)". The process was to train a model on samples from the Alpaca instruction-following dataset until validation loss started rising, to use that instruction fine-tuned model to generate completions to a held-back test set, and then to use an LLM to compare the results from various different models. The details are here; let's call it the IFT eval.
OpenAI's original weights for GPT-2 small consistently beat my own models, even when mine got better results than theirs on a more technical evaluation, where I just measured the cross entropy loss for each model on a held-back set of test sequences. This surprised me; I would have expected a reasonably close correlation between the two evals -- that better test loss would imply better instruction-following.
I have a couple of thoughts about why this might be, and given that I recently
set up poppy, my dedicated LLM training box
I decided to inaugurate her with an experiment to test one of them; further
experiments will come in time -- though I don't think this will be a focus for the blog. More
of a running theme, with occasional posts until either I solve the mystery, or
give up in despair...
In this post I'll give a bit more detail about the nature of the problem, and list some of the things I've been thinking might be the cause. In later posts, I'll dig into some of them.
Benchmarking Qwen 3.6 35B MoE (3B active) on an RTX 3090
I mentioned I'd got a new RTX 3090 on a group chat, and a friend said:
I know this is not really your thing... but let me know how quickly it runs Qwen 3.6 35bn MoE. With only 24gb of VRAM you’ll need to use a 4-bit quantized version and you won’t get a massive context window. But it should still be pretty cool.
He's right that it's not really been my thing -- I've been focusing on my own LLMs recently. I decided to dig in a little, and in particular to play with Llama.cpp, which I haven't used for a while. And then things got a tad out of control, and I wound up doing some relatively detailed benchmarking.
The headline results: I downloaded Unsloth's UD-IQ4_NL_XL quantisation of the model
from Hugging Face. With that,
using the default Arch build of Llama.cpp, which uses Vulkan under the hood:
- Using the GPU only, I was able to get the model to generate at just over 120 tokens per second, and it was able to process the prompt at just less than 2,800 tok/s. However, having the whole model on the GPU didn't leave that much space for the context window -- it was constrained to about 50,000 tokens, compared to the model's native context length of 262,144.
- Offloading the FFNs for the first 12 of the model's 40 layers to the CPU managed to reclaim enough VRAM to be able to get the full context length; however, with that setup, things were -- unsurprisingly -- slower. I got just over 65 tok/s for generation, and 600 tok/s for the prompt.
Compiling Llama.cpp myself, in order to get the full CUDA version, helped a lot:
- With everything on the GPU, I got 140 tok/s for generation and over 3,300 tok/s for the prompt. That was with a context window of 89,600. So everything was better :-)
- It was also easier to get to the full context length; that needed just 10 layers' FFNs to be offloaded, and at that point I was getting 89 tok/s for generation, and about 1,100 tok/s for the prompt.
That was a pretty impressive improvement.
But it also showed that the lack of VRAM on the 3090 really does hurt. The friend who asked me about this is running an Intel Arc B70 Pro, with 32 GiB VRAM. He can (of course) fit the whole 4-bit quantised model on there without any context window reductions, and he says this about throughput:
It starts off at 75-80. but drops into the high 50s as the context window expands
That's worse than the CUDA-with-offload results above, though I did limit my testing to a 2,457-token prompt, with 6,144 tokens generated. The RTX 5090 has 32 GiB VRAM and fast Nvidia processing -- I imagine things would be a lot better there. Downside: it costs more than three times what you pay for a second-hand 3090 (and double the Arc B70 Pro).
Anyway, in the rest of this post, I'll give more details of the benchmark, including charts showing the performance at different numbers of offloaded layers, and comparisons of the Vulkan vs CUDA versions of Llama.cpp for this model.
One caveat: Qwen 3.6 is a multimodal model. That means that it has an extra component
to map graphical inputs to tokens that can be consumed by the LLM. I didn't realise this
as I ran these tests, but you can actually tell it not to use that component if you're only working with
text inputs -- for example, --no-mmproj will do the trick with Llama.cpp. If you're working with text-only stuff, you
might want to play with that to see if you can get better results -- it will free up
some VRAM and might allow you to get longer context lengths, or better performance.
But anyway, as with many of my posts, this is a tidied up version of my lab notes, so you can share my learning journey with me -- but if that sounds pointless, and you've landed here because you want to run this model on your own RTX 3090, you can click this link to jump right to the results.
Building intuition about LLM parameter counts
When I was building my GPT-2 implementation in JAX, I started with just token embeddings for the input, and a separate output head (as I was not using weight tying). It wasn't an LLM -- no Transformer blocks, no attention, no feed-forward networks.
I was somewhat surprised when I noticed that even that stripped-down model had 77 million parameters with the "small" settings I was using to train -- specifically, an embedding dimension of 768. However, I realised I shouldn't be -- with a vocab size of 50,257, each of those components is essentially a matrix, and that is indeed over 38 million numbers.
But the finished LLM at the end of the project was only 163 million parameters -- that meant that the input and output components alone were almost half of it. That felt like a surprisingly large percentage.
I had a similar shock when I was first looking into the feed-forward network, and realised that it had roughly twice as many parameters as the attention layers.
When we learn about the internals of LLMs, a lot of the focus is on the attention mechanism. This makes sense -- it's the hardest part to get your head around. The rest of the setup, at least for simple GPT-2 type models, is fairly standard stuff.
But that means that it is easy to overestimate how much of the total parameter count of the model attention uses up -- especially for smaller models, where the token embeddings and the output head are so large in comparison to the Transformer layers that make up the actual body of the LLM.
OpenAI released GPT 5.6 today, so I decided to take its "Sol" variant for a ride in Codex and asked it to write a visualiser. It shows breakdowns of how the parameters are split between embeddings, attention, the FFNs, and the output head for different sizes of GPT-2 models (or your own custom settings with the same architecture), and you can also add/remove weight tying and QKV bias. It did a really good job -- check it out! Here's a screenshot of what it showed for GPT-2 small without weight tying.
It's well worth a play. In particular, it's interesting to see what happens as the number of tokens in the vocab gets very large (many modern models have hundreds of thousands). You can very easily create a "tiny" model which is almost entirely embeddings and the output head.
poppy the training box, part 1: the beginnings
For a while I've been planning to put together a separate machine for local LLM
training. Until now, I've been using my desktop PC, perry. I have an RTX 3090
installed, and can get useful training runs done (most recently,
a 163M-parameter GPT-2 small style LLM in JAX),
but there are a couple of problems.
perryis my daily driver. If he's doing a training run, then everything is just a little bit sluggish as CPU and GPU alike are busy.- Although I don't play games often, it's annoying to have the option ruled out for days at a time.
- While the GPU is busy with a training run, I can't do other experiments in parallel -- for example, to scope out what the next step might be.
And relatedly to all of those: the two-day limit to the training runs I've been doing is something I set
because that's the maximum amount of time I'm willing to have perry tied up. It
would be really interesting to try longer training runs!
I also have longer-term plans; a multi-GPU box would be interesting to put together -- not just to have more power locally, but so that I could test larger-scale cloud multi-GPU training runs before starting to pay for expensive machines. US$15.92 an hour to rent a machine isn't a lot of money, but it adds up, especially if you're spending it while debugging parallelism issues.
And finally, I've always been interested in putting together a custom water-cooling loop in a PC. I've been building my own machines since 1995 or so, but never got round to that side of things. It sounds fun!
But despite all of those future plans, this is a fairly normal machine-building post -- how I repurposed an old PC, plugged in a second-hand RTX 3090 from eBay, tested it all, accidentally trained an LLM for 11 days, and almost cooked a CPU.
Over time, I expect to be posting more -- and more interesting -- build details. Let's think of this as establishing the baseline.
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.
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!