MiraLM-47M

Why

A hard 50M cap stops being a hyperparameter and becomes the design. I wanted to know what a model still learns when every parameter has to be defended.

And I was tired of "from scratch" meaning "fine-tuned from a checkpoint nobody mentions."

The machine

14 layers, attention and Mamba alternating Jamba-style — GQA (8 query heads over 4 KV, RoPE), SwiGLU, pre-norm RMSNorm. The Mamba blocks are pure PyTorch: no custom CUDA, a Hillis–Steele parallel scan in O(log T) sequential depth rather than a loop over T, under gradient checkpointing — because autograd otherwise retains O(log T) fp32 intermediates per block, and 47M does not fit a 16 GB T4.

Then the part I actually care about: a top-2-of-8 MoE with a semantically seeded router. Every corpus document carries a domain tag. For the first 2,000 steps a guide loss pins each token to its domain expert, then anneals to zero. Total loss is cross-entropy plus three router terms — guide (domain curriculum), z-loss (logit containment, ST-MoE style), aux (Switch-style load balance). I didn't want experts that merely could differ. I wanted a reason for them to.

The budget is a gate, not a hope

50,000,000 is a disqualifier, not a target. So the count is checked twice, and the two must agree: a static projection from the YAML, and a real numel() on the constructed module. The test suite fails if they ever drift.

That makes weight tying non-optional:

47,640,968 + 24,000 × 384 = 56,856,968

— 6.9M over the line untied. Tied, it lands at 47,640,968, with 2.36M of headroom, and only 38,793,608 (81.4%) fire per token.

Training — and the honest part

One T4. fp16 + GradScaler, batch 4 × 1024 = 4,096 tokens/step. Cosine, AdamW. 11,000 steps = 45,056,000 tokens.

Chinchilla says this size wants 20 × 47,640,968 ≈ 9.5e8 tokens. I trained on 4.7% of that. So I won't pretend:

• HellaSwag — 33.0
• ARC-Easy — 27.0
• PIQA — 50.0
• WinoGrande — 50.2
• WikiText-103 PPL — 2837.6

Commonsense multiple-choice sits at its random floor. That is the honest reading, and it is above the fold in the README.

The divergence. At step 17,000 the run diverged — loss climbing to 2.587. I did not quietly ship the newer file. I shipped the step-11,000 best (loss 2.2113, ppl 9.13), and said so.

Two things I'd defend anywhere

Reading the rules. Metric five asks for perplexity on a held-out slice of WikiText-103. lm-eval's stock wikitext task scores wikitext-2 — a different corpus, held out by nobody. So I measured what was actually asked: a deterministic, committed slice of the WikiText-103 test split the training pipeline never reads.

Measuring the router instead of claiming it. Eight experts, zero dead, and routing entropy against its ceiling:

$$H = -\sum p \log p = 2.0783, \qquad \ln 8 = 2.0794$$

— 99.95% of maximum. Top-1 load spread collapsed from 0.268 to 0.024.

What I learned

Constraint is the only design tool that doesn't lie. And at d_model = 384 the model is memory-bandwidth-bound, not FLOP-bound: 6ND ≈ 1.05e16 FLOPs is well under a GPU-hour of pure matmul, yet measured wall-clock is many times that. The cost is elementwise scan ops and kernel-launch overhead. That is the next win, and the honest answer to "training efficiency."

Built With

• accelerate
• bytelevel
• grouped-query
• hub
• hugging
• kaggle
• lm-evaluation-harness
• mamba
• next.js
• notebooks
• numpy
• nvidia
• pytest
• pytorch
• pyyaml
• rmsnorm
• safetensors
• swiglu
• tailwind
• tokenizers
• typescript