ohara-moe-base-d12
The pretrained mixture-of-experts base model from a nanochat-style run with ohara. 332M total parameters, 85M active per token, trained from scratch on 1.48B tokens of ClimbMix on 2xA100-80GB in 2.35 hours.
This is a base language model, not a chat model. It continues text; it does not answer questions and never stops on its own. For something that responds, use joey00072/ohara-moe-chat-d12.
Use this if you want to run your own finetune from a sparse base.
MoE architecture
8 experts per layer, top-2 routed, all 12 layers. No shared expert (so not DeepSeek-style). Each expert is a SwiGLU of width 1024 -- the dense FFN width divided by top-k, so two experts firing cost exactly one dense feed-forward.
| Total parameters | 332M |
| Active per token | 85M |
| Sparsity | 3.9x |
| Layers / hidden / heads | 12 / 768 / 6 |
| Context | 2048 |
| Vocabulary | 50,304 |
Expert shape: up (1024,768), gate (1024,768), down (768,1024) = 2.36M each,
18.87M per layer, 226.6M across the model (68% of all parameters).
Load balancing is quantile balancing (Jianlin Su; used in Kimi K2/K3): the router bias is solved in closed form each optimizer step from batch statistics. No auxiliary loss, no coefficient to tune. All 12 routers converged (MaxVio ~0.03).
The 8 conversation special tokens are reserved in the vocabulary but never seen during pretraining, so a finetune can use them without resizing the embedding.
MoE vs dense, controlled
Identical FLOPs per token, tokens, data, schedule, learning rates and seed.
| val bits/byte | dense | MoE |
|---|---|---|
| step 250 | 1.1900 | 1.1703 |
| step 1000 | 1.0370 | 1.0170 |
| step 2000 | 0.9458 | 0.9280 |
| final (2827) | 0.9062 | 0.8887 |
A steady ~1.9% advantage that never decayed; the MoE reached dense-final quality roughly 500 steps early. Final val accuracy 44.37% vs 43.63%.
For scale: nanochat reaches GPT-2 grade at 0.718 bits/byte using ~28x more compute.
Files
model.safetensors, config.json, tokenizer. config.json records
moe_experts_per_tok, which cannot be recovered from tensor shapes -- without
it the weights load correctly but route with the wrong top-k.
Usage
import json
from safetensors.torch import load_file
from ohara.models.llama import Config, Llama
cfg = json.load(open("config.json"))
cfg.pop("architecture"); cfg.pop("iteration")
model = Llama(Config(**cfg))
model.load_state_dict(load_file("model.safetensors"), strict=False) # rotary rebuilds
Finetune it into a chat model:
python examples/train_sft.py --pretrained-checkpoint moe_d12.pt --moe-experts-per-tok 2
Reproducing
DEPTH=12 bash runs/speedrun.sh # then add --moe-num-experts 8 --moe-experts-per-tok 2
- Downloads last month
- 56