TinyStories-24m: 24.59M BPE GPT trained from scratch on TinyStories (val ppl 8.76, coherent)

README.md

@@ -0,0 +1,88 @@
+---
+license: apache-2.0
+pipeline_tag: text-generation
+language: en
+tags:
+  - tiny
+  - tiny-lm
+  - tiny-model
+  - slm
+  - small-language-model
+  - from-scratch
+  - tinystories
+  - bpe
+  - gpt
+datasets:
+  - roneneldan/TinyStories
+metrics:
+  - perplexity
+---
+
+# TinyStories-24m
+
+A **24.59M-parameter** BPE language model trained **from scratch** on
+[roneneldan/TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories),
+producing coherent short stories with proper dialogue, names, punctuation and
+narrative flow.
+
+## What it is
+
+- **Architecture:** decoder-only GPT, weight-tied embeddings, RMSNorm, fused
+  qkv multi-head causal attention (SDPA), GELU FFN.
+- **Shape:** D=384, L=12 layers, H=8 heads, FFN=1536, SEQ=512, vocab=8192 (BPE).
+- **Params:** 24,585,600 (verified against the safetensors header).
+- **Data:** roneneldan/TinyStories — 447.8M train tokens, 2M held-out val.
+- **Training:** 1 epoch ≈ 13,600 steps, AdamW, cosine LR 6e-4 + 500 warmup,
+  bf16 autocast, on a single RTX 5090.
+
+## Quality
+
+- **Val perplexity:** 8.76 (2.1618 nats/token on the 2M held-out val set).
+- **Generation:** coherent. Sampled 9/9 seeded generations (3 seeds × 3 prompts)
+  produce proper dialogue, character names (Ben, Lily, Mom, Tom, Sarah, Max),
+  punctuation and narrative flow. This model is a story generator for its
+  training domain — it is **not** a general-purpose assistant and will not
+  answer questions it was not trained on.
+
+## Honest caveats
+
+- **Divergence:** the full 13,600-step run diverged to NaN at step 9,350 (LR 6e-4
+  is too hot for a 24M model). The **best** checkpoint (step 6,000, val 2.1618)
+  is what is published here — it is clean and coherent. The divergence is late,
+  so a clean early checkpoint is the right artifact; always sample the best
+  checkpoint, not the final one.
+- **Domain-bound:** trained only on TinyStories. Out-of-domain text (code,
+  questions, general conversation) is out of scope.
+
+## Usage
+
+Not a `transformers` model — load with the bundled `modeling.py`:
+
+```python
+import sys, torch
+sys.path.insert(0, "path/to/this/repo")
+from modeling import TinyStoriesGPT
+from tokenizers import Tokenizer
+
+m = TinyStoriesGPT.from_pretrained("path/to/this/repo", device="cpu")
+tok = Tokenizer.from_file("path/to/this/repo/tokenizer.json")
+
+ids = tok.encode("Ben was playing in the park.", add_special_tokens=False).ids
+x = torch.tensor([ids], dtype=torch.long)
+with torch.no_grad():
+    for _ in range(80):
+        logits = m(x[:, -512:])[:, -1]
+        nxt = torch.multinomial(torch.softmax(logits / 0.8, -1), 1).item()
+        ids.append(nxt)
+        x = torch.tensor([ids[-512:]], dtype=torch.long)
+print(tok.decode(ids, skip_special_tokens=True))
+```
+
+## Files
+
+| file | bytes | what |
+|------|-------|------|
+| `model.safetensors` | 98,349,056 | 75 tensors, float32 |
+| `config.json` | — | architecture + training metadata |
+| `modeling.py` | — | the `TinyStoriesGPT` class (load with `from_pretrained`) |
+| `tokenizer.json` | 560,804 | BPE-8k tokenizer (HF `tokenizers` format) |

config.json

@@ -0,0 +1,20 @@
+{
+  "architectures": ["TinyStoriesGPT"],
+  "model_type": "tinystories-gpt",
+  "D": 384,
+  "L": 12,
+  "H": 8,
+  "FFN": 1536,
+  "vocab_size": 8192,
+  "max_position_embeddings": 512,
+  "n_params": 24585600,
+  "tie_word_embeddings": true,
+  "norm": "RMSNorm",
+  "ffn": "GELU",
+  "attn": "causal SDPA (fused qkv)",
+  "trained_on": "roneneldan/TinyStories",
+  "step": 6000,
+  "val_nats": 2.1618,
+  "val_ppl": 8.76,
+  "dtype": "float32"
+}

modeling.py

@@ -0,0 +1,66 @@
+"""TinyStoriesGPT — 24.59M-param BPE GPT trained on roneneldan/TinyStories.
+Architecture: weight-tied decoder-only GPT, RMSNorm, fused qkv, GELU FFN.
+Not a transformers model — load with this class + safetensors.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class RMSNorm(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.w = nn.Parameter(torch.ones(d))
+    def forward(self, x):
+        return self.w * x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + 1e-6)
+
+class Block(nn.Module):
+    def __init__(self, d, h, ffn):
+        super().__init__()
+        self.ln1 = RMSNorm(d)
+        self.ln2 = RMSNorm(d)
+        self.qkv = nn.Linear(d, 3*d, bias=False)
+        self.proj = nn.Linear(d, d, bias=False)
+        self.fc1 = nn.Linear(d, ffn, bias=False)
+        self.fc2 = nn.Linear(ffn, d, bias=False)
+        self.h, self.d = h, d
+    def forward(self, x):
+        B, T, D = x.shape
+        h = self.ln1(x)
+        qkv = self.qkv(h).view(B, T, 3, self.h, D//self.h).transpose(2,1)
+        q, k, v = qkv[:,0], qkv[:,1], qkv[:,2]
+        q, k, v = q.transpose(1,2), k.transpose(1,2), v.transpose(1,2)
+        att = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        att = att.transpose(1,2).reshape(B, T, D)
+        x = x + self.proj(att)
+        x = x + self.fc2(F.gelu(self.fc1(self.ln2(x))))
+        return x
+
+class TinyStoriesGPT(nn.Module):
+    def __init__(self, vocab_size=8192, d=384, n_layers=12, n_heads=8, ffn=1536, seq=512):
+        super().__init__()
+        self.tok = nn.Embedding(vocab_size, d)
+        self.pos = nn.Embedding(seq, d)
+        self.blocks = nn.ModuleList([Block(d, n_heads, ffn) for _ in range(n_layers)])
+        self.ln_f = RMSNorm(d)
+        self.vocab_size = vocab_size
+    def forward(self, x, targets=None):
+        b, t = x.shape
+        h = self.tok(x) + self.pos(torch.arange(t, device=x.device))
+        for blk in self.blocks:
+            h = blk(h)
+        h = self.ln_f(h)
+        logits = h @ self.tok.weight.t()
+        if targets is not None:
+            return F.cross_entropy(logits.float().view(-1, self.vocab_size), targets.view(-1))
+        return logits
+    @classmethod
+    def from_pretrained(cls, path, device="cpu"):
+        import json
+        from safetensors.torch import load_file
+        cfg = json.load(open(f"{path}/config.json"))
+        model = cls(vocab_size=cfg["vocab_size"], d=cfg["D"], n_layers=cfg["L"],
+                    n_heads=cfg["H"], ffn=cfg["FFN"], seq=cfg["max_position_embeddings"])
+        sd = load_file(f"{path}/model.safetensors")
+        model.load_state_dict(sd)
+        model = model.to(device).eval()
+        return model