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?
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.