ohara-moe-chat-d12

ohara-moe-chat-d12

A 332M-parameter mixture-of-experts chat model trained from scratch with ohara, following the nanochat pipeline: pretrain on ClimbMix, then supervised finetuning on conversations.

Base model: joey00072/ohara-moe-base-d12. Dense counterpart: joey00072/ohara-chat-d12.

Calibrate before building on this. 85M active parameters at 1.4e18 training FLOPs is roughly 3.5% of what nanochat spends to reach GPT-2 grade. It answers in the right shape, holds the conversation format, and stops cleanly. It also confabulates freely. This demonstrates a working pipeline, not a useful assistant.

MoE architecture

8 experts per layer, top-2 routed, on all 12 layers. No shared expert (so not DeepSeek-style). Each expert is a SwiGLU of width 1024 -- exactly the dense FFN width divided by top-k, so two experts firing cost the same FLOPs as one dense feed-forward.

Load balancing is quantile balancing (Jianlin Su; used in Kimi K2/K3): the router bias is solved in closed form from batch statistics each optimizer step. No auxiliary loss and no balancing coefficient to tune. All 12 routers converged to a balanced state (MaxVio ~0.03).

MoE vs dense, controlled

Identical FLOPs per token, token budget, data, schedule, learning rates and seed. The only change was the dense feed-forward becoming 8 experts at half width.

The MoE held a ~1.9% bits-per-byte advantage at every evaluation and reached the dense model final quality about 500 steps early.

On wall clock, be careful: this run took 2.35h against the dense run 3.92h, but that speedup came from pre-tokenized data and torch.compile, not from MoE. At matched FLOPs, MoE routing costs about 2.1x the step time of dense. What MoE bought is the quality delta, not the speed.

Files

model.safetensors, config.json, and the 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, torch
from safetensors.torch import load_file
from ohara.models.llama import Config, Llama
from ohara.chat_engine import ChatEngine, SamplingConfig

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 buffers rebuild

engine = ChatEngine(model, tokenizer)   # tokenizer from this repo
print(engine.generate([{"role": "user", "content": "Why is the sky blue?"}],
                      SamplingConfig(temperature=0.7, max_new_tokens=100)))

Training

Pretraining 2,827 steps, 1.48B ClimbMix tokens, batch 524,288, Muon lr 0.02, warmup-stable-decay, bf16. SFT 800 steps on SmolTalk + MMLU + GSM8K (567,656 conversations), loss on assistant tokens only, conversations packed best-fit so none is split across rows.

Reproduce: DEPTH=12 bash runs/speedrun.sh (dense) or add --moe-num-experts 8 --moe-experts-per-tok 2.