VerySimpleMoE

A working PyTorch implementation of a tiny decoder-only Mixture-of-Experts language model trained from HuggingFace FineWeb.

Architectures

The original architecture is preserved as the v1 preset:

12 experts total per MoE layer
6 active experts per token (top-k=6 router)
each expert has exactly 500,000 parameters
- expert MLP: Linear(500 -> 500, bias=False), GELU, Linear(500 -> 500, bias=False)
- params: 500*500 + 500*500 = 500,000
GPT-style causal self-attention before the MoE block
learned top-k router plus load-balancing auxiliary loss
GPT-2 tokenizer by default

The new research architecture is available as v2-32x1m:

32 experts total per MoE layer
4 active experts per token by default (top-k=4 router)
each expert has exactly 1,000,000 parameters
- expert MLP: Linear(500 -> 1000, bias=False), GELU, Linear(1000 -> 500, bias=False)
- params: 500*1000 + 1000*500 = 1,000,000
router noise during training for exploration
router z-loss to keep router logits stable
optional phased expert training so only a subset of experts are trainable/routeable at a time

The default still uses n_layers=1, which means the model has one MoE block. If you raise --n-layers, each layer gets its own full expert set.

Install

python3 -m venv .venv
source .venv/bin/activate
pip install -e .

Train v1 on FineWeb

This streams FineWeb, so it does not download the full dataset first.

verysimplemoe-train \
  --arch v1 \
  --dataset-name HuggingFaceFW/fineweb \
  --dataset-config sample-10BT \
  --out-dir checkpoints/verysimplemoe-v1 \
  --max-steps 1000 \
  --batch-size 8 \
  --grad-accum-steps 4 \
  --block-size 256

Train v2: 32 experts, 1M params per expert

Recommended laptop-friendly command:

verysimplemoe-train \
  --arch v2-32x1m \
  --out-dir checkpoints/verysimplemoe-v2-32x1m \
  --train-experts-per-phase 16 \
  --expert-phase-steps 500 \
  --batch-size 4 \
  --grad-accum-steps 4 \
  --block-size 256

Phased expert training means:

only 16 of the 32 experts are eligible for routing in a given phase
only those 16 experts have gradients enabled
the expert optimizer is rebuilt per phase, so Adam state is kept only for the current expert subset
the phase window overlaps by default using a half-window stride, e.g. 0-15, 8-23, 16-31, 24-31 + 0-7

Useful router options:

--active-experts 4          # top-k experts per token
--router-noise-std 0.1      # train-time router exploration
--router-z-loss-coef 1e-4   # router logit stabilization
--aux-loss-coef 0.01        # load-balancing loss

Tiny CPU smoke test

verysimplemoe-train --device cpu --max-steps 5 --batch-size 1 --grad-accum-steps 1 --block-size 64

For a tiny v2 smoke test:

verysimplemoe-train \
  --device cpu \
  --arch v2-32x1m \
  --train-experts-per-phase 16 \
  --max-steps 5 \
  --batch-size 1 \
  --grad-accum-steps 1 \
  --block-size 64

Useful speed flags on NVIDIA GPUs:

verysimplemoe-train --amp --compile --max-steps 10000 --batch-size 16

Note: with phased expert training, torch.compile may recompile when the active expert phase changes.

Resume after interruption

If training stops after a checkpoint save begins, resume from the checkpoint directory:

verysimplemoe-train \
  --resume-from checkpoints/verysimplemoe-v2-32x1m \
  --out-dir checkpoints/verysimplemoe-v2-32x1m \
  --max-steps 1000 \
  --batch-size 4 \
  --grad-accum-steps 4 \
  --block-size 256 \
  --train-experts-per-phase 16

--max-steps is the final target step count, not additional steps.

Generate with the final model

verysimplemoe-generate \
  --checkpoint checkpoints/verysimplemoe-v2-32x1m \
  --prompt "The future of open language models is" \
  --max-new-tokens 120 \
  --temperature 0.8 \
  --top-k 50

You can also run modules without installing scripts:

PYTHONPATH=src python -m verysimplemoe.train --max-steps 100
PYTHONPATH=src python -m verysimplemoe.generate --prompt "Hello"

Files

src/verysimplemoe/model.py - model, MoE router, experts, generation
src/verysimplemoe/train.py - FineWeb streaming trainer, phased expert training, checkpointing
src/verysimplemoe/generate.py - checkpoint loader and text generation CLI

4.3 KiB Raw Blame History