No Cluster, Single-GPU experiments using Moedl, an MoE implementation built on Hugging Face Transformers
Jump to
- Addressing the elephant first, which Load Balancing Strategy? This animated heatmaps settles the debate.
- Hit the ground running: Hands-on guide.
- Smaller Experts, More of Them: Resolution and Granularity.
- Shared Experts? Probably not compelling enough to be worth the effort.
- Token Dropping is pretty much irrelevant with proper load balancing. So Don't!
- Scaling #Experts: Well ablated in literature while we run into anomaly which requires further investigation.
- Description of
MoedlandMoedlTrainer - Implementing Autograded Grouped GEMM and Speedup
- Future Plans
2025 is the year of reasoning and agents. It can also be argued as the year Mixture-of-Experts (MoE) models truly hit the mainstream. Virtually every flagship model from frontier labs is an MoE, although Google has been pioneering and popularizing the idea since 2017.
The beginning of 2025 was marked by the release of DeepSeek R1, which demonstrated reasoning learned through pure RL with verifiable feedbacks. Its backbone is DeepSeek V3, a 671B MoE. Qwen3 MoE arrived mid-year, followed by OpenAI's GPT-OSS in August. The year closed with models such as Kimi-K2 and Mistral Large 3, which largely adopt DeepSeek V3–style architectures with larger number of experts.
Personally, I am particularly interested in understanding which components of MoE models contribute most to their efficacy beyond just scale and the conditional compute. My initial questions centered on how effective the load-balance biasing strategy used in DeepSeek V3 is compared to loss-based penalties. While DeepSeek-V3 (Section 4.5.3) ablations verified better model performance with load-balance biasing and improved specialization, specialization inherently creates "hot" experts. I am particularly curious about the generic load-balancing delta between the two, specifically how their routing distributions evolve over the course of training. I also questioned whether shared experts are truly mandatory, given that they introduce asymmetry and potentially additional engineering effort. While some prior work suggests they are not strictly required, I find the existing evidence not yet compelling (though personally favor avoiding unremarkable additions).
I intended to carry out these ablations using modest resources with HuggingFace (HF) Transformers. However, I could not find an implementation that exposes load balance biasing control. More importantly, there was no single MoE model type that allowed only contrasting a single axis of design choices while keeping most components consistent. For example, when comparing loss penalties versus biasing strategies for balancing expert load, I would like the attention layers to remain identical. DeepSeek V3 uses Multi-Latent Attention (MLA), which differs from the standard multi-head attention (MHA/MQA/GQA) used in models such as OLMoE or Qwen3, using/modifying existing implementations make direct comparisons unhygienic.
As a result, I decided to implement a new model type in local HF Transformers, Moedl (no pun intended!😝). It is designed to enable individual design choices to be toggled on or off while keeping the rest fixed, providing the controlled environment needed for ablation studies.
While some features are still work in progress and certain ablations require larger resources, I believe there is now sufficient material to document the observations and findings.
-
Install: clone and
make install-moelabormake install-dev-moelab -
Test:
make run-tests -
Run:
make <experiment-id>, seeMakefile. Most experiments can fit within a single 80GB GPU. Most took around 3-5 hrs on a RTX Pro 6000 gpu.# Available make targets (Experiments) 00_llama2_ref b20_moedl_e4_k1_4ep e1_moedl_cf_1.0 01_moedl_dense c1_moedl_e8_k1 e2_moedl_cf_1.5 a0_moedl_no_lb c2_moedl_e16_k2 e3_moedl_cf_2.0 a1_moedl_lb_penalty c3_moedl_e32_k4 e4_moedl_cf_2.5 a2_moedl_lb_biasing c4_moedl_e64_k8 gen-tinystories b1_moedl_e2_k1 d1_moedl_s0_k4_e32 gpulist-check-busy b2_moedl_e4_k1 d2_moedl_s1_k3_e31 install-dev-moelab b3_moedl_e8_k1 d3_moedl_s2_k2_e30 install-moelab b4_moedl_e16_k1 d4_moedl_s3_k1_e29 run-tests- More customization: use
moelab_main.pylike we use standard HF script. Dopython moelab_main.py --helpto see options. - Find LR: Appending
--sweep_lr <list of comma-limited lr>tomoelab_main.pywill turn it into learning rate sweep for small number of steps, which can be configured with--sweep_lr_steps <num_steps>. For experiments in theMakefile, just append sweep_lr=1 to the make command. e.g.make c1_moedl_e8_k1 sweep_lr=1. A report will be generated in the output folder and metrics of the sweep are also logged to wandb by default.
-
All runs are shared via W&B project (w/ grouped_mm), with CSV export for quick result lookup.
-
Qualitative Eval per TinyStories.
make gen-tinystories ckpt=roneneldan/TinyStories-33M # official dense GPTNeo model make gen-tinystories ckpt=vuiseng9/01_moedl_dense-0119 # moedl dense model make gen-tinystories ckpt=vuiseng9/c3_moedl_e32_k4-0119 # moedl moe model
Our objectives are:
- To enable controlled ablations of key MoE design choices.
- To support small-scale experiments that require only a single GPU and can be completed within a few hours.
- To use a simple, accessible and familiar stack, entirely within the HuggingFace Transformers ecosystem.
All experiments use TinyStories, a synthetic dataset of short stories restricted to vocabulary typically understood by children of ages 3-4, generated using GPT-3.5 and GPT-4.
We choose TinyStories because the original paper demonstrates that models with as few as ~10M parameters can already learn to generate fluent and logically consistent stories. This allows us to bound both model size and training compute, keeping experiments feasible on a single GPU while still uncovering meaningful ablation trends. MoE models in this repo are typically a few hundred million parameters, mostly around 400M total parameters with 12.5% sparsity, ~135M active parameters per token.
TinyStories is also easy to evaluate qualitatively: story coherence and logical consistency are readily observable, making it practical for comparing generation quality across MoE variants. In contrast, prior experience using GPT-2 or OPT models of similar scale trained on large, generic corpora often results in incoherent or unstructured generation due to capacity limits, making cross-model comparisons rely largely on numerical metrics.
We use the llama2 tokenizer for its smaller 32K vocabs. After tokenization, the training set contains approximately 1B tokens. We limit most experiments to 2 epochs, based on the diminishing returns of longer epochs observed in Scaling Data-Constrained LMs.
This repo primarily focuses on the MoE layers. Since we set sight on 32K llama2 vocab, we might as well inherit the llama (llama2) architecture as the starting point. Its attention stack also reflects the current standard: configurable GQA (MHA), RoPE, RMSNorm, and GLU-style MLPs. As a result, the standard llama configurables are directly applicable to Moedl, including hidden_size, intermediate_size, num_attention_heads, num_hidden_layers, and related parameters.
Moedl can be configured as either a dense or an MoE model. As a sanity check, we train an equivalent llama2 dense model (make 00_llama2_ref) and a Moedl dense model (make 01_moedl_dense) on TinyStories. Equivalent training trajectories ensure architectural parity.
For MoE configurations, the following parameters can be varied:
num_experts: number of experts per MoE layer, denoted as E throughout. If set to 0, a dense layer is used.num_active_experts: number of experts activated per token, denoted as K throughout.num_shared_experts: number of shared experts, denoted as ES throughout.lb_coeff: load-imbalance penalty coefficient. If set to 0, no loss-based penalty is added to the training objective.lb_gamma: router biasing update rate. If set to 0, bias-based control is disabled.- No load balancing is applied when both
lb_coeffandlb_gammaare set to 0. And both are mutually exclusive currently, both larger than 0 throw error. capacity_factor: factor controlling expert capacity threshold for token dropping. Setting this to 0 disables token dropping (i.e., dropless routing).
The interplay between these configurables is here. The exact MoE execution logic can be found in the []MoeBlk forward pass](src/moelab/moedl/modeling_moedl.py#L179). Tests are developed and validated against OLMoE and DeepSeek-V3 to ensure functional correctness across configurations.
We subclass the HF Trainer to add MoE-specific training callbacks for bookkeeping and logging, including routing statistics, expert load tracking, and fine-grained expert load heatmap generation with GIF collation. A key component is the LoadBalanceBiasController, responsible for router biasing control (Eq.3) in load balancing; encourage to review the code.
As of Jan 2026, most HuggingFace-based MoE implementations (including our initial impl.) still iterate over experts and execute them sequentially. This leads to poor GPU utilization: when token count per expert is small or expert dimensions do not saturate tensor cores, serialized expert execution leaves significant compute idle.
Grouped GEMM addresses this by tiling the output of all experts such that these output tiles are scheduled/executed independently and in parallel (with their corresponding inputs). This enables concurrent execution across SMs instead of waiting for one expert to finish before launching the next. Historically this required custom kernels. In PyTorch 2.10 released Jan 2026, torch.nn.functional.grouped_mm provides an official primitive.
To enable training, we implemented an autograded GroupedMMFunc and GroupedGLU module to wrap around grouped_mm. See codes. While the integration is straightforward, a few practical notes are worth highlighting:
-
Partial coverage of expert GEMMs. Synonymous to linear layer, experts requires 3 grouped GEMMs
Y, grad_x, grad_w. The currentgrouped_mmAPI assumes a fixed contraction (inner) dimension across groups. This allows us to implement theyandgrad_xGEMMs via appropriate transpose, but notgrad_ybecause the contraction dim is the token dim, and token counts per expert vary dynamically. As a result,grad_yare still computed per expert in a loop. In essense, we accelerate 2/3 expert GEMMs. -
The execution pattern of
GroupedGLU: Usinggrouped_mm, gate, up and down projections are grouped and dispatched in parallel across all activated experts. This replaces the vanilla expert-by-expert execution. -
Reduced precision:
grouped_mmreturns outputs in the dtype of input A. Under our BF16 training regime, per-expert outputs are BF16. This slightly affects the weighted aggregation over top-k experts due to rounding. Across all experiments, final eval loss degrades by ~0.85% on average, while ablation trends and relative comparisons remain stable.
Acceleration Results on 1xRTX Pro 6000: Across all experiments, GroupedGLU yields:
- 23% average training (incl. eval) speedup, up to 46% speedup in high-expert settings.
The benefit is strongly correlated to expert count. The more experts you have, the more wasteful sequential execution becomes, so grouped execution helps more.
TLDR: Router biasing is surprisingly effective, easier to implement, requires less tuning!
Load balancing is a fundamental requirement for MoE models as it directly determines whether the underlying devices are utilized effectively. By load balance, it simply means router's ability to distribute the incoming tokens evenly across all experts. Proper load balance allows computational work to be shared evenly across experts/devices, enabling parallelism and improving both training and inference efficiency.
In practice, however, load balancing is not built-in. The standard training objective applied on MoEs provides no explicit incentive for routing tokens evenly. The router is optimized to minimize training loss, which can naturally lead to imbalanced routing, where a small subset of experts becomes overloaded while others remain underutilized (as we will see in our ablations). So how should load balancing be enforced in MoE models?
Google has pioneered the use of an auxiliary loss (
where
-
$E$ is the number of experts -
$f_e$ is the fraction of tokens routed to expert$e$ -
$p_e$ is the average router probability assigned to expert$e$ (router's softmax output, averaged over tokens and$k$ experts)
In essense, the auxiliary loss backpropagates through the router probabilities, shaping the router's gradients to discourage imbalanced expert utilization during training. Conceptually, this resembles a form of regularization, applied to routing behavior rather than model parameters.
In our implementation, lb_coeff.
2. Router Load Biasing (DeepSeek v3)
Dubbed as auxiliary loss-free load balancing, DeepSeek V3 introduced an alternative approach that directly modifies the router logits (sigmoid output) with an additive load-biasing term before the Top-k selection of experts. This biasing term is updated periodically based on the observed expert loads, effectively nudging the router to favor less-utilized experts.
where
-
$s \in \mathbb{R}^{E}$ is the biased router score vector over all experts (per token). -
$i_K$ is the indices of selected experts and$s_K$ is thier respective score in$s$ . -
$b_{lb} \in \mathbb{R}^{E}$ is the load-balancing bias vector, initialized to$\mathbf{0}$ and updated periodically as:
where
-
$\mathbf{f} \in \mathbb{R}^{E}$ is the observed expert-load fraction averaged over$N$ steps. Simplest would be N=1, i.e. after every optimizer stepping. -
$\bar{f} = \tfrac{1}{E}\mathbf{1}$ is the uniform target load. -
$\gamma$ is the bias update rate.
It is important to note that
Setup: We ablate on Moedl with 8 experts (E=8) and 1 active expert per token (K=1) on TinyStories for 2 epochs. We compare no load balancing (baseline), imbalance penalty, and router biasing.
| make [exp. id] | Eval Loss |
|---|---|
a0_moedl_no_lb |
1.127 |
a1_moedl_lb_penalty |
1.137 |
a2_moedl_lb_biasing |
1.130 |
At first glance, all three strategies converge to similar final eval loss, with the best result achieved without load balancing, followed closely by router biasing, and finally the imbalance penalty. While the final loss differences are small, the load-balancing dynamics differ. To illustrate this, we examine expert load over time. The following plots, logged in W&B, show expert load averaged across layers.
Without load balancing, expert imbalance is immediately apparent. A small subset of experts dominates the routing, as reflected by the disproportionate heights in the stacked plot and the uneven distributions in the % overlay view. Observe how e002 and e006 are less assigned than the rest throughout training.
Imbalance penalty improves load distribution gradually over training. All series fluctuate around balance point 1/E=0.125 throughout. Between steps 2k and 10k, the % plot exhibit noticeably higher variance, indicating noisy and unstable routing before the auxiliary loss sufficiently regularizes expert utilization.
Router biasing yields the most stable behavior. Expert load converges rapidly toward a uniform distribution, with obvious tighter variance bounds in the % plot throughout training.
At this point, router biasing edges out the imbalance penalty (lower eval loss, tighter distribution variance), but not convincingly so. The take-my-money moment comes next: by examining distribution at granular level, i.e. per-expert, per-layer load deviations from balance point. We visualized them as heatmaps animated over training.
Examining the animated heatmaps, the advantage of router biasing becomes strikingly clear. Notice how plain and less hot the color distribution remains throughout training. Router biasing rapidly achieves near-perfect uniform balance across experts and layers, with minimal deviation over time. In contrast, the imbalance penalty shows signs of expert collapse at the later stages of training, where certain experts (L7E6, L1E7, L6E1) remain consistently overloaded while others are underutilized. As expected, the no-load-balancing baseline exhibits imbalance hotspots across layers throughout training.
Why does router biasing work better? The auxiliary loss in Eq.1 is a globally reduced scalar objective, a few localized imbalance signals may be too weak to meaningfully influence the overall loss. Router biasing (Eq.2&3), by contrast, applies control directly at a per-router level. Each expert is adjusted independently via a dedicated bias term, enabling more precise and effective correction.
While expert-specific coefficients could be introduced for the imbalance penalty, doing so requires additional tuning. Router biasing is simpler to implement and requires minimal tuning in practice. In my experience, tuning is straight forward, basically ensuring the bias update rate
Based on these results, we adopt router biasing as the default load-balancing strategy for the remaining ablations.
With a stable and effective load-balancing strategy in place, we now turn to the question of scale. Specifically, we examine how increasing the number of experts impacts overall model performance. Following Switch Transformer study, we fix K=1, one activated expert per token to keep computation roughly constant to dense counterpart (aside from minor router overhead) and scale the number of experts by doubling from 1 to 16.
| make [exp. id] | # Params | # Experts (E) | Eval Loss |
|---|---|---|---|
01_moedl_dense |
137M | 1 | 1.111 |
b1_moedl_e2_k1 |
175M | 2 | 1.118 |
b2_moedl_e4_k1 |
251M | 4 | 1.125 |
b3_moedl_e8_k1 |
402M | 8 | 1.130 |
b4_moedl_e16_k1 |
704M | 16 | 1.138 |
Contrary to most ablations in the literature (e.g., Fig. 1 in Switch Transformer), we observe monotonic performance degradation as expert count increases. We believe this is largely due to the sample inefficiency of sparse models. With a fixed budget of 1B tokens (TinyStories), increasing
- Iso-tokens-per-expert: Train the MoE for E$\times$ more epochs so that each expert sees roughly the same amount of tokens as the dense model. If this fixes the loss, it validates undertraining and data starvation (dataset is probably too small).
- Scale E by dividing dense capacity: An alternative scaling strategy that increases expert count while reducing per-expert capacity, potentially alleviating data starvation while keeping the overall ablation compute manageable.
- Proportional Data Scaling: Simply throw more data at it and move out of the TinyStories regime. This is conceptually straightforward, but deviates from our original goal of small-scale ablations.
We will revisit this once we improve the underlying kernel efficiency. At the moment, MoE layers are implemented by naively looping over experts (the standard HF approach). We plan to integrate a more efficient grouped GEMM implementation. 26'Feb 14: Grouped gemm integrated. More experiments to follow.
DeepSeekMoE is among the first to propose fine-grained expert segmentation, also referred to as expert granularity. The core idea is to use smaller experts but more of them. We supplement granularity with the term resolution, as it directly reflects the actual E:K config, and implicitly conveys the sparsity ratio.
The effectiveness of higher resolution can be reasoned about combinatorially. For a fixed sparsity ratio and approximately constant total model parameters, increasing resolution dramatically increases the number of possible expert combinations (
| make [exp. id] | Expert Dff | Resolution (E:K) | Combinations ( |
Eval Loss |
|---|---|---|---|---|
c1_moedl_e8_k1 |
2048 | 8:1 | 8 | 1.1296 |
c2_moedl_e16_k2 |
1024 | 16:2 | 120 | 1.0658 |
c3_moedl_e32_k4 |
512 | 32:4 | 35,960 | 1.0581 |
c4_moedl_e64_k8 |
256 | 64:8 | 4,426,165,368 | 1.0584 |
We ablate MoE models at a fixed sparsity ratio of 12.5% while increasing resolution: E:K = 8:1, 16:2, 32:4, and 64:8. Total model parameters are kept approximately constant at 400M by adjusting the expert hidden size. Activated param count stays constant as well.
The benefit of higher resolution is clearly observed empirically: as resolution increases, model performance improves. However, the gains eventually saturate, we attribute the diminishing returns to under-training of smaller experts or data starvation. Our observed trends are consistent with prior results reported in DeepSeekMoE's Table 1 and OLMoE ablations (Fig. 5).
In addition to finer-grained experts, DeepSeekMoE also advocates the use of shared experts, where a subset of experts is activated for every token. The rationale is that certain forms of common knowledge may be universal, and sharing experts could reduce parameter redundancy and improve learning efficiency.
While this intuition is appealing, OLMoE (4.1.3) highlights a potential confound to the combinatorial argument. Shared experts are essentially fixed routing paths and contribute no additional combinations. At iso-active-expert settings, introducing shared experts effectively reduces the number of unique combinations, which may be counter-productive for model expressivity (see Table below).
Empirical results in the literature are mixed. The proposer DeepSeekMoE (Fig.3,6) reports benefits from shared experts whereas OLMoE (4.1.3) observes limited or negative impact. NVIDIA's ablations show similar convergence behavior with and without shared experts, though in the specific context of upcycling a Nemotron-4 dense model into an MoE. Notably, most prior studies explore shared experts in a binary on/off setting.
Given the lightweight nature of our setup, we are able to explore this design axis more thoroughly by varying the number of shared experts while holding number of active experts and total parameters constant. Our baseline is E=32, K=4, ~400M total parameters with ~135M actived. We vary the number of shared experts (ES) from 0 to 3, adjusting the K and E to maintain a total of 4 active experts.
| make [exp. id] | Shared (ES) | E | K | Active Experts | Combinations ( |
Eval Loss |
|---|---|---|---|---|---|---|
d1_moedl_s0_k4_e32-0119 |
0 | 32 | 4 | 4 | 35,960 | 1.0581 |
d2_moedl_s1_k3_e31-0119 |
1 | 31 | 3 | 4 | 4,495 | 1.0573 |
d3_moedl_s2_k2_e30-0119 |
2 | 30 | 2 | 4 | 435 | 1.0600 |
d4_moedl_s3_k1_e29-0119 |
3 | 29 | 1 | 4 | 29 | 1.1014 |
Results: The overall trend suggests that more shared experts lead to worse performance aligning to the combinatorial argument. Yet the best result is achieved with 1 shared expert (ES=1) by just tiny margin (0.0008) over no shared experts (ES=0). This small gain may be attributed to dataset characteristics such as limited scale or diversity in TinyStories, leading to under-training of non-shared experts.
Personal take: In my view, the benefit of shared experts, if any, appears marginal. Rather than allocating capacity to shared experts, I would prefer to expand the main trunk by widening layers or adding depth to preserve symmetry within MoE blocks. Conceptually, these approaches serve a similar purpose, applying a common transformation to every token.
Moreover, shared experts can be viewed as inherently load-imbalanced, as they are activated for every token by design. The practical ramifications of this behavior remain unclear (at least to me at the moment). As such, we do not recommend shared experts as a default design choice.
Since the early days of MoE research, Google has employed fixed expert capacity, largely because TPUs and the XLA compiler require tensor shapes to be known statically; dynamically varying token counts at runtime are challenging to support. By expert capacity, each expert is allowed to process only a limited number of tokens during routing. Tokens exceeding this capacity are dropped, i.e., they are not processed by any expert.
Concretely, the capacity factor (CF) controls this limit:
expert_capacity = (#Tokens_per_batch * K / E) × CF
In plain terms, CF = 1.0 restricts each expert to the average expected token load under perfectly balanced routing. Increasing CF allows experts to handle more tokens beyond this average. For example, CF = 2.0 permits each expert to process up to 2x the average load, reducing the likelihood of token dropping during imbalance routing.
Following the vein of MegaBlocks, we ablate the effect of token dropping on Moedl with resolution E:K = 8:1 by sweeping the capacity factor (CF) from 1.0 to 2.5 in increments of 0.5. Note that our setup differs in load balancing: we use router biasing, whereas MegaBlocks uses an auxiliary load-balancing objective. Run make e1 .. e4 for the CF sweeps, and refer a2 / b3 for the dropless setting (CF disabled).
Dropping tokens is detrimental. The dropless setting (CF disabled) outperforms all CF configurations. Among the CF runs, the highest CF (2.5) which minimizes token dropping performs best, consistent with the idea that fewer dropped tokens leads to better learning.
To make this concrete, we enclose the total number of dropped tokens over training. With CF=1.0, the model continues to drop tokens throughout training, ending at ~6K+ dropped tokens aggregated across layers (roughly ~1% of all tokens). In contrast, CF=1.5/2.0/2.5 quickly converge to near-zero dropped tokens. This also suggests router biasing already routes tokens effectively, experts do not hit extra expert capacity (higher CF).
A little nuance though, CF=1.5 and CF=2.0 converge slightly worse than CF=1.0. We suspect this is variance, or that early-stage token dropping can have lasting impact on model learning. Regardless, the trend is clear: token dropping is not beneficial. Also crucially, router biasing effectively mitigates expert overload, making capacity-based token dropping unnecessary. So, don't drop tokens!
Our ablations suggest: Use router biasing for load balancing. Prefer MoE with higher resolution - small experts, more of them, but watch for diminishing returns. Skip shared experts and token dropping.
- Grouped GEMM kernel integration for more efficient training and larger-scale ablations. (done Autograded Grouped Gemm around Mid Feb 2026.)
- Revisit scaling number of experts ablations to understand the anomaly we observed.
@misc{chua_moe_lab_2026,
author = {Vui Seng Chua},
title = {moe-lab: Rigorous MoE Design Ablations You Can Run at Home},
year = {2026},
url = {https://github.com/vuiseng9/moe-lab},
note = {Single-GPU, Hugging Face-based ablations of MoE design choices}
}







