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.