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.