🐿️ Tiny-Storyteller-MoE-59M-A19M

An 8-expert dropless Mixture-of-Experts language model, trained from scratch in 20 minutes.

Every layer's feed-forward network is replaced by 8 SwiGLU experts with top-2 routing, run through MegaBlocks' dropless block-sparse dMoE kernels — so no token is ever discarded, however skewed the router gets. The tokenizer is a byte-level BPE trained from raw bytes; the transformer is plain PyTorch. Nothing is imported from a model library.

Trained on TinyStories-Instruct for 2,000 steps — 20.1 minutes on a single H100.

Results

⚠️ CE and perplexity above are training-set numbers. No held-out split was logged for this run, so treat them as optimization progress, not generalization.

Quickstart

The prompt format is not optional. This is a base model trained on exactly one document shape. You must close the instruction and open the story, or the model will finish writing the instruction instead of the story.

from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "Omkarsatapathy/Tiny-Storyteller-MoE-59M-A19M"
model = AutoModelForCausalLM.from_pretrained(MODEL, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(MODEL)

instruction = (
    "Write a story that includes the words spring, pitch, flexible, "
    "features moralvalue, and follows this summary: A squirrel wants to play "
    "with a ball on the pitch but is too small, so a flexible rabbit helps him."
)
prompt = f"<|startofinstruction|>{instruction}<|endofinstruction|><|startofstory|>"

inputs = tokenizer(prompt, return_tensors="pt")
output = model.generate(**inputs, max_new_tokens=300, do_sample=True,
                        temperature=0.8, top_k=50)

story = tokenizer.decode(output[0, inputs.input_ids.shape[1]:], skip_special_tokens=True)
print(story)

trust_remote_code=True is required — this is a custom architecture that ships its own modeling code. The repo ships a pure-PyTorch MoE implementation, so it runs on CPU, MPS, or CUDA — MegaBlocks is not needed for inference.

• Weights are mixed precision, as trained: attention, embeddings and norms in float32 (5.9M params), MoE expert and router weights in bfloat16 (53.1M) — MegaBlocks kept them that way. Loading is exact; pass torch_dtype=torch.float32 to promote everything.
• Batched generation works. Set tokenizer.padding_side = "left", as with any decoder-only model.

<|startofinstruction|>INSTRUCTION<|endofinstruction|><|startofstory|>STORY<|endofstory|>

Every INSTRUCTION came from one template:

Write a story that includes the words {words}, features {features}, and follows this summary: {summary}

Straying from that template degrades output as much as dropping the markers does. The model emits <|endofstory|> on its own and generation stops there.

Architecture

59.0M total parameters, 19.2M active per token. The 8 experts hold 53.1M of that; the top-2 router touches 2 of them per token, so capacity scales with total params while compute scales with active ones (5.9M shared + 13.3M of experts).

Why dropless matters: classic MoE gives each expert a fixed capacity and throws away overflow tokens. MegaBlocks expresses the whole layer as one block-sparse matmul instead, so no token is dropped no matter how imbalanced routing becomes.

AdamW (β=0.9/0.95, weight decay 0.1 on 2D params only), fused; cosine LR 8.5e-4 → 8.5e-5 after 200 warmup steps; grad clip 1.0; bf16 autocast; effective batch ~524K tokens/step across 2,000 steps. Load-balancing loss weight 0.01, pulled from MegaBlocks' global state via batched_load_balancing_loss() — it arrives pre-scaled and is added to CE directly.

Training dynamics

Cross-entropy falls 7.8× — most of it inside the first 500 steps.

The auxiliary load-balancing loss does its job: all 8 experts converge onto an even split and stay there, ending within ±0.004 of the ideal 0.125.

Throughput climbs through warmup and settles at a median of 1.0M tokens/sec.

Sample outputs

Verbatim generations, temperature=0.8, top_k=50. Both stopped on their own at <|endofstory|>.

Instruction: include the words spring, pitch, flexible, features moralvalue, summary: A squirrel wants to play with a ball on the pitch but is too small, so a flexible rabbit helps him.

The squirrel saw a friend and asked, "Do you want to play with me?" The rabbit said, "Yes, I want to play. But I am too small." The rabbit thought for a moment and said, "I can help you. You can spring and jump to the ball." The rabbit showed the squirrel how to play with the ball. They had fun and played together. The squirrel was happy and learned that it is good to help friends. The squirrel learned that it is good to help friends.

The squirrel tried to play with the ball, but it was too big. He tried to jump and spring, but he could not reach the ball. The squirrel felt sad. He wanted to play with the ball, but he was too small. Then, a flexible rabbit came to help. The rabbit said, "I can help you, but you must be careful." The rabbit was happy and said, "Thank you, rabbit!" The rabbit and the rabbit played with the ball. They had fun and were happy. The moral of the story is to be careful and help others, even if you are small.

Limitations

• Metrics are training-set, not held-out. ~1.9 epochs over the corpus, so some memorization is plausible and the reported PPL likely flatters generalization.
• Repetition loops. Sentences and whole clauses repeat ("The squirrel learned that it is good to help friends" twice above); greedy decoding is worse, and the model sometimes confuses its own characters ("The rabbit and the rabbit").
• Coherence decays past ~3–5 sentences.
• Instruction-following is formal, not semantic. It matches the template's shape and usually works the required words in, but ignores length constraints and often misses the summary's actual plot.
• TinyStories vocabulary only. No facts, no reasoning — off-domain prompts produce confident nonsense. photosynthesis alone costs 7 tokens.
• No KV cache, so generation recomputes the whole prefix each step and is slow for long outputs.
• No instruction tuning, no alignment, no safety filtering. Research artifact for studying MoE training dynamics at small scale — not a product.
• The 300K→1.0M tok/s comparison against the earlier hand-written MoE is directional, not a controlled benchmark.

Citation

@misc{satapathy2026tinymoe,
  title  = {Dropless Mixture-of-Experts in a Tiny Language Model},
  author = {Satapathy, Omkar},
  year   = {2026},
  url    = {https://huggingface.co/Omkarsatapathy/Tiny-Storyteller-MoE-59M-A19M}
}

📄 Technical report: GPT-MoE_MegaBlocks_Technical_Report.pdf — full derivation, training run, and results.

🧑‍💻 Code: github.com/omkarsatapathy/Building_Language_models — under experiments/code/MoE_MegaBlocks. Built from scratch, tokenizer included.

📚 Built on: MegaBlocks · TinyStories · RoPE · SwiGLU

License

MIT.