OliverSundaram/MoE-Study-Remastered

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.

<details> <summary><b>Full List of Modifications</b></summary>

• —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.

</details>

At a glance

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

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 lambadaopenai 19.02% acc vs 0.00% (perplexity 201 vs 766,416). It also leads on arceasy (35.44 vs 28.37 accnorm), 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 arcchallenge 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

Training setup

Data

Hyperparameters

Hardware and cost

Final losses

Train is the mean of the final 100 optimizer steps

[image]

Evaluation

Harness: lm-evaluation-harness v0.4.12

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.

[image] [image] [image] [image] [image] [image] [image] [image]

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](https://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](https://github.com/EleutherAI/lm-evaluation-harness) — evaluation framework used for all benchmark scores.
• —[openbmb/Ultra-FineWeb-L1](https://huggingface.co/datasets/openbmb/Ultra-FineWeb-L1) — training corpus.
• —[Ultra-FineWeb: Efficient Data Filtering and Verification for High-Quality LLM Training Data](https://arxiv.org/abs/2505.05427) — paper describing the filtering pipeline behind the dataset above.
• —[RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864) — source of the RoPE implementation used for positional encoding.
• —[GLU Variants Improve Transformer](https://arxiv.org/abs/2002.05202) — introduces SwiGLU, used in the FeedForward/expert module.
• —[Root Mean Square Layer Normalization](https://arxiv.org/abs/1910.07467) — RMSNorm, used in place of a custom norm layer.
• —[Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://arxiv.org/abs/1701.06538) — 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](https://arxiv.org/abs/2101.03961) — informed the auxiliary load-balancing loss formulation and the fix to how it's scaled/averaged.
• —[Hugging Face Transformers](https://github.com/huggingface/transformers) — PreTrainedModel/PretrainedConfig base classes, tokenizer utilities, and the evaluation/inference tooling around the model.
• —[Hugging Face Tokenizers](https://github.com/huggingface/tokenizers) — trained the custom byte-level BPE tokenizer used for this model.
• —[PyTorch](https://pytorch.org/) — training and model implementation.
• —[EleutherAI/pythia-160m](https://huggingface.co/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](https://arxiv.org/pdf/2304.01373) — EleutherAI paper.

License

MIT — see LICENSE.