Moe Guide Scale
Captured source
source ↗Cerebras Skip to main content
Cerebras Announces First Quarter 2026 Results >>
Sep 03 2025 MoE at Scale: Making Sparse Models Fast on Real Hardware Daria Soboleva Quentin Anthony
MoE Fundamentals | Router Wars | Debugging Dead MoE Models | MoE at Scale | MoE Math Demystified
In this video we discuss scaling MoE models on modern hardware and address key optimization challenges. If you can’t open the video displayed above, please use this link to open it on YouTube: https://youtu.be/MXo9LEYzwkg Mixture-of-Experts (MoE) models allow you to increase total parameter count without proportional increase in compute, letting you train bigger and better models efficiently (Soboleva, 2025a). You might wonder if extracting theoretical benefits from MoE models requires significant engineering work. After all, your part 3 implementation (Soboleva and Tiwari, 2025) trained perfectly fine on a small acceleration node (and even your laptop). An important point here is that you used only 4 experts and 124M backbone parameters, but production systems like DeepSeek-V3, Qwen3, etc., use hundreds of experts and huge backbones. Try scaling to their sizes with our previous implementation on the GPU, and you will quickly hit your device’s memory limit.
Let’s understand why this happens. Remember the code on Figure 1 that we implemented in part 3? There we warned you that it isn’t efficient and we only used it for pedagogy. Why? On Figure 1, you can see that there is a sequential loop over all experts, even though we only need the experts that were selected for a given token. You might wonder, why so wasteful? You already know that the industry standard routing is not deterministic (Soboleva, 2025b), and thus it is impossible to predict which expert will be activated for a given token in advance. This means that we have to load all experts into memory just in case we might need them later (moving expert weights to the GPU memory on the fly for each token would be too slow given memory-bandwidth and latency limits). Because of that, we have a linear growth in memory requirements as you add more experts, and super quickly it becomes impossible to train on a single GPU (even a dozen of experts won’t fit with our previous implementation).
The rest of this guide covers solutions to this scaling challenge. First, we will talk about GPU solution and its limitations, second, we’ll cover how Cerebras Wafer Scale Engine (WSE) solves this problem, and finally how to combine both to maximize the benefits. Figure 1: Expert mixing for batch processing (see train train_gpt moe.pyL208-L225 ). GPU Solution and Its Problems When your MoE exceeds a single GPU’s memory capacity, the most popular solution is to use some sort of model parallelism, an example of which is expert parallelism, or EP in short (Lepikhin et al., 2020; DeepSeek-AI et al., 2024). With EP, we typically assign an equal number of experts to each GPU, while all other layers (attention, router, etc.) are replicated on each device. In this workflow, the router predicts which tokens should be routed to which experts, then an all-to-all communication operation shuffles tokens to the correct devices based on the router assignments, executes each expert in parallel and then shuffles results back to tokens’ original devices to perform expert mixing. Given that we shard only expert layers across devices, for all other layers GPUs are doing repetitive work.
To make EP efficient, your experts have to be load-balanced across all layers in the network. This creates a tension between model quality and infrastructure efficiency. ML researchers optimizing for the best possible model quality often prefer routing strategies that allow experts to specialize heavily (Soboleva, 2025b), even if this creates load imbalance. Meanwhile, infrastructure teams need predictable, balanced workloads to achieve optimal hardware utilization, and minimize training costs. Both goals are critical, but they pull in different directions. Aggressive load balancing can hurt the model’s quality by forcing tokens to suboptimal experts (like the case we’ve seen with hash routing). But imbalanced experts create expensive infrastructure bottlenecks, where some GPUs sit idle, while others are overloaded, significantly decreasing overall hardware utilization. Additionally, EP introduces communication overheads that worsen as we increase number of experts. Modern MoE architectures activate many small experts per token rather than fewer but larger ones for better compute efficiency (Krajewski et al., 2024). However, in this setup, intensive load balancing is required to parallelize with EP, which impacts model quality. If you can’t trade off quality for speed with aggressive load balancing, EP’s communication/computation ratio worsens, as we spend most of the time moving tokens around, rather than experts performing useful operations with them. Overall, training MoE models on GPUs remains challenging, as EP alone is typically not enough to generate sufficient parallelism. Cerebras WSE vs GPU (Architectural Differences) GPUs running large MoE models typically need multiple parallelism strategies in addition to expert parallelism. For example, DeepSeek-V3 used a combination of pipeline, expert, and data parallelism (DeepSeek-AI et al., 2024). This creates a complex 3D parallel implementation that you must carefully tune, and that you must re-tune when either your model or the cluster changes. With Cerebras WSE, distributed computing is not required. We utilize data parallelism to train and scale our models. What makes this possible? We have about 900 times more on-chip memory (SRAM) than a latest single GPU, which allows us to store much bigger models on the chip directly (roughly up to 1B in total parameter count). When scaling to larger models, we employ a technique called weight streaming (Hall et al., 2023) that disaggregates memory and compute on WSE. With weight streaming, we remove model parameters (those heavy tensors) from the wafer entirely. They now live in the external memory units, and we stream them to the wafer during training to compute gradients. The wafer streams gradients back to the memory units to update the weights. This technique allows us to train today’s trillion parameter MoE models (Kimi Team, 2025) on just a single device. Memory Problem Solved, Now Compute We’ve established that Cerebras WSE can fit large MoE models on the chip directly with weight streaming and no...
Excerpt shown — open the source for the full document.
Notability
notability 5.0/10Substantive research post on MoE technique by Cerebras.