Fun with low-rank vocab matrices (and a bonus test loss reduction?)

I was nerdsniped! On the HF discussion page for one of the models I created for my previous post, AndrewThompson1233 asked if I'd considered trying out factorised embeddings -- something he's using for his model, Maba, and which was previously used in some other models, like ALBERT.

It's a really nifty idea; you use a similar trick to LoRA as a way of reducing the number of parameters used for your embeddings and your output head. And for small models, those can be a disproportionate number of the total.

For example, with my 163 million parameter GPT-2 small-style models, the embeddings are about 39 million of them; the output head is the same size, so with those two taken together, that's half the model just getting stuff in and out rather than actually doing the thinking. Even if you use weight tying, like the original GPT-2 models did, you'll wind up spending 30% of your "budget" on the single matrix shared between embeddings and the output head.

Similarly, while larger models have a smaller percentage spent that way, even medium-sized MoE models can wind up spending a lot of the active parameters on embeddings; I calculated that for Qwen 3.6 35B MoE, with 3B active parameters, there were about 1B of them used across both the input embeddings and the output head.

So anything that can reduce that -- so long as it doesn't come at a high cost in terms of the model's capabilities -- is worth considering. You save on your parameter count, which either means smaller models and potentially quicker training, or allows you to "invest" the savings in more thinking parameters -- that is, a wider network with a higher embedding dimensionality, or a deeper one with more layers.

Andrew reported that the impact of using this trick was minimal:

At rank 128 on a 50k vocabulary, the penalty on validation cross-entropy is negligible (typically within +0.02 to +0.04 loss, or <0.5 perplexity delta).

Now, that suggests that he's getting a loss of around 2.5 (that being the point where a loss increase of 0.04 implies a perplexity change of 0.5), which is much lower than I typically get with my GPT-2 architecture (3.5 is more like it for me), but still, perhaps I'd get good numbers too? I decided to train a few models to see what happened.

The results were interesting! I found that using this trick on both the embeddings and the output head caused a somewhat larger increase in loss than Andrew described -- about 0.07 in my first test, 0.10 in the second, and 0.14 in the third. Those are relatively large, though potentially recoverable from further training or larger models.

But more surprisingly, I found that using the trick on the input embeddings only seemed to reduce test loss. Before I'd started, I'd expected it to be less harmful on the input side than it was on the output side, but seeing an improvement was certainly unexpected. Getting loss down by reducing your number of parameters is not what normally happens!

So, it's definitely worth digging in a bit. Let's start with the theory: what are these factorised embeddings, and how do they work? Like LoRA, they rely on low-rank factors, so firstly I'll define those.

[ 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 ]

Putting my JAX-trained models on the Hugging Face Hub

I hadn't uploaded the models that I trained using JAX to the Hugging Face Hub because Transformers has been PyTorch-only since version 5 (though they say they're working to add interoperability with JAX in the future), so it would have been tough to get them working natively with AutoModelForCausalLM and the like.

But then it dawned on me that I'd already written a conversion script that could take my JAX safetensors files and convert them into ones compatible with my PyTorch code. It's actually those converted models that I use for my evals -- so I could use my existing PyTorch script to upload them.

So, I've now uploaded PyTorch-compatible versions of all of my JAX-trained models:

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

• gpjt/jax-no-mha-bias-no-dropout -- the first full LLM trained in the post, in the "Adding LayerNorm" section.
• gpjt/jax-no-mha-bias-with-dropout -- the second full LLM trained in the post, in the "Dropout" section.
• gpjt/jax-with-mha-bias-no-dropout -- the third full LLM trained in the post, in the "Adding bias to the MHA output projections" section.

"Why do OpenAI's GPT-2 weights beat mine? Part three: testing overtraining"

• gpjt/jax-with-mha-bias-no-dropout-extended -- the single-epoch, double-Chinchilla-tokens model.
• gpjt/jax-with-mha-bias-no-dropout-2-epoch -- the model trained on two epochs over the Chinchilla-optimal number of tokens.

"A quick(ish) Chinchilla check"

• gpjt/jax-with-mha-bias-larger-chinchilla-1 -- the slightly-larger model.
• gpjt/jax-with-mha-bias-larger-chinchilla-2 -- the slightly-smaller model.

I've also added links to the posts in question.

Use the built-in GELU, don't roll your own!

Unsurprisingly, PyTorch's own built-in GELU function is faster than the hand-rolled one I've been using to date. But I was surprised at how much faster using it made things when training my models. I discovered this accidentally just now while working on something unrelated, but am logging the details here for anyone else that might find it useful.

The headline numbers: the same code, training the same model on the same data, ran at about:

• 21,000 tokens per second using the hand-rolled GELU from Sebastian Raschka's book "Build a Large Language Model (from Scratch)".
• 25,000 tokens per second using PyTorch's built-in GELU with no arguments.
• 25,000 tokens per second using the built-in GELU with approximate="tanh", which uses the same maths as Raschka's version under the hood.

That's a 20% increase in throughput for both of the built-in versions -- definitely nothing to be sneezed at.

And what is particularly interesting is that there aren't that many GELUs going on -- it's a GPT-2 small-style model, with 12 layers. So that's 12 GELUs handling tensors shaped (batch_size, seq_len, 4 * d_emb), which is (6, 1024, 3072) for my training setup. Given that the rest of the model is doing all of the normal full attention stuff for GPT-2, it's really surprising that the GELUs alone must have been taking up so much of the time. The throughput numbers mean that we must have been spending about 17% of our time on the extra overhead from the hand-rolled version, so that sets a lower bound for how much time the GELUs were taking up.

[ Read more ]

On first looking into JAX

Much have I travell'd in the realms of gold,
And many goodly states and kingdoms seen;
Round many western islands have I been
Which bards in fealty to Apollo hold.
Oft of one wide expanse had I been told
That deep-brow'd Homer ruled as his demesne;
Yet did I never breathe its pure serene
Till I heard Chapman speak out loud and bold:
Then felt I like some watcher of the skies
When a new planet swims into his ken;
Or like stout Cortez when with eagle eyes
He star'd at the Pacific -- and all his men
Look'd at each other with a wild surmise --
Silent, upon a peak in Darien.

John Keats, On First Looking into Chapman's Homer

I've been working with PyTorch quite a lot for the last couple of years, and feel like I've come to a reasonably solid understanding of how it all fits together. Working through Sebastian Raschka's book "Build a Large Language Model (from Scratch)", training my own LLMs locally and in the cloud, rebuilding Andrej Karpathy's 2015-vintage RNNs -- over time, it all adds up!

But, of course, there are other frameworks, and one I kept hearing about was JAX. While it's less dominant than PyTorch, it has a reputation for a certain cleanliness, a certain purity. And having spent time over the last couple of weeks working through the tutorials, and translating small PyTorch examples into it, I've been really impressed.

In this post I want to give an overview -- to report back to beginners like me, still living in PyTorch-land, on my new discovery. Less like Herschel discovering Uranus, and more like a 16th-century European coming back after having discovered something that the people who lived there were perfectly well aware of. What is this JAX thing, and how does it differ from PyTorch?

[ 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 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 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 ]

Writing an LLM from scratch, part 32h -- Interventions: full fat float32

This is the last of the interventions I'm trying out to see if I can improve the test loss for a from-scratch GPT-2 small base model, trained on code based on Sebastian Raschka's book "Build a Large Language Model (from Scratch)".

Back when I did my first training run for a base model, on my local RTX 3090, I used two optimisations:

• Setting the 32-bit floating point matrix multiplication precision to "high" rather than to "highest", which means that it uses lower-precision (but still technically 32-bit) TF32 for those operations rather than normal float32.
• Using PyTorch's Automated Mixed Precision (AMP), which allows it to use 16-bit calculations rather than 32-bit in places where it makes sense to do so.

The first of those boosted training speed from 12,599 tokens per second to 15,402 in my test harness, while AMP on its own boosted it to 19,921 tps (and also allowed me to increase the batch size from 5 to 6). Doing both appeared to hit some kind of diminishing returns -- it maxed out at 19,997 tps, only a little better than AMP on its own.

But intuitively, you'd expect that might come at a cost. While I'm sure the PyTorch developers have solid understanding of where switching to 16-bit will have a minimal impact on training quality, it seems too good to be true that it would have no impact at all.

Let's see what happens if we switch both of these optimisations off!

[ Read more ]

Writing an LLM from scratch, part 32g -- Interventions: weight tying

In Sebastian Raschka's book "Build a Large Language Model (from Scratch)", he writes that weight tying, while it reduces the parameter count of a model, in his experience makes it worse. As such, apparently people don't use it in modern LLMs. Intuitively, that makes sense -- I'll explain why in this post.

But as I'm trying various interventions to see if I can get my model -- based on Raschka's code, but trained for a fraction of the time that the original GPT-2 model was -- to perform as well as the original in terms of the loss it gets on a test set, I thought it would be worth seeing if it really is a negative for this particular tiny model of 163M parameters.

After all, the original weights use weight tying, and I did find that QKV bias appeared to help -- and that's another old-school technique that they used, which has since dropped out of fashion. Might this one help too?

Worth a try! Let's give it a go.

[ Read more ]