MoE-Study-Remastered

---

Overview

A 201M-parameter sparse Mixture-of-Experts language model (69M active per token), trained from scratch in 21 hours on a single RTX 4060 with 8GB of VRAM, on 645M tokens of Ultra-FineWeb-L1. It is a corrected rebuild of my earlier MoE-Study: the auxiliary load-balancing loss is now averaged and scaled, positional encoding is RoPE, the FFN is SwiGLU, normalization is RMSNorm, and the tokenizer is a custom 32,768-vocab BPE. Evaluated with lm-evaluation-harness v0.4.12 against pythia-160m at a matched token budget.

Where MoE-Study trained a dense and a sparse model side by side, this study drops the dense arm and focuses on the MoE alone. The goal was as much to correct the mistakes and un-optimized structures of that first attempt as to deepen my own understanding of them. For the wider context, MoE-Study is linked above and worth reading first.

Full List of Modifications

• Averaging aux loss: In MoE-Study, during training, all MoE layers in the LLM calculated their aux loss in an effort equalize tokens per expert. This loss was correctly summed over all MoE layers during training; however, it was then mistakenly added directly onto the cross entropy loss, without averaging it across all MoE layers. This caused the total loss to be extremely high, with the summed aux loss dominating the cross entropy loss, which is unconventional in LLM training.

• Scaling down aux loss: In MoE-Study, during training, not only was the aux loss summed and not averaged, but it was also not scaled down by a small hyperparameter. This, on top of the summed aux loss, caused the total loss to be catastrophically high (in the mid 10s), ultimately overshadowing the cross entropy loss.

• *Implementing RoPE (Rotary Positional Embeddings):* In MoE-Study, both models were capped to a context length of 1024 due to the positional embeddings having a size cap of 1024. For training this caused no major issues, but at inference it capped any conversation at 1024 tokens. RoPE removes the learned-table cap: this run still builds a 1024-position cache to match the training context, but the cache can be rebuilt or rescaled at longer lengths without retraining the embeddings.

• *Using SwiGLU (Swish Gated Linear Unit) instead of GELU (Gaussian Error Linear Unit):* SwiGLU is an advanced activation mechanism to replace traditional activations like ReLU or GELU. It combines the Swish (SiLU) activation function with a Gated Linear Unit (GLU) structure, introducing a multiplicative gating mechanism that allows the LLM to dynamically route and scale features to produce accurate token predictions.

• *Using RMSNorm (Root Mean Square Layer Normalization):* In MoE-Study, the models used a custom-built Norm class for normalization. That worked, but PyTorch ships nn.RMSNorm, a fused and well-optimized implementation that is also the standard choice in modern LLMs.

• Using a custom BPE (Byte Pair Encoding) tokenizer: In MoE-Study, the training used the GPT2 tokenizer, which had a vocab size of 50257. Lowering the vocab size to 32768 shrinks the tied token-embedding/output matrix and cuts the cost of the final softmax.

• Training on Ultra-FineWeb-L1: In MoE-Study, both models trained on data from nampdn-ai/tiny-textbooks. This dataset was far too small to train a model of this size to any meaningful quality.

• Reduced few-shot count for hellaswag and arc_challenge: In MoE-Study, hellaswag was evaluated at 10-shot and arc_challenge at 25-shot. Both prompts overflowed the 1024-token context and were truncated, which can cut off the question itself and skew the scores. Here they run at 5-shot and 15-shot respectively, so every prompt fits.

At a glance

Architecture	RoPE, SwiGLU FeedForward, MoE, Multi-Query Attention, decoder-only transformer
Total parameters	201,267,712
Active parameters / token	69,147,136 (~2.9x sparsity)
Experts / top-k	8 / 2
Context length	1024
Training data	openbmb/Ultra-FineWeb-L1 [CC-Main-2025-30], 500M words
Hardware	NVIDIA GeForce RTX 4060, 8GB VRAM
Training time	75,871s (21.08 h)
Weights	Model

---

Research questions

Can one person engineer and locally train a LLM on a RTX 4060 that holds its own on standard benchmarks against other models in the same weight class?

At a matched trained-token budget, does a sparse top-2 MoE (201M total / 69M active parameters) outperform a dense model of comparable size?

---

Key results

Benchmark	Few-shot count	Main metric	This model	pythia-160m (step256)
arc_easy	0	acc_norm	35.44	28.37
piqa	0	acc_norm	62.35	51.58
wikitext	0	word_perplexity	75.20	2101.14
lambada_openai	0	acc	19.02	0.00
winogrande	5	acc	51.22	50.36
hellaswag	5	acc_norm	29.22	25.46
arc_challenge	15	acc_norm	22.78	24.15
Inference speed	N/A	tokens/sec	17.92	108.69

Headline finding: Against pythia-160m at step256—a dense baseline matched on trained tokens (537M vs 645M)—this model wins 5 of 7 benchmarks, ties 1, and loses 1. The clearest gaps are language modelling: wikitext word perplexity 75.20 vs 2,101.14 and lambada_openai 19.02% acc vs 0.00% (perplexity 201 vs 766,416). It also leads on arc_easy (35.44 vs 28.37 acc_norm), piqa (62.35 vs 51.58) and hellaswag (29.22 vs 25.46); winogrande is a tie inside one standard error (51.22 vs 50.36) and arc_challenge is a loss (22.78 vs 24.15). The cost is speed: 17.92 tok/s against 108.69, a 6.1x gap that comes entirely from having no KV cache.

---

Architecture

Component	Setting	Why
Layers	14	Set model depth. Chosen to give reasonable representational depth while inside the 8GB VRAM budget.
Embedding dim	512	Set model width. Keeps the token-embedding/output table and expert FFN weights small enough to fit in 8GB VRAM.
Attention heads	8	With emb_dim=512, this yields an even head dim of 64.
Attention type	Multi-Query Attention	A single shared key-value projection (sized head_dim) is used across all query heads, cutting K/V projection parameters and KV memory 8x versus multi-head attention.
FFN hidden dim	1024	2x emb_dim. Since the router selects top 2 experts, the hidden dim was kept below the typical 4x-dense ratio so all 8 expert weights fit in VRAM.
Experts	8	Amount of specialized FFN sub-networks, chosen to fit GPU VRAM and the top-2 router.
Router top-k	2	Number of experts each token is routed to. Two gives the router a blend of specialists per token while keeping active parameters low.
Load-balancing loss weight	0.01	Scales the router's load-balancing aux loss before it is added to the cross-entropy loss; kept small so expert load evens out gradually without dominating the primary loss.
Context length	1024	Sets the length of the RoPE cache and the quantity of tokens per step fed to the model during training—limited by VRAM.
Vocab size	32768	Size of the custom BPE tokenizer, chosen to shrink the tied embedding/output matrix versus GPT-2's 50257-token vocab, directly cutting parameter count and VRAM use.

---

Training setup

Data

Source	openbmb/Ultra-FineWeb-L1
Words used	500,000,000
Tokens	657,681,194
Train / val / test split	98% / 1% / 1%
Tokenizer	Tokenizer
Preprocessing	Wrote tokens as bytes into .bin, then read with memmap

Hyperparameters

Optimizer	AdamW
Learning rate	6e-4
Schedule	OneCycleLR, 0.03 * total_optim_steps
Weight decay	0.1 (params with dim > 1)
Grad clipping	1.0
Batch size	2 × 16 accum = 32
Precision	bfloat16
Epochs / steps	1 / 19,669 optimizer steps (314,705 micro-batches)
Seed	42

Hardware and cost

GPU	NVIDIA GeForce RTX 4060 Dual, 8GB VRAM
Peak VRAM used	6.7GB
Throughput	~4.15 micro-batches/sec, ~8,500 tok/s
Total training time	75,871s (21.08 h)

Final losses

	Train	Val	Test
Loss	3.2815	3.3155	3.3070

Train is the mean of the final 100 optimizer steps

---

Evaluation

Harness: lm-evaluation-harness v0.4.12

Benchmark	Few-shot count	Main metric	What it measures
arc_easy	0	acc_norm	Easy science questions. Tests basic reasoning.
piqa	0	acc_norm	Everyday physical commonsense. Which action works.
wikitext	0	word_perplexity	Raw language modeling. Lower is better.
lambada_openai	0	acc	Predicts the last word of a passage. Needs context.
winogrande	5	acc	Pronoun resolution. Needs commonsense.
hellaswag	5	acc_norm	Picks the most likely next sentence. Tests commonsense.
arc_challenge	15	acc_norm	Harder science questions. Requires advanced reasoning.
Inference speed	N/A	tokens/sec	How fast the model generates text. Not accuracy.

Benchmark	Metric	This model	pythia-160m
arc_easy	acc	37.88	27.40
piqa	acc	62.68	53.21
wikitext	byte_perplexity	2.243	4.181
wikitext	bits_per_byte	1.166	2.064
lambada_openai	perplexity	201.27	766,416.04
hellaswag	acc	27.95	25.86
arc_challenge	acc	17.83	18.34

Baseline: EleutherAI/pythia-160m [step256]. I chose this model for two reasons. First, its 160M parameters are close to this model's 201M. Second, pythia-160m was trained on 300B tokens but EleutherAI published intermediate checkpoints, and at step256 it had seen roughly 537M tokens — close to this model's 645M.

Speed measurement: Both models were given a 128-token prompt and greedily decoded for 256 tokens, hand-written to run the identical loop for each—no generate(), no KV cache. The custom LLM has no cache to begin with, so the pythia baseline was forced to use_cache=False to match it; both therefore recompute the full prefix on every step, which understates the baseline's real speed substantially. One warmup round was discarded to absorb CUDA context setup and kernel autotuning, then 5 rounds were run with the two models interleaved—each round timed one full pass of both, rather than 5 passes of one model followed by 5 of the other. The reported figure is the median decode-phase tokens/sec across those 5 rounds: 17.92 tok/s (min 17.36, max 18.18) for this model against 108.69 tok/s (min 106.25, max 109.67) for pythia-160m. Full per-round numbers are in speed.json.

---

Results

The model wins 5 of 7 benchmarks—wikitext word perplexity 75.20 vs 2,101.14, lambada_openai acc 19.01 vs 0.00, arc_easy acc_norm 35.44 vs 28.37, piqa acc_norm 62.35 vs 51.58, and hellaswag acc_norm 29.22 vs 25.46. Both models tie on winogrande (51.22 vs 50.36 acc), and pythia wins arc_challenge (22.78 vs 24.15 acc_norm). However, this model sacrifices speed: 17.92 tok/s against pythia's 108.69, a 6.1x gap.

---

Challenges and debugging

• Incorrect SwiGLU FeedForward Module: After finishing what I thought would be the one and only training run, I noticed the FeedForward module had implemented SwiGLU incorrectly — the SiLU was applied to the product of the gate and up projections instead of to the gate alone:

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.gate = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False)
        self.up = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False)
        self.down = nn.Linear(cfg["hidden_dim"], cfg["emb_dim"], bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x) * self.up(x)))

So I fixed the module and re-ran training from scratch—another 20+ hours—to get the benefit of a correctly gated SwiGLU:

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.gate = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False)
        self.up = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False)
        self.down = nn.Linear(cfg["hidden_dim"], cfg["emb_dim"], bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

---

Limitations

• No ablations. Every change listed above was made before a single training step was run, so it is unknown which of them actually improved the training loss and by how much.

• Undertrained. 644,516,564 training tokens against 201,267,712 parameters is 3.20 tokens per parameter (9.32 per active parameter), well under the roughly 20:1 ratio a model this size should have. Training was a single epoch of 19,669 optimizer steps, and it stopped because the training tokens ran out.

• Scale. This model is limited to my 8GB of VRAM. This includes: Batch Size, Embedding Dim, Hidden Dim, N-Layers, N-Heads, Context Length, Tokens Trained, N-Experts, etc.

• Non-standard few-shot counts. Every benchmark ran on its full split, but the few-shot counts deviate from leaderboard convention—hellaswag at 5-shot instead of 10 and arc_challenge at 15-shot instead of 25—because the 1024-token context truncates the standard prompts.

• No KV cache. LLM.forward returns logits and no past_key_values, so generation recomputes the entire prefix at every step. This is why the model measures 17.92 tok/s against pythia-160m's 108.69, and why the baseline had to be forced to use_cache=False to keep that comparison equal.

• Base model. The model has only been pre-trained—no fine-tuning and certainly no RLHF.

---

Model weights

Weights and tokenizer: huggingface.co/OliverSundaram/MoE-Study-Remastered

The checkpoint is in Hugging Face format—config.json, model.safetensors, the tokenizer, and a copy of modules.py—so from_pretrained works. generate() does not. LLM subclasses PreTrainedModel without GenerationMixin, so model.can_generate() returns False and the method does not exist on the object. Sampling has to be written by hand:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

torch.manual_seed(42)

MODEL = "OliverSundaram/MoE-Study-Remastered"  # or a local runs/checkpoint/final

tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, trust_remote_code=True)
model.to("cuda").eval()


@torch.inference_mode()
def generate(prompt, max_new_tokens=60, temperature=0.8, top_k=50):
    ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")

    for _ in range(max_new_tokens):
        # No KV cache, and the RoPE table stops at 1024: crop to the last 1024 ids.
        logits = model(ids[:, -1024:]).logits[:, -1] / temperature

        top_logits, top_ids = torch.topk(logits, top_k, dim=-1)
        probs = torch.softmax(top_logits, dim=-1)
        next_id = top_ids.gather(-1, torch.multinomial(probs, 1))

        ids = torch.cat([ids, next_id], dim=1)

    return tokenizer.decode(ids[0])


print(generate("To make a cup of tea, you first"))

trust_remote_code=True is required: model_type is custom_llm, and the architecture is loaded from the modules.py shipped inside the checkpoint.

Sample output

Both samples come from the snippet above, at temperature=0.8, top_k=50, seed 42.

1.

Prompt: To make a cup of tea, you first

To make a cup of tea, you first need to start by making a cup:
- Preheat your oven to 350 degrees Fahrenheit.
- Combine the milk and tea powder with the water to a frother, stirring gently.
- Add the milk and tea to the water. Gradually add the milk and tea mixture to the water

2.

Prompt: The Industrial Revolution began in Britain because

The Industrial Revolution began in Britain because of the American Revolutionary War
(1940-1979). Many cities had to move to accommodate workers, and some of the most
significant buildings in the city were now being built.
The American Revolution, then, was an experiment for the United States. It was a
revolutionary experiment that brought a

The model answers both prompts comically badly, but it clearly did learn something: fluent English with correct syntax—bullet lists, dates, and consistent capitalization throughout.

---

Citation

@misc{moe-study-remastered,
  author       = {Oliver Sundaram},
  title        = {MoE-Study-Remastered: A Corrected Sparse Mixture-of-Experts Language Model, Trained From Scratch},
  year         = {2026},
  publisher    = {GitHub},
  howpublished = {\url{https://github.com/OliverSundaram/MoE-Study-Remastered}}
}

---

Acknowledgements and references

• lm-evaluation-harness — evaluation framework used for all benchmark scores.
• openbmb/Ultra-FineWeb-L1 — training corpus.
• Ultra-FineWeb: Efficient Data Filtering and Verification for High-Quality LLM Training Data — paper describing the filtering pipeline behind the dataset above.
• RoFormer: Enhanced Transformer with Rotary Position Embedding — source of the RoPE implementation used for positional encoding.
• GLU Variants Improve Transformer — introduces SwiGLU, used in the FeedForward/expert module.
• Root Mean Square Layer Normalization — RMSNorm, used in place of a custom norm layer.
• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer — origin of the sparsely-gated top-k MoE layer and load-balancing objective this model's router is based on.
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity — informed the auxiliary load-balancing loss formulation and the fix to how it's scaled/averaged.
• Hugging Face Transformers — PreTrainedModel/PretrainedConfig base classes, tokenizer utilities, and the evaluation/inference tooling around the model.
• Hugging Face Tokenizers — trained the custom byte-level BPE tokenizer used for this model.
• PyTorch — training and model implementation.
• EleutherAI/pythia-160m — Pythia model on Hugging Face, used as the evaluation baseline.
• Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling — EleutherAI paper.

---

License

MIT — see LICENSE.