moe-500m-fineweb-edu (Zhouaz-MoE)
A custom mixture-of-experts decoder pretrained from scratch (random init) on
HuggingFaceFW/fineweb-edu
(sample-10BT, streaming), ~65M tokens in one ~82-minute run on a single NVIDIA A10G.
This model requires trust_remote_code=True — the architecture source ships in
modeling_zhouaz_moe.py.
Architecture (chosen by CPU benchmark A/B against 4 candidates)
- 16 layers, dim 1024, GQA attention (16 query heads : 4 KV heads), RoPE, RMSNorm (pre-norm), no biases
- FFN = MoE: 8 experts, top-2 routing, SwiGLU expert hidden 1024
- Switch-style load-balancing auxiliary loss, coef 0.01
- Tied input/output embeddings, vocab 50257 (GPT-2 BPE), context 1024
- Total params 496,224,256 (~496M) · active params per token ~194M
- Benchmark motivation: fastest per-token decode of all candidates (38.7ms vs 75.4ms dense on CPU), ~2x faster than an equally-sized dense model, while carrying the most total (capacity) parameters.
Training
- CLM loss + aux loss, bf16, packed 1024-token blocks
- AdamW β=(0.9, 0.95), weight decay 0.1, grad clip 1.0
- lr 3e-4, cosine, 200-step warmup, effective batch 32×1024 = 32,768 tokens/step
- 2,000 optimizer steps ≈ 65.5M tokens
Results
| Metric | Value |
|---|---|
| final eval loss | 4.7106 |
| eval perplexity | ≈ 111 (exp of eval loss) |
| final train loss (real) | 4.66 |
Heavily undertrained (65M tokens for ~194M active params — Chinchilla would want ~4B+). It demonstrates the architecture and pipeline, not quality: expect repetitive, low-coherence text.
Usage
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("zhouaz/moe-500m-fineweb-edu")
model = AutoModelForCausalLM.from_pretrained(
"zhouaz/moe-500m-fineweb-edu", trust_remote_code=True).eval()
ids = tok("The history of the Roman Empire begins", return_tensors="pt").input_ids
for _ in range(60):
with torch.no_grad():
logits = model(ids).logits
ids = torch.cat([ids, torch.multinomial(torch.softmax(logits[0, -1] / 0.8, -1), 1)[None]], 1)
print(tok.decode(ids[0], skip_special_tokens=True))
Note on logged loss
The Trainer's loss field in logs is the pre-gradient-accumulation value (accum=4);
divide by 4. Eval loss is unaffected.