{"slug": "extending-raschka-s-gpt-2-an-moe-trained-from-scratch-on-an-rtx-3090", "title": "Extending Raschka's GPT-2: an MoE trained from scratch on an RTX 3090", "summary": "A developer extended Sebastian Raschka's GPT-2-style code from the book \"Build a Large Language Model (from Scratch)\" to add mixture-of-experts support and trained a 446M-parameter model with 220M active parameters from scratch on an RTX 3090, using 6 experts with 2 active per token. The training run took just under eight days, roughly four times longer than the developer's prior runs, and the resulting model beat the original OpenAI GPT-2 small (124M parameters) on the test set but fell short of GPT-2 medium (345M parameters), despite having more total parameters and fewer active ones. The write-up details the code added to GPT-2, including an auxiliary loss needed to ensure all experts are used during training.", "body_md": "Mixture-of-experts models are really nifty. You get inference speed close to\na small model's, with a lot of the smarts and knowledge of a large one. While they\nuse as much memory as an equivalently-sized dense (non-MoE) model, they're much faster.\nThe frontier labs don't publish their architectures, but Claude and ChatGPT are widely\nrumoured to be MoEs these days -- and certainly many large open-weights models like\n[DeepSeek](https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro-0813) and\n[Kimi K3](https://huggingface.co/moonshotai/Kimi-K3) are.\n\nIn this post, I'll show you how I added MoE support\nto the GPT-2-style code from\n[Sebastian Raschka](https://sebastianraschka.com/)'s book\n\"[Build a Large Language Model (from Scratch)](https://www.manning.com/books/build-a-large-language-model-from-scratch)\",\nthen used that to train a 446M-parameter model with 220M active parameters\nfrom scratch on my RTX 3090 -- essentially, GPT-2 small with 6 experts, 2 active per token.\nI wanted to understand how MoEs work, and as always, felt that the best way to do\nthat is to build one (and then to write it up like this). Hopefully because it's all\nfresh in my mind as I write this, I should be able to explain things clearly for others\nwho've finished Raschka's book.\n\nThe training run took just less than eight days, and the resulting model\ngot a better loss on my test set than any of the other models I've trained\nso far -- which was a good thing, given that it took [four times longer](/2025/12/llm-from-scratch-28-training-a-base-model-from-scratch) to train! It was also better\nthan the original OpenAI GPT-2 small (124M parameters), but not as good as GPT-2 medium\n(345M parameters), although it was close. That second result was interesting, as my\nmodel had more total parameters than GPT-2 medium, but fewer active ones.\n\nOn an instruction fine-tuning test, it did\nbetter than any of my other models, but worse than both OpenAI models (why OpenAI's models are\nso good on that particular test is a [mystery](/gpt-2-mysteries) I'm digging into separately).\n\nI'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.\n\nBefore we get into the weeds, though, let's start off with the basics; how do MoEs work in Transformers-based LLMs?\n\nThe phrase \"mixture of experts\" evokes an idea of a model which has separate parts that are knowledgeable in different domains. Maybe one part would know about coding, another about history, another about philosophy, and so on. You can imagine something that had separate LLMs for different topics, and routed incoming inputs appropriately.\n\nThat's actually not a bad design for a system -- Sakana.ai got a lot of interest for\ntheir [Fugu](https://sakana.ai/fugu/) system back in June of this year, and it works rather like that. But the \"experts\"\nin an MoE are at a much lower level, and (as with [so many things](/2025/05/llm-from-scratch-13-taking-stock-part-1-attention-heads-are-dumb)\nin LLMs) their expertise is in some weird and alien thing\nthat they determined was helpful for modelling language during their training.\n\nLet's look at how they fit in mechanically. The GPT-2-style LLM that we will start from -- the one Raschka describes in his book -- looks like this:\n\n**Diagram 1**: a GPT-2-style LLM at the top level\n\nThe MoE magic happens inside those Transformers layers, so let's zoom in on one of them:\n\n**Diagram 2**: a GPT-2-style Transformers block\n\nSpecifically, what we do to make our LLM an MoE is to replace that feed-forward network (FFN) with multiple separate FFNs -- our experts. Different context vectors get routed to different experts based on their contents.\n\nThe FFNs themselves -- in both the dense and MoE versions -- [are surprisingly simple](/2025/08/llm-from-scratch-17-the-feed-forward-network). In GPT-2, they\ntake in the incoming context vectors, run them through a normal linear layer\nto expand the number of dimensions by four, run the result through a GELU activation\nfunction, then project it back into the original incoming dimensionality with another\nlinear layer.\n\nIt's a really simple two-layer neural network, and at first glance seems somewhat arbitrary. Tutorials about LLMs tend to spend pages and pages explaining attention mechanisms, and pretty much gloss over the FFNs.\n\nBut the FFNs take up [twice as many parameters](/2026/07/llm-parameter-counts) as the attention mechanism in GPT-2.\nThey're clearly highly important, and my (very loose) metaphor for why that is,\nis that attention is how the LLM works out what to think about, and the FFNs are where\nit does its thinking. It's not a perfect match for what's going on, but I think it's\na decent working model for intuition.\n\nAn MoE leverages this. Instead of having just one FFN per Transformers block, we have multiple, and we use a subset of them for each context vector. We have what is called a router, or a gating network. It takes the incoming context vectors, and for each one, decides which of these FFNs -- which experts -- to use. Then we feed the input into the experts that were selected for it, combine their results, and that's our final output -- like this:\n\n**Diagram 3**: an outline of an MoE block\n\nDoing this gains us more \"space\" in the LLM for it to remember facts and ways of thinking about things -- we have multiple experts for that knowledge to be spread over. We could, of course, do that by dedicating more space to the FFNs -- for example, by having one bigger one. But with an MoE, because we only activate a subset of the experts for each context vector, we save on the amount of computing we do for each context vector. We need to keep all of the experts in RAM -- remember that we're routing to them per-context vector, so in a batch of sequences we're likely to be using most if not all of them. But we don't have to feed everything through all of them.\n\nSo, that's the basics -- nothing conceptually difficult. What becomes more tricky is the implementation. How, concretely, does the router choose which experts to use for a given context vector, how do we implement that choice -- and how can we do all of that efficiently? And how do we train the router to make its choices?\n\nWhen I started on this project, my first step was a Google search for useful references.\nI came across this [excellent summary](https://www.ibm.com/think/topics/mixture-of-experts)\nfrom IBM.\n\nIn that post, they mention a number of papers; four seemed relevant [1](#fn-1):\n\nI decided to take a look at them, and was pleasantly surprised about how readable they all were. It looked to me like I'd be able to put together a working MoE model by following what they (and the IBM summary) said, and that turned out to be true -- the model worked, and trained well. Of the four papers, I found the Switch Transformers one the most useful, but the original 1991 paper also clarified a lot to me. I really do recommend skimming through them if you want to learn more background about this stuff.\n\nAnyway, once I'd got\nit all put together, and trained the model,\nI compared what I'd written to\n[the Hugging Face source code](https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py)\nfor [Mixtral](https://mistral.ai/news/mixtral-of-experts/), which was the first open\nMoE model I remember hearing about. Mixtral has a bunch of other improvements over GPT-2 beyond its MoE\nsupport (RMSNorm, RoPE, and so on), and it's a much larger model than mine. But\nit turns out that the specific way\nit handles the MoE side of things is largely the same as mine, apart from one\ndifference in how it calculates auxiliary loss (about which more later). So that\nwas reassuring.\n\nI also ran the final codebase past three LLMs -- ChatGPT, Claude, and Kimi K3 -- just to make sure that I hadn't drifted from the mainstream or screwed up in other ways. I did this in an anonymous/private session so that they wouldn't use any memories they have of conversations we'd had in the past, and asked them\n\nPlease take a look at the attached and tell me what you see. I'm particularly interested in the MoE stuff.\n\nAll of them came back with some variation of \"it's a pretty standard MoE implementation on top of GPT-2, with load-balancing adapted from Switch Transformers\". Even more reassuring!\n\nSo I'm comfortable that what I've built is a normal implementation, and that's what I'll describe in the rest of this post. It's all fresh in my mind, so hopefully by describing it in terms that would have made sense to me when I just finished Raschka's book, the result will be useful to other people coming to this for the first time.\n\nLet's get started by digging into the mechanics of how the router works.\n\nThe problem we want to solve is that we have a context vector coming into our MoE\nblock -- as per [the diagram above](#diagram-3) -- and we want to decide which experts we want to\nroute it to. Let's say that we have total experts, and we want to send\neach input to of them (where , of course). We'll also say that our context\nvectors are of dimensionality .\n\nTaking a -sized input and categorising it into probabilities across a number of options is a pretty standard job for a neural network. In fact, we can use a single linear layer for it. Imagine that we create one with inputs and outputs. For a given context vector, it might produce (for ) something like this:\n\n```\nIn [5]: router(inputs)\nOut[5]:\ntensor([ 0.0418, -0.1140, 0.4254, 0.1342, 0.5106, -0.1385],\n grad_fn=)\n```\n\nSo we can imagine picking out the indexes of the top of those.\n\nLet's be specific, and say that . That means that we have indexes 4 and 2 -- the positions of the two largest numbers -- so we want to route our original context vector to experts 4 and 2. They do their calculations, we get output context vectors from them, and we can combine them.\n\nHow might we combine them? Well, we do a lot of adding context vectors together in our GPT-2-style LLM code, and treat the results as meaningful -- token embeddings get added to position embeddings, shortcuts around attention and the FFN get added back in to their results, and so on. So perhaps we could do that?\n\nThat kind of setup has a problem, though -- it doesn't train the router.\n\nLet's think about the MoE block again; here's the diagram again:\n\n**Diagram 3** (repeated): an outline of an MoE block\n\nThink about the flow of data through it. The context vectors flow into the router. We use the output of the router to select the two experts, and then the original context vectors -- the ones that were fed into the router -- flow into those experts, then are combined, and we get our result.\n\nNow, consider what happens when you're training a model. You run some training data through it, then calculate the loss -- how good your model's results for that data are. You then use that to work out the gradients that you want to apply to the parameters to make it better. You work out the gradients by using back-propagation, working back from the loss, through the network, retracing the computation graph in reverse.\n\nWhen your backprop gets to this part of the model, it will start with the output context vectors, trace back through the combination step, then back through the two chosen experts, then back to the input context vectors -- and then it will go back to whatever step came before the MoE block.\n\nThe calculations inside the router that selected our two experts did actually happen in our forward pass -- but they're not in the computation graph as we trace it backwards from the loss. It's kind of a dead-end. There's nothing in there to connect our selection of which experts we used to the path through the computation that ended up at the loss. (Interestingly, the effort of doing that diagram made that particularly clear to me -- that slightly-messy labelling of the arrows at the top, with numbers to show the sequence, is a direct result of the same issue.)\n\nWhat that all means is that the router will not be trained. The backward pass of our training will completely ignore it, and so there will be no gradients to apply to it. A starting model with random weights will randomly allocate context vectors to experts, and it will continue to do that.\n\nClearly, we need to somehow modify our computation graph so that the router is connected -- so that when the backward pass flows up from the output context vectors, it sees the router -- and ideally sees some aspect of it that will allow it to be trained to do its job better.\n\n\"Adaptive Mixtures of Local Experts\" starts by describing some previous work which treats the router's outputs as weights for each of the experts. That's a nice trick, and while it does have some problems (as we'll see in a bit), it fixes the backprop issue. (Interestingly, they actually decided to do things differently -- but later work, like Switch Transformers, goes back to doing things this way, with some important tweaks.)\n\nAs I said earlier, back in 1991 they weren't thinking of MoEs as being a way to save on computation. Instead, their focus was on training better models to handle specific tasks, and they felt that having some way to split those tasks up might make that easier.\n\nSo, for their case, they might take those outputs from above and normalise them:\n\n```\nIn [7]: logits\nOut[7]:\ntensor([ 0.0418, -0.1140, 0.4254, 0.1342, 0.5106, -0.1385],\n grad_fn=)\nIn [8]: torch.softmax(logits, dim=-1)\nOut[8]:\ntensor([0.1459, 0.1249, 0.2141, 0.1600, 0.2332, 0.1218],\n grad_fn=)\n```\n\n...and then run their inputs through all six experts, then add together the outputs weighted by those numbers -- 0.1459 times expert 0's output, 0.1249 times expert 1's, and so on. Rather like this:\n\n**Diagram 4**: An MoE with all experts active\n\nWith that, we have something where the router is on the backward pass through the computation graph from the output context vectors (and thus from the loss). The weights provide a route back, so the router will get trained.\n\nBut that version, of course, is running all of the experts for every token, so it doesn't have the computational benefit of a modern sparse MoE.\n\nIt also has an issue, as they point out in the paper, that if you imagine that you want each expert to have a well-defined responsibility, things get messy. Imagine that expert 2 in the above example was the one that knew how to solve a particular task; all of the other experts would be contributing too, and unless you find some way to train their weights down to exactly zero for that task, to get a good loss on your training run they'd need to learn to balance each other out.\n\nTheir solution was to run all of the experts then to use a gating function on the output -- that is, it would drop all of the results from the non-selected experts (see their figure 1). But the modern solution is a bit different; for that, let's move on to \"Outrageously Large Neural Networks\".\n\nWhat we really want is an output from the router where some of the weights are zero. Specifically, with our active experts, total, we want all but the top of them to have a weight of zero. Then when we're running the network, we can skip those zero-weighted experts and get a result that will be identical to what we would have got without skipping them.\n\n\"Outrageously Large Neural Networks\" does this with a trick that will be familiar from\nthe [causal mask](/2025/03/llm-from-scratch-9-causal-attention) in the GPT-2 code.\nLet's call the \"raw\" weights from our router `logits`:\n\n```\nIn [7]: logits\nOut[7]:\ntensor([ 0.0418, -0.1140, 0.4254, 0.1342, 0.5106, -0.1385],\n grad_fn=)\n```\n\nLet's run that through softmax again:\n\n```\nIn [9]: torch.softmax(logits, dim=-1)\nOut[9]:\ntensor([0.1459, 0.1249, 0.2141, 0.1600, 0.2332, 0.1218],\n grad_fn=)\n```\n\nThat gives us some initial weights. But we want all but the top to be zero.\n\nNow, if we want something to be zero after softmax, then it needs to be on\nthe way in. I'll show the code to do this later, but for now, let's just assume that\nwe have some magic to set all but the top- values to that. In our\nconcrete example with , the original `logits` with\nthat change applied would be:\n\n```\nIn [15]: masked_top_k_logits\nOut[15]:\ntensor([ -inf, -inf, 0.4254, -inf, 0.5106, -inf],\n grad_fn=)\n```\n\nNow we can run that through softmax:\n\n```\nIn [17]: torch.softmax(masked_top_k_logits, dim=-1)\nOut[17]:\ntensor([0.0000, 0.0000, 0.4787, 0.0000, 0.5213, 0.0000],\n grad_fn=)\n```\n\nAnd we have some weights! With those, we can conceptually run this \"all-experts-active\" kind of network:\n\n**Diagram 4** (repeated): An MoE with all experts active\n\n...but skip the experts whose weights are zero. The weights that we provide through the steps above will mean that the router will get trained.\n\nThat's pretty neat!\n\nThere is one thing to highlight, though. Imagine if we have one active expert -- that is, . We'd mask out all of the other ones:\n\n```\nIn [15]: masked_top_k_logits\nOut[15]:\ntensor([ -inf, -inf, -inf, -inf, 0.5106, -inf],\n grad_fn=)\n```\n\n...and softmax:\n\n```\nIn [17]: torch.softmax(masked_top_k_logits, dim=-1)\nOut[17]:\ntensor([0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000],\n grad_fn=)\n```\n\nThe weight will always be one, regardless of the original logits. And if a function can only ever return the same result, its derivative will always be zero, so there will be no gradients and it can't be trained.\n\nInterestingly, that's something that is -- almost silently -- addressed in the Switch Transformers paper. They are specifically looking at MoEs, and they mention the \"Outrageously Large Neural Networks\" paper's routing system, then present their own calculations which instead of replacing non-top- values with , and then softmax, do the softmax first, then zero out the non-top-. They don't highlight the difference.\n\nI decided that for this experiment, I'd go ahead with the non-Switch Transformers\ncalculations, anyway, and do the replace-with--then-softmax system, with some\nguard code to prevent it from accidentally being set to . Most\nrecent MoE models I've seen have at least two active experts, after all.\nThe good news is that not only did it work -- when\nI checked later, it also\nmatched the choice taken in Mixtral, which feels like a solid endorsement. [2](#fn-2)[3](#fn-3)\n\nSo, at this point, we know how an MoE works in theory. Let's start coding.\n\nI started with the code that I had from\n\"[Build a Large Language Model (from Scratch)](https://www.manning.com/books/build-a-large-language-model-from-scratch)\".\nIn that, the Transformers block looked like this:\n\n``` python\nclass TransformersBlock(nn.Module):\n\n def __init__(self, cfg):\n super().__init__()\n self.att = MultiHeadAttention(\n d_in=cfg[\"emb_dim\"],\n d_out=cfg[\"emb_dim\"],\n context_length=cfg[\"context_length\"],\n num_heads=cfg[\"n_heads\"],\n dropout=cfg[\"drop_rate\"],\n qkv_bias=cfg[\"qkv_bias\"],\n )\n self.ff = FeedForward(cfg)\n self.norm1 = LayerNorm(cfg[\"emb_dim\"])\n self.norm2 = LayerNorm(cfg[\"emb_dim\"])\n self.drop_shortcut = nn.Dropout(cfg[\"drop_rate\"])\n\n def forward(self, x):\n shortcut = x\n x = self.norm1(x)\n x = self.att(x)\n x = self.drop_shortcut(x)\n x = x + shortcut\n\n shortcut = x\n x = self.norm2(x)\n x = self.ff(x)\n x = self.drop_shortcut(x)\n x = x + shortcut\n\n return x\n```\n\nI wanted to keep the capability to run this in MoE or non-MoE mode, and decided\nthat I'd do that by extending the model config in `cfg` so that it would have an\noptional `moe` section, which would include MoE-specific stuff. If that wasn't present,\nthen I'd create a normal dense LLM. To do that, I replaced the line in `__init__`\nthat assigned to `self.ff` with this:\n\n```\n self.is_moe = \"moe\" in cfg\n if self.is_moe:\n self.ff = MixtureOfExperts(cfg)\n else:\n self.ff = FeedForward(cfg)\n```\n\nNext, it was time to implement the `MixtureOfExperts` class itself. The obvious\nconfig that it would need was how many experts there were in total, and how many\nwere active per token:\n\n``` python\nclass MixtureOfExperts(nn.Module):\n\n def __init__(self, cfg):\n super().__init__()\n self.num_experts = cfg[\"moe\"][\"num_experts\"]\n self.num_active_experts = cfg[\"moe\"][\"num_active_experts\"]\n```\n\nNow, there's that issue where having just one active expert will pin the softmaxed weight to one, so it won't train -- so I decided to protect against that (and against another obvious mistake that the config could contain):\n\n```\n if self.num_active_experts < 2:\n raise Exception(\n f\"Can't train with ``num_active_experts`` < 2 (got {self.num_active_experts})\"\n )\n if self.num_active_experts > self.num_experts:\n raise Exception(\n f\"{self.num_active_experts=} is larger than {self.num_experts=}\"\n )\n```\n\nNext, it was time to create our router, mapping from the incoming context vectors\n(of size `cfg[\"emb_dim\"]` in this code) to the number of experts:\n\n```\n self.router = nn.Linear(cfg[\"emb_dim\"], self.num_experts, bias=False)\n```\n\nI decided to make it unbiased because I have a vague impression that that is the fashion these days -- nothing more principled than that :-)\n\nNext, we needed the `num_experts` experts themselves, each of which would be one of\nthe same `FeedForward` modules as we were using in non-MoE mode:\n\n```\n self.experts = nn.ModuleList([\n FeedForward(cfg) for _ in range(self.num_experts)\n ])\n```\n\nThey needed to go into an `nn.ModuleList` because if they were just in a regular\nlist (as I discovered when I tried it) they would not be registered by PyTorch as\nthings containing parameters belonging to the module. We create our optimiser\nfor training with code like this:\n\n```\n optimizer = torch.optim.AdamW(\n model.parameters(),\n lr=learning_rate, weight_decay=weight_decay\n )\n```\n\n...and so anything that isn't in `model.parameters()` will never get updated, which\nwould be a Bad Thing.\n\nThat was enough to have the pieces in place. It was time to write the forward pass -- to get the router logits, do the top-, softmax, and then to use those results to run the experts.\n\nGetting the logits was simple enough. Looking at the `forward` method:\n\n``` python\n def forward(self, xs):\n```\n\n...we have an incoming set of context vectors, `xs`.\nThat is shaped `(batch_size, sequence_length, d_emb)`.\n\nLet's try to visualise that. Understanding tensor operations is something you can kind of short-cut with intuition, but I think that in order to really understand the next steps it's best to have something a bit more tangible in mind.\n\nYou can think of an order-3 tensor like this as a cuboid. With `batch_size` of 3,\n`sequence_length` of 5, and `d_emb` of 7 -- artificially small values to keep things simple -- it might look like this:\n\n**Diagram 5**: The `xs` tensor as a 3-D cuboid\n\nEvery dot in that cuboid is a single number. A single context vector is the numbers\nas you follow a line from the \"front\" of the cube -- the `batch_size` `sequence_length`\nface to the left -- to the \"back\".\n\nThat 3-D representation was fiddly to get right, and is likely to get confusing if\nwe keep using it. So instead, let's look at two 2-D views of the same thing,\nfrom the front (where you are looking at a `batch_size` `sequence_length`\nface), and from one of the sides, where it's `batch_size` `d_emb`:\n\n**Diagram 6**: `xs` as two 2-D views\n\nSo in this view, each of the circles in the \"Front view\" is the first number in a specific context vector, and each of the rows in the \"Side view\" is the set of numbers that make up a specific context vector. Hopefully that's easy to visualise.\n\nNow, a PyTorch `nn.Linear` layer like our `self.router` operates on the last dimension\nin the tensor you pass in. For our `xs`, shaped `(batch_size, sequence_length, d_emb)`,\nit will work on the `d_emb` dimension -- that is, it will operate on each context\nvector independently, which is what we want.\n\n```\n routing_logits = self.router(xs)\n```\n\nThat gives us a set of logits shaped `(batch_size, sequence_length, num_experts)`.\nIn our visualisation, that's a `batch_size` `sequence_length` `num_experts`\ncuboid; let's diagram that with four available experts: `num_experts = 4`:\n\n**Diagram 7**: `routing_logits` as two 2-D views\n\nAgain, if we look at the front view, each circle represents the first element of the\nrouting logits for one of the context vectors, and in the side view, each row represents\nall of the routing logits for a context vector. So we have the right data in our\ncuboid. From what we're calling the front view, it's got the same dimensions as `xs`, which\nis useful.\n\nIn order to work out which of our\nexperts each context vector should go through, we need to get the top- values\nfor each context vector in `routing_logits`. PyTorch has\na [`topk` function](https://docs.pytorch.org/docs/2.14/generated/torch.topk.html) to do exactly that:\n\n```\n top_k_values, top_k_indices = torch.topk(\n routing_logits,\n k=self.num_active_experts,\n dim=-1\n )\n```\n\nThe `dim=-1` tells it to work across the last dimension, which is the\ndimension of length `num_experts`. It returns the top `self.num_active_experts` values\nand their positions, as tensors -- both shaped `(batch_size, sequence_length, num_active_experts)`.\n\nSo, we have two new tensors, which we can visualise as cuboids, both of `batch_size` `sequence_length` `num_active_experts`.\nLet's show that with `num_active_experts = 2`:\n\n**Diagram 8**: `top_k_values` (or, equivalently, `top_k_indices`) as two 2-D views\n\nBoth `top_k_values` and `top_k_indices` are shaped that way. As normal, the\n\"front\" face is the same; each circle corresponds to a context vector, and for\n`top_k_values` it holds the highest value in that context vector's logits (because\n`topk` returns results sorted), while\nfor `top_k_indices` it holds the index that that highest value sits at in the logits\nlist. Looking at the side view, we're seeing a list of top-k logit values or indices\nfor a given context vector.\n\nNow, we want a version of `routing_logits` that we can run through softmax to get\nsome weights, and in order to do that we need to replace the non-top- values\nwith for each of the context vectors.\n\nThe solution that I hit on eventually (after various other attempts [4](#fn-4)) was this.\nLet's start with a tensor identical in size to `routing_logits`, but\nfull of s:\n\n```\n top_k_routing_logits = torch.full_like(routing_logits, -torch.inf)\n```\n\nNow, `top_k_values` contains the values that we want to have in there, and `top_k_indices`\ncontains the indices in the logits lists where they should go. All of the other values\ncan remain as .\n\nThis will, of course, be the same shape as `routing_logits`:\n\n**Diagram 9**: `top_k_routing_logits` as two 2-D views\n\nBoth `top_k_values` and `top_k_indices` are compatible in their\nfirst two dimensions with `routing_logits`, and thus with `top_k_routing_logits` --\nthat is, their front faces in our visualisations are the same -- compare the front\nface of the above with the one for [diagram 8](#diagram-8).\n\nFor all of our tensors, the first two dimensions correspond ultimately to an incoming\ncontext vector in `xs`. It's just the last dimension and the data they contain that differ.\n\nPyTorch has a [`scatter_`](https://docs.pytorch.org/docs/2.14/generated/torch.Tensor.scatter_.html)\nfunction on its `Tensor` class, which takes a list of positions and list of values.\nIt takes one specified dimension, and treats that one essentially as a set of lists.\nThen it takes two other tensors with the same number of dimensions as each other, each of which matches in size in all but the specified\ndimension, and treats the other dimension as being a list of values for one parameter,\nand list of indices for the other. It overwrites the data at the specified indexes\nwith the specified values.\n\nThat sounds useful! Let's make it concrete. Remember that in all of these\ncuboids we're visualising, the front view we're looking at corresponds ultimately to a context\nvector. In `routing_logits`, it's the raw logits for expert routing for that vector --\nthe number on the front face is the first of those (corresponding to the logits for\nrouting to the first expert). If we look at the side view, each row would correspond\nto the set of logits for a given context vector.\n\nNow, let's consider just that. Inside `routing_logits`, our specific context vector\nmight have logit values corresponding to it like this:\n\n**Diagram 10**: A single CV's routing logits in-place in the cuboid\n\nVisualised the same way, with the equivalent part of `top_k_indices` would look like this:\n\n**Diagram 11**: A single CV's top-k routing logits indices in-place in the cuboid\n\nThat is, the `topk` function identified that index 2 in the original logits was the\nhighest value, and 0 was the second-highest.\n\nLikewise, `top_k_values` would have this:\n\n**Diagram 12**: A single CV's top-k routing logits values in-place in the cuboid\n\nThese diagrams are getting a bit unwieldy; let's look at this specific context vector's\ndata as lists. For our\nselected context vector, we have these logits from [diagram 10](#diagram-10):\n\n```\n[0.4254, 0.1342, 0.5106, -0.1385]\n```\n\n...these top- indices from [diagram 11](#diagram-11):\n\n```\n[2, 0]\n```\n\n...and these values from [diagram 12](#diagram-12):\n\n```\n[0.5106, 0.4254]\n```\n\nOur `top_k_routing_logits` tensor is just full of s, and is the same shape\nas `routing_logits`, so its corresponding part is this:\n\n```\n[-inf, -inf, -inf, -inf]\n```\n\nWhat `scatter_` will do is select indices 2 and 0 (because of the values in `top_k_indices`),\nand copy the corresponding numbers from `top_k_values` on top of whatever is already there:\n\n```\n[0.4254, -inf, 0.5106, -inf]\n```\n\nIt will do that for every one of the positions in that front view. So now we have what we wanted: a grid of routing logits for each context vector, where the non-top- ones have been replaced by .\n\nThat's pretty nifty! And so here's the code to do it:\n\n```\ntop_k_routing_logits.scatter_(dim=-1, index=top_k_indices, src=top_k_values)\n```\n\nI was initially a little worried that the whole \"replace the logits with a 'static' tensor and then scatter the values in there\" approach might break the computation graph, but tests showed that it didn't. My intuition is that because the s were summoned out of nowhere, they are dead ends for backprop, but because the logits that we're copying in come from previous calculations, they are not.\n\nSo once we've done that, we have our top- logits, ready for a softmax -- so it's time to do that:\n\n```\n expert_weights = torch.softmax(top_k_routing_logits, dim=-1)\n```\n\n...and we have our weights, yet another cuboid like this:\n\n**Diagram 13**: `expert_weights` as two 2-D views\n\nJust as before, a given number on our front face is the first of the expert weights for a specific context vector, and each row in the side view is all of the expert weights for that context vector across all experts. Post-softmax, the weights for active experts for that context vector are positive numbers, and the weights for the inactive ones are zero.\n\nThe next step was to actually feed the context vectors into their experts. I decided to keep this simple, and just iterate over the experts, one by one, and for each one to find which incoming context vectors wanted to go to it, run them all through, and then reassemble the results.\n\nI believe that larger MoE systems have\nsmarter routing systems -- for example, for a huge model that can't fit on a single\nGPU you might have different experts on different GPUs or even machines, and route to\nthem in parallel. But for my toy-sized models, this felt simplest, and felt like it would be efficient enough [5](#fn-5).\n\nThe way I decided to do this was to use an \"accumulator\" model. We know that\nthe shape of the outputs is the same as the shape of the inputs -- that is, if we have\nour incoming `xs` shaped `(batch_size, sequence_length, d_emb)`, then the output\nwill have the same shape. So I started off by creating a tensor of zeros of that\nshape:\n\n```\n all_outputs = torch.zeros_like(xs)\n```\n\nSo, just like `xs` in [diagram 6](#diagram-6), it will look like this:\n\n**Diagram 14**: `all_outputs` as two 2-D views\n\n...but it would be filled with zeros in every position.\n\nThe plan was that each expert would be run on the appropriate context vectors, its outputs\nwould be scaled by the weights that were calculated by the router and its associated\ntop- and softmax for that context vector/expert pair, and then the results could\nbe added into `all_outputs`. That would give us the result we wanted: after all experts\nhad been run on their respective context vectors, `all_outputs` would have the weighted\nsums.\n\nSo the next step was to iterate over the experts:\n\n```\n for expert_ix, expert in enumerate(self.experts):\n```\n\nNow, we need to know which context vectors wanted to be fed to this expert.\n\nLet's take a look at our representation of `expert_weights` again:\n\n**Diagram 13** (repeated): `expert_weights` as two 2-D views\n\nWe can see it as a bunch of `num_experts` \"slices\", each the same shape as that\nfront face, `(batch_size, sequence_length)`, where each one is the weights for a given expert.\nThe front face itself (corresponding to the rightmost column on the side view) is\nthe weights for expert 0, the next one \"back\" is for expert 1, and so on.\n\nAnd in code, we\ncan get the slice for the expert with index `expert_ix` like this:\n\n```\nexpert_weights[:, :, expert_ix]\n```\n\nThat will be a simple 2-D grid of numbers like this:\n\n**Diagram 15**: `expert_weights` sliced for one expert\n\nYou can see that it's a grid of one number for each context vector in our input --\nthe same shape as the front face in all of the diagrams so far. Each number is the\nweight that the expert with index `expert_ix` has for the context vector in question.\n\nWe can now do a comparison:\n\n```\n this_expert_mask = expert_weights[:, :, expert_ix] > 0\n```\n\nThat will give us a new grid of the same shape, but the numbers have been replaced\nwith booleans -- `True` if the corresponding context vector has a weight for this\nexpert that is greater than zero, `False` otherwise. We've got a mask that identifies\nexactly the context vectors that we want to run through this expert.\n\nLet's imagine it looks like this for some particular set of context vectors and some specific expert; I've coloured in the circles representing\n`True` and left the `False` ones white.\n\n**Diagram 16**: what `this_expert_mask` might look like for one expert\n\nNow comes the clever part :-) Remember that our original incoming context vectors,\nin `xs`, looked like this:\n\n**Diagram 6** (repeated): `xs` as two 2-D views\n\nWe can use our mask -- the grid of `True` s and `False` s in [diagram 16](#diagram-16) -- to\nselect a subset of the context vectors in there, like this:\n\n```\nxs[this_expert_mask]\n```\n\nThat will return us the subset of the incoming context vectors that we want to run through this expert. In terms of our diagrams above, it will be selecting the context vectors in the front face that are \"selected\" by our mask, and taking the \"cores\" as it goes back through the cuboid from there.\n\nThe question is, what shape will it be? You can see that there's no simple 3-D shape it could be. The first sequence in our batch -- the first row on the front face -- has two selected context vectors, while the second has one, and the third three.\n\nYou could imagine a world where the output would be the same shape as the input -- that\nis, a tensor the same shape as `xs` -- but with the non-selected numbers replaced with `None` s\nor something like that. But instead, PyTorch produces what amounts to a\nlist of the selected context vectors for this expert. We'll call the number of selected\nvectors `num_selected_context_vectors_for_this_expert`,\nand it's six in the example diagram above, so the shape will be `(num_selected_context_vectors_for_this_expert, d_emb)`,\nlike this:\n\n**Diagram 17**: the \"selected\" context vectors\n\nNow, note that this is a big change from all of our tensors so far. It has lost the\nconnection back to the original `xs` tensor's shape. In all of the other tensors, we could\nmap something back to the original context vectors in `xs`. But with this new\none, we have a bunch of context vectors with no inherent connection to where they\noriginally came from in `xs`. We'll come back to that.\n\nBut for now, we have the data that we want to run through the expert with index\n`expert_ix`. The expert is an FFN -- our two linear layers with a GELU in between them --\nand will treat all but the last dimension in whatever we feed it as \"batch\" dimensions,\nso it will work over that `d_emb` dimension, as we want it to. So the code to\nactually select the context vectors -- our `xs[this_expert_mask]` above -- and then\nto run it through the expert itself just becomes this:\n\n```\n this_expert_results, _ = expert(xs[this_expert_mask])\n```\n\nWe get the results of running our selected context vectors through the expert,\nstill shaped `(num_selected_context_vectors_for_this_expert, d_emb)`.\n\n(If you're familiar with the GPT-2 code, that `_` in the code might seem odd. We'll come back\nto why the expert is returning a tuple and why we are ignoring the second item in\nit later.)\n\nSo we have our results -- we've run the appropriate context vectors for this expert through it. But they're in a slightly funny shape, and we're going to need to fix that later on. But first, we need to apply the appropriate weights for this expert to each result.\n\nRemember that `this_expert_mask` is a grid of booleans, `(batch_size, sequence_length)`, like this:\n\n**Diagram 16** (repeated): what `this_expert_mask` might look like for one expert\n\n...where `True` -- filled in the diagram -- means that this expert is active for the corresponding context vector, and\n`False` means it isn't.\n\n`expert_weights` looks like this:\n\nNow, previously we did this:\n\n```\nxs[this_expert_mask]\n```\n\n...to pluck out the context vectors from `xs` where the mask was true. If we were to\ndo\n\n```\nexpert_weights[this_expert_mask]\n```\n\n...then we would be doing something similar. We'd get a grid like the one of\ncontext vectors in [diagram 16](#diagram-16), with one row for every selected context\nvector for this expert, except that instead of each row containing\na context vector, it would contain that context vector's per-expert weights:\n\n**Diagram 18**: all expert weights for the \"selected\" context vectors\n\nNow, on its own, that's not particularly useful -- but you can hopefully see that\ncolumn `this_expert_ix` is the weights for our expert for all of the context\nvectors that are going to be run through it! So if we do the same masked lookup as\nbefore, but also select that column, we get code like this:\n\n```\nexpert_weights[this_expert_mask, expert_ix]\n```\n\nWhat we're saying is \"pick each item in `expert_weights` that matches a `True` in\n`this_expert_mask`, then take the `expert_ix` th element of it\".\n\nIt will look like this:\n\n**Diagram 19**: this expert's weights for the \"selected\" context vectors\n\nThat is exactly\nwhat we need to get the weights for our results! We need to multiply the th\nelement in `this_expert_results` -- an output context vector that is the results for\na particular input context vector -- with the th element of `expert_weights[this_expert_mask, expert_ix]`.\n\nHowever, there is one tensor-compatibility issue. What we have is shaped\n`(num_selected_context_vectors_for_this_expert,)` -- that is, it's essentially just\na list of numbers.\n\nNow, we want to multiply our results by these weights, which means that we want to\nbroadcast them across `this_expert_results`, which\nis shaped `(num_selected_context_vectors_for_this_expert, d_emb)`.\n\nNaively you might think that if you try to multiply a `(num_selected_context_vectors_for_this_expert, d_emb)`\ntensor by a `(num_selected_context_vectors_for_this_expert,)` one, PyTorch would match\nup the two dimensions of the same size and it would just work. But it would actually\ntry to match up across dimensions from right to left, and would complain that\n`num_selected_context_vectors_for_this_expert` does not equal `d_emb`. So instead,\nit's best to feed it an explicit `(num_selected_context_vectors_for_this_expert, 1)`\ntensor so that it knows which dimensions we're trying to match up, and which ones\nit should broadcast over:\n\n```\n this_expert_weights = expert_weights[this_expert_mask, expert_ix].unsqueeze(1)\n```\n\n...and that will give us `this_expert_weights` of shape `(num_selected_context_vectors_for_this_expert, 1)`\nThat would look identical to [diagram 19](#diagram-19) in the way I've been diagramming these things, but it's technically\ndifferent and necessary.\n\nSo now, with the\n`unsqueeze` having fixed the dimensionality so that we can do a broadcast, we can do this:\n\n```\nthis_expert_results * this_expert_weights\n```\n\n...and with that we have our weighted results for this expert -- all of the context vectors that should have been run through it have been, and they've been multiplied by the weights that the router gave for them, so we have our backprop channel for training the router.\n\nThe next step is that we need to somehow put these results into `all_outputs`. Specifically,\nbecause it's initially all zeros, and we want it to hold the results of adding together the\nresults from each active expert for each context vector, we need to add them to whatever is\nalready there at this point in the iteration through the experts.\n\nEven though -- as noted earlier -- we have, by this point in the calculations, lost\nthe connection between the tensors we've been working with and the original \"front-facing\"\ngrid of our original tensors like `xs`, where we could link something directly to the\nincoming context vector that it related to, it's actually surprisingly easy to patch that\nup :-)\n\nRemember that outside our loop through the experts, we created `all_outputs` like this:\n\n```\n all_outputs = torch.zeros_like(xs)\n```\n\nSo, just like `xs` in [diagram 6](#diagram-6), it looked like this:\n\n**Diagram 14** (repeated): `all_outputs` as two 2-D views\n\nNow, previously we had code to extract the context vectors that we wanted to go through\nour expert from `xs`, and it looked like this:\n\n```\nxs[this_expert_mask]\n```\n\nThat gave us what amounted to a list of context vectors, like this:\n\n**Diagram 17** (repeated): the \"selected\" context vectors\n\nNow, the interesting thing about a masked lookup into a tensor like `something[some_mask]`\nis that not only can you do it to extract a subset of the values from the tensor like\nwe did with `xs` -- you can also use it to assign to the masked subset of the elements in\nthe `something` tensor.\n\nTo make that concrete: when we did\n\n```\nxs[this_expert_mask]\n```\n\nWe got a tensor sized `(num_selected_context_vectors_for_this_expert, d_emb)`.\n\nThat means that if we were to do:\n\n```\nall_outputs[this_expert_mask]\n```\n\n...then given that the mask is the same, and `all_outputs` has the same shape as\n`xs`, then we'd also get a `(num_selected_context_vectors_for_this_expert, d_emb)`\nresult.\n\nBut the thing with assignment means that if we had some tensor shaped `(num_selected_context_vectors_for_this_expert, d_emb)`\n-- let's call it `foo` -- then we could do this:\n\n```\nall_outputs[this_expert_mask] = foo\n```\n\nWith that, instead of selecting the parts of `all_outputs` and using them in the future,\nwe're overwriting them with whatever is in `foo`.\n\nFurthermore, we can use the same trick with the augmented assignment operators like\n`+=`:\n\n```\nall_outputs[this_expert_mask] += foo\n```\n\n...means \"select the elements from `all_outputs` that have `True` in `this_expert_mask`,\nand then increment them by the elements of `foo`\".\n\nSo finally, we get our code:\n\n```\n all_outputs[this_expert_mask] += this_expert_results * this_expert_weights\n```\n\nThat multiplies the results from the expert by the corresponding weights, and then adds\nthem in to the running totals we're keeping in `all_outputs`. Because `all_outputs` is the\nsame shape as `xs`, and thus `all_outputs[this_expert_mask]` is the same shape as\n`xs[this_expert_mask]`, and we've kept that same shape as we went through the expert itself\nand multiplied by the weights, it will work.\n\nAnd with that, we've done all of the calculations that we need for this specific expert, so we can go back round the loop for the next one. When we've finished with all of the experts, we can return the result:\n\n```\n return all_outputs\n```\n\nPhew! That was quite a lot of explanation, but I think that what is going on in the code needs it. I originally wrote it in a kind of flow state of inspiration, and was very pleased with it, but it feels like now that I've explained it, I actually understand what my subconscious must have known while I was writing it. And I hope it's reasonably clear for anyone reading this.\n\nBut there was one extra thing I wanted to keep track of before I started running this: the balance between the different experts.\n\nWe'll come back to this in some detail later, but a problem with MoEs is that they can wind up depending heavily on specific experts, and ignoring the others -- the auxiliary loss I mentioned way back in the intro to this post is required to avoid that.\n\nI didn't want to implement that yet, but I wanted to log enough data to see if it really would be necessary with my setup.\n\nA good way to keep track of how much each expert is being used is to record the logits -- the original results from the router, before the top- and the softmax -- and the actual post-top-, post softmax expert weights that we actually used.\n\nThe changes to do this are dotted around a bit, so I've linked to the appropriate lines on GitHub and if you have the screen real-estate to do so, I'd recommend that you use that to follow along. However, I've tried to put enough code inline in this post below that it should be comprehensible without that.\n\nI decided that I'd return the logits and the expert weights from the `MixtureOfExperts` module's forward pass as\nan extra output [here](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/gpt.py#L152):\n\n``` python\nclass MixtureOfExperts(nn.Module):\n ...\n def forward(self, xs):\n ...\n return all_outputs, (routing_logits, expert_weights)\n```\n\nNow, back in the `TransformersBlock` we were setting\n`self.ff` to either a `FeedForward` object or a `MixtureOfExperts` depending on\nwhether MoE was enabled for this model [here](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/gpt.py#L169):\n\n``` python\nclass TransformersBlock(nn.Module):\n ...\n def __init__(self, cfg):\n ...\n self.is_moe = \"moe\" in cfg\n if self.is_moe:\n self.ff = MixtureOfExperts(cfg)\n else:\n self.ff = FeedForward(cfg)\n```\n\nThat meant that in our forward pass where we previously just did this (I won't link to this because it's an old version and having links to different versions would be confusing):\n\n``` python\nclass TransformersBlock(nn.Module):\n ...\n def forward(self, x):\n ...\n x = self.ff(x)\n```\n\n...then if the model was an MoE, we'd be getting a tuple -- the actual outputs from\nthe module and then that extra routing information we were returning. But if it was\na `FeedForward`, we'd just get the outputs.\n\nI decided that the simplest fix was to change the `FeedForward` module so that it returned\nvalues that were compatible with `MixtureOfExperts`. The old `forward`, which was\nthis:\n\n``` python\n def forward(self, x):\n return self.layers(x)\n```\n\n...became [this](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/gpt.py#L107):\n\n``` python\n def forward(self, x):\n return self.layers(x), None\n```\n\nThat meant that we could do [this](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/gpt.py#L191)\nin place of the `TransformersBlock.forward` code above:\n\n```\n x, this_block_moe_routing_info = self.ff(x)\n```\n\nIf we had an MoE in `self.ff`, then we'd get some real MoE routing info in\n`this_block_moe_routing_info`, whereas if we were running a normal dense model then\nwe'd get `None`.\n\nThis, by the way, explains the odd bit of code in the `forward` for `MixtureOfExperts`\nthat you might remember from earlier -- [this bit](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/gpt.py#L148):\n\n```\n this_expert_results, _ = expert(xs[this_expert_mask])\n```\n\n`expert` there is one of the `FeedForward` modules, so what we're doing there with the\nunderscore is just\nignoring the `None` that we know we're returning from it.\n\nThe next step was to work out how to get the extra `(routing_logits, expert_weights)` information\nthat we had in `this_block_moe_routing_info` out of the `TransformersBlock`, if\n`self.ff` was a `MixtureOfExperts`.\n\nNow, `TransformersBlock` is used in an `nn.Sequential` in the `GPTModel` class [here](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/gpt.py#L210):\n\n``` python\n def __init__(...):\n ...\n self.trf_blocks = nn.Sequential(\n *[TransformersBlock(cfg) for _ in range(cfg[\"n_layers\"])]\n )\n ...\n def forward(self, inputs):\n ...\n x = self.trf_blocks(x)\n ...\n```\n\nThat means that the outputs from the `forward` method of the first one are fed directly in as the inputs\nto the second one, and so on. Additionally, `nn.Sequential` assumes that the forward\nmethod just takes a single input.\n\nWhat I really wanted at the end of this was a list of `(routing_logits, expert_weights)` pairs,\none for each Transformers layer. So I decided to pass an accumulating list into\n`TransformersBlock.forward`\n(allowing it to be `None`), like [this](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/gpt.py#L179):\n\n``` python\nclass TransformersBlock(nn.Module):\n ...\n\n def forward(self, inputs):\n x, moe_routing_info = inputs\n if self.is_moe and moe_routing_info is None:\n moe_routing_info = []\n```\n\n...and then append whatever routing info came from the (potentially `MixtureOfExperts`)\nFFN to that [here](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/gpt.py#L192):\n\n```\n if self.is_moe:\n moe_routing_info.append(this_block_moe_routing_info)\n```\n\n...and then return it from `TransformersBlock.forward` for the next layer [here](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/gpt.py#L197):\n\n```\n return x, moe_routing_info\n```\n\nThen, in `GPTModel` I wanted to return it to whatever called the model for inference.\nI didn't want to change the implicit API that I was using -- pass in inputs, get\nnext-token logits -- so I decided to just attach it to the `logits` like [this](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/gpt.py#L230):\n\n``` python\nclass GPTModel(nn.Module):\n ...\n\n def forward(self, in_idx):\n ...\n\n x = self.drop_emb(x)\n x, moe_routing_info = self.trf_blocks((x, None))\n x = self.final_norm(x)\n\n logits = self.out_head(x)\n logits.moe_routing_info = moe_routing_info\n\n return logits\n```\n\nI'm not sure, in retrospect, that \"smuggling\" the routing info out like that was\nthe right choice. Perhaps a Hugging Face-like model where the LLM returns some kind\nof \"output\" class that contains `logits` and other stuff like this as explicit fields would be better. But this setup seemed\nto work for my use case, so I've left it as-is for now, with a mental note to revisit\nlater.\n\nNext, I modified my training script to store the `moe_routing_info` that we got\nwhen we did our forward pass in a list, and then to store that in the metadata associated\nwith my checkpoints for later analysis.\n\nThe training script I'm using is something I developed after working through Raschka's\nbook, and it's kind of complicated -- although the core is essentially the same as the\none we use to train our model in chapter 5, I've built it out to allow training\n[across multiple GPUs](/2026/01/llm-from-scratch-29-ddp-training-a-base-model-in-the-cloud), with all kinds of\n[tweaks](/2026/04/llm-from-scratch-32m-interventions-conclusion) that I've learned about\nsince. So I won't dig into that code in any depth in this post; if you've been following\nalong with my various training posts in the past, though, and want to see how this\nall fits in, you can see the code that accumulates the routing information\n[here](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/ddp_train.py#L444)\n(note that it uses `detach` so that we don't accumulate compute graph information\nover time),\nthe code that passes those details into the checkpointing function [here](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/ddp_train.py#L541),\nand the updated checkpointing function [here](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/checkpointing.py#L37).\n\nSo: with those changes, I had a plausible-looking MoE setup. I was confident that it would be able to train, but suspected that it would not be able to balance load between the experts well and would collapse to using a subset of them.\n\nIt was time to give it a go.\n\nThe first question was, what total number of experts, and how many active ones, should I have? I was pretty limited by what I could fit into my RTX 3090's VRAM, and what an appropriate training speed might be, but after a bit of fiddling around I came to the conclusion that six experts with two active would fit into memory and train reasonably quickly, and that didn't sound like a crazy balance. Mixtral is 8 experts, two active, for example.\n\nI wouldn't be able to\nfit in the batch size of 6 that I had used in the past when training 163M-parameter dense models;\nthe largest batch size I could fit in was 3.\n(For those who've been following my previous LLM training,\nbecause I'm using\n[gradient accumulation](/2026/04/llm-from-scratch-32k-interventions-training-our-best-model-locally-gradient-accumulation) over\n16 steps with the microbatch size of 6 -- for a global batch size of 96 -- I could get the same effect,\nand fit the MoE into VRAM, by going down to a microbatch size of 3, with 32 gradient\naccumulation steps.)\n\nThe next question was how many tokens to train for. As an experiment I kicked it off with the same 3.2 billion tokens that I would normally use to train my 163M-parameter models. I knew that this larger model would need to be trained on more tokens than that, but I just wanted a reasonably serious run to see how the model behaved.\n\nThe training script predicted that it would take just over two days to complete this experimental run, which was perfect. I had reached this point during the late afternoon on a Friday, and had stuff to do over the weekend that would keep me away from my computer, so that run length would mean that it would be ready for me on Monday.\n\nOn Monday, I came back to see that it had completed properly:\n\n```\nTraining complete in 217,540.703 seconds\nTokens seen: 3,260,252,160\nThroughput: 14,987 tokens/second\n```\n\nOver the training run, the loss declined nicely:\n\nOut of interest, I ran an evaluation script to see how it performed on the held-back test set that I normally use to compare my models:\n\n``` bash\ngiles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run test_loss.py datasets/ runs/1xrtx3090-moe-no-balancing/model.json runs/1xrtx3090-moe-no-balancing/checkpoints/latest/model.safetensors\nFetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 2215.98it/s]\n100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3200/3200 [05:28<00:00, 9.73it/s]\nLoss against our test dataset: 3.501952\n```\n\nThat ranked it better than the 163M-parameter models I'd trained using PyTorch, but\nworse than the original GPT-2 small, and also worse than the models I've\n[trained with JAX](/2026/07/llm-from-scratch-34b-building-and-training-gpt-2-small-in-jax)\n(which I believe got lucky with their initial untrained random weights). That was\na pretty solid result, but as this was an exploratory training run, there's no point in\nputting much weight on it.\n\nFor a start, this model was clearly undertrained. The\nnumber of tokens, 3.2B, was [Chinchilla-optimal](/2026/08/chinchilla-check) for my 163M-parameter\nmodels, but this one had more than twice the total parameters at 446M, and had\n220M active for each token. Larger models need more tokens to be well-trained.\n\nThe more important result was the load-balancing between experts.\nI put together a [notebook](https://github.com/gpjt/ddp-base-model-from-scratch/blob/fb7d62a682e58d864a51d4d59ec76f95f105de7d/plot_router_saturation.ipynb) to\ningest the metrics that I was writing to the checkpoints, and to create two charts for each\nlayer:\n\nIn the charts, for each global step, I plotted a line for each expert; if it had an average probability less than uniform / 3 -- that is, it was getting less than a third of the tokens that a perfectly flat distribution across all experts would give it -- then the line was white, meaning that it was being \"starved\". If it was getting a number of tokens between uniform / 3 and uniform * 2 -- then the line was green, because it was getting what felt like a reasonable number of tokens. And finally, if it was higher than uniform * 2, it was plotted in red: it was being overfed.\n\nThe charts showed exactly the problem I was expecting to see. Let's look at layer 0:\n\nYou can see that it started off pretty well-balanced, but rapidly came to prioritise expert 5; the remainder of the probability distribution looks like it must have been scattered over the other experts, with a slight preference for expert 3.\n\nThe other layers came back with similar issues: certain experts were strongly preferred while others were starved:\n\nSo the problem was real. The model was relying heavily on some experts and ignoring others. We needed to do extra stuff to balance the load across the experts.\n\nThe problem with load-balancing across experts has been clear for quite some time, and\nit was covered in \"Outrageously Large Neural Networks\" back in 2017. The solution they used is simple, and clever:\ndefine an \"auxiliary\" loss, which goes up as the experts become more unbalanced, and\ndown as they become more balanced. You can then add that (scaled by some amount) on to the normal\n[cross entropy loss](/2025/10/llm-from-scratch-20-starting-training-cross-entropy-loss) that you get\nby comparing the outputs from your model to the training targets, to get a combined\nloss number -- and then you back-propagate using that combined loss.\n\nBy reducing the combined loss, then you optimise both for correctness -- is the model doing its job as an LLM -- and for balance across the experts. Of course, you need to be careful about the scaling factor that you use to adjust the auxiliary loss before adding it. If it's too small, then the load-balancing will be de-prioritised and you can wind up with imbalanced experts anyway. But if it's too large, the training process will prioritise balance over quality of the output, and you'll wind up with a crappy model that balances load across the experts beautifully. We'll come back to that shortly.\n\nThe \"Outrageously Large Neural Networks\" paper's particular calculations for the auxiliary loss are a little complicated -- they generate two separate numbers and combine them -- and it was simplified in the GShard paper, and then again in Switch Transformers. I found the last of those reasonably easy to understand, and decided to adapt it.\n\nLet's start off with their formalisation -- but keep in mind that they had only one active expert per context vector, so we'll need to make some changes.\n\nFirstly, they define a per-expert number, -- the subscript means that it's for expert .\n\nThat's a little intimidating-looking but is much simpler than it looks. They define it as \"the fraction of tokens dispatched to expert \", and in that light we can interpret it.\n\nThe that we're iterating over in the is basically:\n\n```\nfor x in batch_context_vectors:\n```\n\nSo, the stuff inside the is done once for each of our input context vectors across all sequences in a batch.\n\nThe (which should be rendering with a kind of doubled-1 -- Chrome, as of this writing, doesn't handle that, though Firefox does) is meant to mean \"1 if the condition in the braces is true, 0 if it's false\".\n\nIn , the function is the \"raw\" probabilities from our router. As I mentioned earlier, they were doing softmax and then zeroing out the non-top- values, so they had those raw numbers available. We'll come back to that in a moment.\n\nBut for now, the condition inside the just means \"true if this expert was the top one for this particular incoming context vector, false otherwise\".\n\nPutting that all together, then we're just counting the number of tokens in our batch for which expert is the top pick of the routing network. The at the start then divides that by the number of tokens in the batch (which is what they mean by ), and for a network with one active expert like the Switch Transformers ones, we've got -- as they say -- \"the fraction of tokens [in this batch] dispatched to expert \".\n\nSo that works nicely in the one-active Switch Transformers world.\n\nBut we can fairly simply extend it to handle multiple active experts, while keeping similar semantics. Let's say that we calculate how many of the context vectors in a batch were sent to expert ; we can divide that by the number of tokens to get something equivalent. It still means \"the fraction of tokens dispatched to expert \".\n\n(Ab)using their notation, if we say that our number of active experts is , then it might look something like this:\n\nLet's run with that for now.\n\nAs well as , they define another per-expert number, :\n\nThey describe this as \"the fraction of the router probability allocated for expert \".\n\nAgain, we're iterating over every context vector in every sequence in the batch -- but here we're just adding together all of the raw probabilities for the expert for each one, then dividing by the number of elements in the batch. In other words, we're working out the expert's average probability across the entire batch. Much easier! And in our multiple-active-expert world, their equation works -- it means exactly the same thing.\n\nNow, both of these calculations, as they're expressed in the maths, depend on something that we're not currently calculating. Remember, Switch Transformers worked out the weights for each expert by doing a softmax across all of the raw logits that came out of the router's linear layer, then zeroing out the ones that were not in the top- -- that is, for their setup, all but the highest (argmax) one. So they had those \"raw\" softmaxed probabilities knocking around.\n\nBut because we are replacing the non-top- logits with and then doing softmax, our router weights aren't the same -- they're just the relative probabilities of our selected experts.\n\nFor that doesn't matter; we can use our weights and easily identify the experts that were routed to for a given element in the batch, because they're the only ones that are not zero.\n\nBut for we do need the pre-top-'s logits from the router. And, not entirely\ncoincidentally, we already have the code to make that available! In order to do those charts above\n-- the ones showing where experts were being starved and when they were being over-fed\nin the non-load-balanced training run --\nwe passed the raw, non-top-'ed, non softmaxed router logits, and the expert weights\nafter the top- and the softmax, out of the `MixtureOfExperts` module's `forward`:\n\n``` python\nclass MixtureOfExperts(nn.Module):\n ...\n def forward(self, xs):\n ...\n return all_outputs, (routing_logits, expert_weights)\n```\n\n...and then added code to feed that through to the training loop because it was needed for the metrics that we used to generate those charts. The numbers that we kept for the \"raw\" side of things were the logits rather than the actual probabilities, but we can fix that with a simple softmax.\n\nSo that means that in our training code, we already had the numbers to work out and for our experts. They needed to be combined to make up a single scalar auxiliary loss for the model when run over a batch, and Switch Transformers does that like this:\n\nIf you imagine as a vector containing all of the per-expert s, and as a similar one containing the s, then that is a simple dot product, .\n\nThey then scale that up by the number of experts , multiply by a scaling factor -- the one I mentioned earlier to balance load-balancing auxiliary loss against \"real\" cross entropy loss -- which they call . (We'll come back to later.)\n\nSo, we have a set of calculations to work out an auxiliary loss; there are a few extra things I'd like to highlight before we dive in to the code.\n\nLet's imagine what a \"perfect\" router would look like from the perspective of these calculations, firstly for the Switch-style one-active-expert model, and then for our own more general case.\n\nStarting with ; it's the number of tokens in the batch for which expert is the chosen one, divided by the number of tokens in the batch. We want all of our experts to get the same amount of \"traffic\", so for experts, one active per token, clearly each one will be active of the time. So that should be the value of if everything is balanced.\n\nNow let's think about . Again, we want each expert to receive of the tokens -- so its average probability should also be .\n\nSo, that means that for perfect balance, all of our s and all of our s should be . That means that when we do the in that loss calculation to work out the dot product, then for a perfectly balanced router, each one will contribute:\n\nThere will be of them, so that will come to\n\n...and then we're multiplying by to get the loss, so the result (disregarding the scaling factor ) will be .\n\nSo, when balance across the experts is perfect, the Switch Transformers auxiliary loss has a value of .\n\nHowever, they have only one active expert, and our equation is slightly different to allow for the fact that we have of them.\n\nI won't go through the boring derivation again, but if we replay the maths above with active experts, we get an auxiliary loss for the ideal, perfectly balanced router of .\n\nThat's not a problem in and of itself -- after all, this is just a number we're trying to minimise, and it's not super-important what we're trying to minimise it to. But it does matter when we're talking about , because if the auxiliary loss is larger, we'll need to scale it down more to stop it from \"drowning out\" the signal from the actual training loss. The Switch Transformers paper explains what values they used for the scaling, but ours will be different because of the different number of active experts.\n\nThe auxiliary loss calculation above only covers what happens in one layer, so we need to work out how to combine the contributions from all of the layers. In the paper, the only mention they make of multiple layers in this context is:\n\nFor each Switch layer, this auxiliary loss is added to the total model loss during training\n\nI took that to mean that we just sum up the scaled auxiliary loss across all layers and then add that sum to the normal cross entropy loss. I think that's the most natural interpretation of what they are saying.\n\nSo -- given the value of for a perfectly-balanced router that I worked out above -- for the model I was planning to train, with 2 active experts, 12 layers, the auxiliary loss before scaling by would be 24.\n\nThis, by the way, is where my implementation differs from Mixtral's -- or, at\nleast, [the Hugging Face source code](https://github.com/huggingface/transformers/blob/5f8ab9bb53ec9e0c9329153d18bd825ff1db80f9/src/transformers/models/mixtral/modeling_mixtral.py)\nfor it as of this writing.\nIn that, they do something that feels a bit odd. They treat (for example) expert 1\non layer 1 as being the same as expert 1 on layer 2, and so on throughout the layers,\nthen do the calculations just once. That feels\na bit dodgy. After all, you can imagine that expert 1 on layer 1 might be being\nstarved but its equivalent on layer 2 might be getting overfed, and the two would\nbalance out.\n\nI don't know if that's an error in that implementation, or if there's something I'm missing. Conceivably I might test it some day by training another model using their loss function, but I suspect I won't get around to that.\n\nThere's also something that made me hesitate a bit in the Switch Transformers paper, and which I think I still need to ponder a bit. That calculation for :\n\n...did not look differentiable to me. Things like (and our own equivalent's ) are generally not.\n\nIndeed, they confirmed that shortly after defining it, but in a way that gave me pause:\n\nThe objective can also be differentiated as the -vector is differentiable, but the -vector is not.\n\nThe \"objective\" they're referring to is the auxiliary loss, and it makes me a bit uncomfortable that something that is defined in terms of A and B is differentiable if A is, but B isn't.\n\nMy hand-wavy way of thinking about it right now is that the undifferentiable bit can be treated as a constant, so as long as part of the calculation is differentiable, the whole thing can be treated as such. But I'm not 100% happy with that, and need to think further.\n\nBut now, I think, we've covered the maths for the auxiliary loss, so it's time to dive into the code!\n\nMy old training loop had the following\n[code](https://github.com/gpjt/ddp-base-model-from-scratch/blob/2a9a25fb0488dcbadb7d17738ead3bbde4ebb539/ddp_train.py#L335) to do the forward then the backward pass:\n\n```\n if use_amp:\n with torch.amp.autocast(device_type=device.type, dtype=torch.float16):\n logits = model(inputs)\n train_loss = calculate_loss(logits, targets)\n else:\n logits = model(inputs)\n train_loss = calculate_loss(logits, targets)\n\n is_last = accumulation_step == gradient_accumulation_steps - 1\n with model.no_sync() if not is_last else nullcontext():\n if scaler is not None:\n scaler.scale(train_loss / gradient_accumulation_steps).backward()\n else:\n (train_loss / gradient_accumulation_steps).backward()\n```\n\nLet's strip out all of the extra enhancements that I have accumulated there on top of the simple\ntraining code from the book; there's\n[AMP](/2026/04/llm-from-scratch-32h-interventions-full-fat-float32),\n[DDP](/2026/01/llm-from-scratch-29-ddp-training-a-base-model-in-the-cloud) and\n[gradient accumulation](/2026/04/llm-from-scratch-32k-interventions-training-our-best-model-locally-gradient-accumulation)\nand if we remove that it would simply look like this:\n\n```\n logits = model(inputs)\n train_loss = calculate_loss(logits, targets)\n\n train_loss.backward()\n```\n\nHopefully that's familiar!\n\nWhat I wanted to do was add in the auxiliary loss (if we were training an MoE), scaled by that scaling factor. What I came up with (and again, here I've stripped out all of the stuff required by those enhancements):\n\n```\n logits = model(inputs)\n train_loss = calculate_loss(logits, targets)\n if hasattr(logits, \"moe_routing_info\") and logits.moe_routing_info is not None:\n moe_router_loss = calculate_moe_router_loss(logits.moe_routing_info)\n moe_router_losses.append(moe_router_loss.item())\n else:\n moe_router_loss = 0\n\n total_loss = train_loss + moe_router_loss * moe_router_loss_scale\n total_loss.backward()\n```\n\nNote that I wanted to keep track of the router losses in that `moe_router_losses` list as well in order to monitor\nthem as the training run progressed, just like I normally monitor training loss (as you\ncan see in the loss chart above).\n\nThat's all pretty nice and simple -- if we are getting MoE routing info back from\nthe model, then we call this new `calculate_moe_router_loss` to work out the loss\nfrom the maths in the last section, and then add it on, scaled by\n`moe_router_loss_scale`, which is what I decided to call the\nsomewhat-opaquely-named from the Switch Transformers paper.\n\nAdding all of the\nAMP, DDP and gradient accumulation gubbins back in, the final code looked like\n[this](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/ddp_train.py#L425):\n\n```\n with torch.amp.autocast(device_type=device.type, dtype=torch.float16) if use_amp else nullcontext():\n logits = model(inputs)\n train_loss = calculate_loss(logits, targets)\n if hasattr(logits, \"moe_routing_info\") and logits.moe_routing_info is not None:\n moe_router_loss = calculate_moe_router_loss(logits.moe_routing_info)\n moe_router_losses.append(moe_router_loss.item())\n else:\n moe_router_loss = 0\n\n is_last = accumulation_step == gradient_accumulation_steps - 1\n with model.no_sync() if not is_last else nullcontext():\n total_loss = train_loss + moe_router_loss * moe_router_loss_scale\n if scaler is not None:\n scaler.scale(total_loss / gradient_accumulation_steps).backward()\n else:\n (total_loss / gradient_accumulation_steps).backward()\n```\n\nSo that was simple enough (for LLM-training values of simple).\n\nThe code to route the `moe_router_loss_scale` from the\ntraining configuration file is not really worth going through, and nor is the\ncode to save average, minimum and maximum values from `moe_router_losses` into the\ncheckpoint metadata, or to chart those (though I'll show the charts later).\n\nThe interesting bit is, of course, that `calculate_moe_router_loss` function.\n\nIt's [here](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/ddp_train.py#L345)\nand looks like this:\n\n``` python\ndef calculate_moe_router_loss(moe_routing_info):\n total_routing_loss = 0\n for routing_logits, expert_weights in moe_routing_info:\n batch_size, seq_len, num_experts = expert_weights.shape\n\n flattened_expert_weights = expert_weights.view((batch_size * seq_len, num_experts))\n expert_active_counts = (flattened_expert_weights > 0).sum(dim=0)\n expert_frequencies = expert_active_counts / (batch_size * seq_len)\n\n raw_routing_weights = torch.softmax(routing_logits, dim=-1)\n flattened_raw_routing_weights = raw_routing_weights.view((batch_size * seq_len, num_experts))\n expert_prob_allocation = flattened_raw_routing_weights.sum(dim=0) / (batch_size * seq_len)\n\n layer_routing_loss = num_experts * torch.dot(expert_frequencies, expert_prob_allocation)\n\n total_routing_loss += layer_routing_loss\n\n return total_routing_loss\n```\n\nLet's look at it from the outside in.\n\nWe start with a list called `moe_routing_info`. Remember, this has been passed\nback from the MoE model for a single forward pass of a batch. It contains one item\nfor each Transformers layer in the model, and those items are pairs of\n`(routing_logits, expert_weights)`.\n\nWe start off with a total routing loss of zero, and then for each layer, we do some stuff to work out its routing loss, and then at the end of the loop, we add it on to our running total. Finally, we return the total.\n\nObviously, the fun stuff is inside the loop :-) It's time for some of those tensor diagrams again.\n\nLet's remind ourselves of the shape of `expert_weights`:\n\nEach cell on the front face relates to one context vector in one sequence in our batch, and the \"core\" going into the cuboid from there (horizontally along the side view) is the weights we actually used when routing the context vector in question to the experts -- zero for unselected experts, some value between zero and one for the selected ones.\n\nNow let's go back to the code for a moment.\nWe start off by using the shape of `expert_weights` to work out what our different\ndimensions are:\n\n```\n batch_size, seq_len, num_experts = expert_weights.shape\n```\n\nThen we do our first block of calculations, trying to work out from the maths, \"the fraction of tokens dispatched to expert \" We want to do this efficiently with a vector calculation, working out all of the s (remember, there's one for each expert) in parallel.\n\nOur first step is to flatten out the `(batch_size, seq_len)` grid into a single\ndimension -- essentially stacking all of the front view's columns on top of each other to make just one long\ncolumn, like this:\n\n**Diagram 20**: `expert_weights` flattened, as two 2-D views\n\n...or, more simply, as it's now just a 2-D Tensor, like this:\n\n**Diagram 21**: `expert_weights` flattened, as one 2-D view\n\nWe can use PyTorch's [`view`](https://docs.pytorch.org/docs/2.14/generated/torch.Tensor.view.html)\nmethod to do that without having to copy any data around in memory -- as the name suggests,\nit just returns a different view on the same data:\n\n```\n flattened_expert_weights = expert_weights.view((batch_size * seq_len, num_experts))\n```\n\nNow, remember that these are the expert weights. Each one of those cells contains a number -- zero if the expert was not selected for the context vector that corresponded to it in the original layout, or some non-zero number if it was. So if we do this:\n\n```\nflattened_expert_weights > 0\n```\n\n...then we'll get a tensor of the same shape, where we have `True` if the weight was\nmore than zero (that is, the expert was active for that context vector), `False`\notherwise.\n\nAnd that means that if we sum down those columns in [diagram 21](#diagram-21), we'll get a new grid of one row,\nand `num_experts` columns, which represents the total number of times each expert was active\nin the given batch:\n\n**Diagram 22**: `expert_active_counts` as one 2-D view\n\nSo, in code, we can just do this:\n\n```\n expert_active_counts = (flattened_expert_weights > 0).sum(dim=0)\n```\n\nIf we divide that by the number of context vectors in the batch, we've got a new\ntensor, a row with `num_experts` columns -- the same shape as [diagram 22](#diagram-22) -- containing exactly what we want:\n\n```\n expert_frequencies = expert_active_counts / (batch_size * seq_len)\n```\n\nThat is, `expert_frequencies` is a vector in terms of the maths, containing\nall of the s that we want for our auxiliary loss calculation -- that is, all\nof the result across all experts for this:\n\nThe calculations for the s are very similar. We start off with `routing_logits` like this:\n\n**Diagram 7** (repeated): `routing_logits` as two 2-D views\n\nSo, each cell on the front face relates to one context vector, and the \"core\" going into the cuboid from there is the set of raw routing logits for that context vector, one number per expert.\n\nFirstly we need to convert\nthe `routing_logits` into probabilities by running them through softmax, along the\nlast dimension -- the `num_experts` one that is the horizontal axis on the side view:\n\n```\n raw_routing_weights = torch.softmax(routing_logits, dim=-1)\n```\n\nThen we do the same trick with `view` to convert the result to a single column\nof lists of length `num_experts` (metaphorically) -- the same shape as in [diagram 21](#diagram-21):\n\n```\n flattened_raw_routing_weights = raw_routing_weights.view((batch_size * seq_len, num_experts))\n```\n\nNow if we add them up across the rows like this:\n\n```\nflattened_raw_routing_weights.sum(dim=0)\n```\n\nThen we get a single row, `num_experts` column result that has, for each expert, the sum of\nall of the probabilities it had across all of the context vectors in the batch,\njust like the one we had in [diagram 22](#diagram-22).\n\nWe can divide that by the number of context vectors in the batch:\n\n```\n expert_prob_allocation = flattened_raw_routing_weights.sum(dim=0) / (batch_size * seq_len)\n```\n\n...and that's our vector containing all of the s, where each is one of these:\n\nFinally, we can multiply all of the s and their corresponding s by each other, and sum the results, by using a dot product, and then multiply the result by the number of experts:\n\n```\n layer_routing_loss = num_experts * torch.dot(expert_frequencies, expert_prob_allocation)\n```\n\n...and that's this bit done (apart from the ):\n\nAnd that's our auxiliary code wrapped up! We've been through `calculate_moe_router_loss`,\nand we've already seen the code that scaled it by aka `moe_router_loss_scale`,\nso we have a training script with MoE auxiliary loss using the maths in the last section!\n\nAgain, I hope that the diagrams helped with that workthrough. The code is the kind of thing where it's easy to scan through and get a vague understanding, but I think that visualising what's going on step-by-step is important if you want it to really stick.\n\nWith the code in place, the next thing to do was to explore what the right value might be for .\n\nIn the \"Switch Transformers\" paper, they say:\n\nFinally, a hyper-parameter is a multiplicative coefficient for these auxiliary losses; throughout this work we use an which was sufficiently large to ensure load balancing while small enough to not to overwhelm the primary cross-entropy objective. We swept hyper-parameter ranges of from to in powers of and found balanced load quickly without interfering with training loss.\n\nBut, as I noted earlier, that worked for them with their single active experts, but because I had multiple, my auxiliary loss would be larger. Now, given that their \"ideal\" per-layer loss was 1, and mine was 2 for my planned training run, it sounded like using half of their recommended value, , would be appropriate.\n\nHowever, I was also a little concerned about the number of layers interfering with things.\n\nThe auxiliary loss, as we saw above, was a single value per layer, all of which were added together. With my \"ideal\" per-layer loss of 2, and 12 layers, that meant that the ideal across all layers was 24.\n\nNow, they mentioned a specific value for , but didn't mention the number of layers they had, or if they swept for different possibilities across different numbers of layers. That seemed strange!\n\nI decided to work on the hypothesis that they had found that as the number of layers increased, you needed the total contribution of the auxiliary loss to scale up in proportion. That was only a guess based on trying to fit together the info in the paper, though, and could well be wrong.\n\nI was in enough doubt, though, that I felt it would be wise to do my own, minimal sweep over a few values for . For each, I'd do a one-hour training run. At the end I'd check the training loss -- that is, the pure cross entropy loss saying how well the model was doing at its real purpose of modeling language -- and the auxiliary loss, showing how well-balanced its usage of experts was.\n\nI got these results:\n\n| | Training loss | Auxiliary loss | \n|---|---|---|\n| 0 | 6.152867 | 53.38285 | \n| 0.001 | 6.166725 | 30.00601 | \n| 0.005 | 6.106290 | 25.21015 | \n| 0.01 | 6.160989 | 24.78106 | \n\nI also generated router usage maps like the ones I gave way back in this post for the two-day training run with no auxiliary loss. I won't put them all in there, but:\n\nSo, on the basis of those results (and, of course, the fact that it fit well with the Switch Transformers paper's recommendation), I decided to use .\n\nIt was time to train this thing!\n\nI decided not to think too hard about the right number of tokens to train it on. The Chinchilla number, 20 times as many tokens as parameters, is a heuristic that works for dense models, but is not meant for MoEs. You'd intuitively think that the \"ideal\" number of tokens to train a model with 446,410,752 parameters, 219,697,152 active per token would be somewhere between the Chinchilla-optimal numbers for those two parameter counts.\n\nBut overtraining isn't necessarily a bad thing, so long as it's on non-duplicated\ndata (or [less than four epochs over the same data](/2026/07/why-do-openai-gpt2-weights-beat-mine-3-overtraining#fn-2)),\nand I had a 10B token dataset all set up from my previous experiment. So training\nit on what would be the Chinchilla-optimal number of tokens if it was just a\n446,410,752-parameter dense model didn't sound like a bad idea, so long as I could\ndo it in a reasonable amount of time.\n\nThat meant 8,928,215,040 tokens, which I had in my normal training dataset without needing multiple epochs.\n\nA quick check -- running the training script with that number of tokens configured -- told me that it would take 90,823 global steps to complete, over about eight days.\n\nIt looked like each checkpoint would take up 5.3GiB, and I had about 348 GiB free on my disk, so I calculated that I could checkpoint every 1,500 global steps -- that would work out as roughly once every three hours. Not great, but losing a maximum of three hours work in the case of a power outage (or cat jumping onto the PC's power button) is not the end of the world.\n\nSo I set things up with [this model configuration](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/runs/1xrtx3090-moe-first-proper-run/model.json)\nand [this training configuration](https://github.com/gpjt/ddp-base-model-from-scratch/blob/9891589eade90f51b3b5c0747b426edc44c2e6fa/runs/1xrtx3090-moe-first-proper-run/train.json),\nand on my dedicated training box, [`poppy`](/poppy-the-training-box), I kicked it off:\n\n``` bash\ngiles@poppy:~/Dev/ddp-base-model-from-scratch (main)$ uv run torchrun --nproc_per_node=1 ddp_train.py 1xrtx3090-moe-first-proper-run datasets/\nFetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 3517.97it/s]\nmoe_router_loss_scale=0.005\nStarting rank 0 training at global step 0\n 0%| | 0/90823 [00:07