Moe Guide Debug

Cerebras Skip to main content

Cerebras Announces First Quarter 2026 Results >>

Aug 19 2025 Debugging Dead MoE Models: A Step-by-Step Guide Daria Soboleva Aman Tiwari

MoE Fundamentals | Router Wars | Debugging Dead MoE Models | MoE at Scale | MoE Math Demystified

This video shows a complete step-by-step walkthrough of training a small MoE model and debugging router issues. If you can’t open the video displayed above, please use this link to open it on YouTube: https://youtu.be/phXUzFt7hrs?si=6FAHkA00Tvpjz4LR What We Expect to See at the End So I bet when you hear Mixture-of-Experts (MoE), you immediately think “another thing that only Google can afford to train”, right? That’s exactly the myth we want to bust today. Yes, the famous MoE models are huge - we’re talking trillion parameter scale (Kimi Team, 2025). But this is like avoiding neural networks because GPT-4 exists. You can create a perceptron network from scratch in less than 20 lines of code. Unfortunately, it is a common myth that training MoE models is not really accessible to majority of people. In fact as we were working on this blog, this was the first question that Aman raised. Do we need a huge cluster of compute for this? But here is a counterintuitive fact - MoE is not just about scale, it’s about a technique that you can scale down to observe the same benefits. At the end of this blog you will learn how to train a small MoE model and beat a dense model at the GPT-2 scale. As for compute – you only need your laptop (okay, maybe a small acceleration node, but definitely not a datacenter!). Let’s dive in. Prerequisites If you are unfamiliar with dense GPT-2 style models, we highly recommend watching Let’s reproduce GPT-2 (124M) by Andrej Karpathy first and then get back to this blog where we build on top of it, specifically focusing on MoEs. Setup Let’s get ready for the big runs! We’re building on the modded-nanogpt codebase (Jordan et al., 2024a), which gives us a solid 124M GPT-2-like baseline trained on FineWeb dataset (Penedo et al., 2024). Our model is a decoder-only transformer which uses (So et al., 2022), RoPE (Su et al., 2023), RMSNorm (Zhang and Sennrich, 2019) and QK-Norm (Henry et al., 2020), trained using the Muon optimizer (Jordan et al., 2024b). It is basically a modernized GPT-2 that trains in 11 minutes. If you have access to Cerebras hardware you can follow along there, otherwise it’ll work fine on GPUs too. To get started, download the MoE fork of NanoGPT: After the above downloads the FineWeb training files, open train_gpt_moe.py in your favorite editor or IDE (this file contains our complete MoE implementation for those who want to run experiments immediately). Alternatively, if you prefer to code along and implement the MoE components step-by-step, open train_gpt.py instead. Also, everything can be scaled down (including the model size!) and we encourage you to play with it – consider this blog as your starting point. Break things, try weird settings, and see what happens. Where Experts Networks Live Alright, you’ve got your modded-nanogpt codebase pulled and you’re staring at the code. Here comes the exciting part - let’s turn our MoE concept introduced in (Soboleva, 2025a) into something you can actually train! The first step is pretty straightforward. We need to swap out single feedforward network (FFN) for multiple expert networks. Figure 1: Each expert is just a GPT-2’s FFN. You can find the expert class definition in train_gpt_moe.py#L184-L196 and the instantiation logic in train_gpt_moe.py#L263 . As we can see on Figure 1 each expert is literally GPT-2’s FFN copy-pasted (two linear layers separated with nonlinearity). Now, how do we use these expert networks? Figure 2 shows how we combine experts outputs when we’re processing batches of tokens, extending on the single-token approach that we introduced in (Soboleva, 2025a). Note that this implementation is not fast and wastes computation, but is useful for pedagogy. Figure 2: Expert mixing for batch processing (see train train_gpt_moe.py#L208-L225) . As was noted in the part 2 of this series MoE routing is tricky and can completely negate benefits that these models offer theoretically. In the next sections, we'll keep the experts architecture and mixing logic unchanged, while playing with routing solely and observing how it changes our model's quality. Fair warning: we are going to mess things up first and make mistakes, just like how it actually happens in real life. Learned Routing Learned routing is the industry standard and it powers all these great MoE models released by Google, OpenAI, Anthropic, etc. Now you’ll build it yourself! We’ve covered the theory in the part 2 and here we will focus on its actual implementation. Figure 3 shows that learned routing has two objectives: first, finding the best expert for each token through learnable weights, and second, keeping all experts utilized through load balancing. This dual objective is usually the place where we can shoot ourselves in the foot, so let’s find out what happens when we launch this training run. Plumb the auxiliary loss through the model, sum it across all layers, and add it to the cross-entropy loss with a weight of 0.01. Figure 3: Learned routing implementation: (top) router logic for token-expert assignments and (bottom) auxiliary loss for load balancing. Locate router definition in train_gpt_moe.py#L228-L240 , load balancing logic can be found in train_gpt_moe.py#L242-L250 . Time to check on the loss curves. Figure 4 shows that our MoE has already beaten the dense GPT-2 baseline. Congratulations! You implemented your first MoE model that works! But wait – we mentioned shooting ourselves in the foot, yet everything just worked. So why does this guide exist? You might have noticed that we promised you approximately a 2% improvement against the dense GPT-2 baseline with learned routing at 4 experts back in part 2 of our MoE 101 series. We are only getting half of that gain now. Where did the other 1% go? This is the moment that most of the researchers experience when implementing someone else’s paper and it doesn’t work as promised. In this guide we are going to fix it. Let’s debug this properly together. Figure 4: MoE with 4 experts and learned routing outperforms the dense GPT-2 model. Loss improvement widens throughout training, reaching a 1% by our target 2.4B tokens. To reproduce this plot, please run plot.ipynb . Remember, our router has two...