ohara-moe-base-d12

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.

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.

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