Add model card, config, and loader for tinystories-50m
#3
by Compactbot - opened
- README.md +55 -0
- config.json +17 -0
- model_loader.py +62 -0
README.md
ADDED
|
@@ -0,0 +1,55 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: apache-2.0
|
| 3 |
+
pipeline_tag: text-generation
|
| 4 |
+
language: en
|
| 5 |
+
tags:
|
| 6 |
+
- tiny
|
| 7 |
+
- tiny-lm
|
| 8 |
+
- tiny-model
|
| 9 |
+
- slm
|
| 10 |
+
- SLM
|
| 11 |
+
- small-language-model
|
| 12 |
+
- from-scratch
|
| 13 |
+
- tinystories
|
| 14 |
+
metrics:
|
| 15 |
+
- perplexity
|
| 16 |
+
- val-loss
|
| 17 |
+
---
|
| 18 |
+
|
| 19 |
+
# tinystories-50m
|
| 20 |
+
|
| 21 |
+
A **54,804,992-parameter** from-scratch BPE GPT trained on [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories). This is the scale-up of my [tinystories-24m](https://huggingface.co/Compactbot/tinystories-24m) — same narrow children's-story domain, ~2.3x the parameters.
|
| 22 |
+
|
| 23 |
+
## What it is
|
| 24 |
+
- **Architecture:** custom GPT — 16 layers, d=512, 8 heads, FFN 2048, seq 512, vocab 8192 (BPE). Weight-tied embeddings, RMSNorm, GELU, fused qkv. **Not** a stock transformers architecture — load it with the included `model_loader.py`.
|
| 25 |
+
- **Data:** 447.9M TinyStories tokens, 1 full epoch (13,668 steps, batch 64, seq 512).
|
| 26 |
+
- **Compute:** RTX 5090 (32 GB), ~8 steps/s, ~28 min.
|
| 27 |
+
- **Params:** **54,804,992** exactly (verified against the safetensors header: 99 tensors, weight-tied so `tok.weight` is not double-counted).
|
| 28 |
+
|
| 29 |
+
## Training
|
| 30 |
+
- LR 3e-4, cosine decay, warmup 500. No divergence (unlike the 24m at 6e-4).
|
| 31 |
+
- Best val loss **1.6566** @ step 13,250; final **1.6371** @ step 13,668.
|
| 32 |
+
- In-domain perplexity (held-out TinyStories): **~5.24** (best ckpt).
|
| 33 |
+
- ~8.2 tokens/param.
|
| 34 |
+
|
| 35 |
+
## Quality (honest)
|
| 36 |
+
Coherent, on-theme children's stories with proper punctuation and narrative structure. Sample (seed 0):
|
| 37 |
+
|
| 38 |
+
> Once upon a time, there was a little girl named Lily. She loved to play in the garden with her dog, Buddy. One day, Lily's mom taught her a new game called wash. Lily was very excited and asked her mom if she could bring her favorite toy back. Lily went outside and filled a bucket with water. As she was catching the bucket, she noticed that it was wet and her clothes were wet. She quickly dried herself off and her clothes were all wet. She ran inside and asked her mom for help.
|
| 39 |
+
|
| 40 |
+
Occasional artifacts are expected at this scale on a narrow domain: a garbled word here and there, pronoun slips, and short-range repetition. It is a **story generator**, not a general model — do not expect general English, reasoning, or instruction following.
|
| 41 |
+
|
| 42 |
+
## What it is NOT good at
|
| 43 |
+
- General-domain text (out of distribution; expect degradation).
|
| 44 |
+
- Reasoning, math, multi-turn dialogue, instruction following.
|
| 45 |
+
- Long coherence beyond ~140 generated tokens (trained at seq 512 but quality degrades with length).
|
| 46 |
+
|
| 47 |
+
## Usage
|
| 48 |
+
```python
|
| 49 |
+
from model_loader import load, generate
|
| 50 |
+
model, tok = load(".") # needs torch, safetensors, tokenizers
|
| 51 |
+
print(generate(model, tok, "Once upon a time, there was a little girl named Lily.", seed=0))
|
| 52 |
+
```
|
| 53 |
+
|
| 54 |
+
## Lineage
|
| 55 |
+
- [tinystories-24m](https://huggingface.co/Compactbot/tinystories-24m) (24.59M, coherent at 18.2 tok/param) → **tinystories-50m** (54.80M, 8.2 tok/param).
|
config.json
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": ["CustomGPT"],
|
| 3 |
+
"model_type": "gpt",
|
| 4 |
+
"vocab_size": 8192,
|
| 5 |
+
"n_positions": 512,
|
| 6 |
+
"n_embd": 512,
|
| 7 |
+
"n_layer": 16,
|
| 8 |
+
"n_head": 8,
|
| 9 |
+
"n_inner": 2048,
|
| 10 |
+
"activation": "gelu",
|
| 11 |
+
"norm": "rmsnorm",
|
| 12 |
+
"tie_word_embeddings": true,
|
| 13 |
+
"attention": "causal-fused-qkv",
|
| 14 |
+
"n_params": 54804992,
|
| 15 |
+
"dtype": "float32",
|
| 16 |
+
"note": "Custom GPT (not a transformers architecture). Load with model_loader.py. Weights in model.safetensors (99 tensors, weight-tied: tok.weight serves both embedding and output projection)."
|
| 17 |
+
}
|
model_loader.py
ADDED
|
@@ -0,0 +1,62 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Load and generate with Compactbot/tinystories-50m (custom GPT, 54.8M params).
|
| 2 |
+
Requires: torch, safetensors, tokenizers. GPU optional (CPU works, slower).
|
| 3 |
+
"""
|
| 4 |
+
import torch, torch.nn as nn, torch.nn.functional as F
|
| 5 |
+
from safetensors.torch import load_file
|
| 6 |
+
from tokenizers import Tokenizer
|
| 7 |
+
|
| 8 |
+
VOCAB,D,L,H,FFN,SEQ=8192,512,16,8,2048,512
|
| 9 |
+
|
| 10 |
+
class RMSNorm(nn.Module):
|
| 11 |
+
def __init__(self,d): super().__init__(); self.w=nn.Parameter(torch.ones(d))
|
| 12 |
+
def forward(self,x): return self.w*x*torch.rsqrt(x.float().pow(2).mean(-1,keepdim=True)+1e-6)
|
| 13 |
+
|
| 14 |
+
class Block(nn.Module):
|
| 15 |
+
def __init__(self,d,h):
|
| 16 |
+
super().__init__()
|
| 17 |
+
self.ln1=RMSNorm(d); self.ln2=RMSNorm(d)
|
| 18 |
+
self.qkv=nn.Linear(d,3*d,bias=False); self.proj=nn.Linear(d,d,bias=False)
|
| 19 |
+
self.fc1=nn.Linear(d,FFN,bias=False); self.fc2=nn.Linear(FFN,d,bias=False)
|
| 20 |
+
self.h,self.d=h,d
|
| 21 |
+
def forward(self,x):
|
| 22 |
+
B,T,Dd=x.shape; h=self.ln1(x)
|
| 23 |
+
qkv=self.qkv(h).view(B,T,3,self.h,Dd//self.h).transpose(2,1)
|
| 24 |
+
q,k,v=qkv[:,0].transpose(1,2),qkv[:,1].transpose(1,2),qkv[:,2].transpose(1,2)
|
| 25 |
+
att=F.scaled_dot_product_attention(q,k,v,is_causal=True)
|
| 26 |
+
x=x+self.proj(att.transpose(1,2).reshape(B,T,Dd))
|
| 27 |
+
return x+self.fc2(F.gelu(self.fc1(self.ln2(x))))
|
| 28 |
+
|
| 29 |
+
class GPT(nn.Module):
|
| 30 |
+
def __init__(self):
|
| 31 |
+
super().__init__()
|
| 32 |
+
self.tok=nn.Embedding(VOCAB,D); self.pos=nn.Embedding(SEQ,D)
|
| 33 |
+
self.blocks=nn.ModuleList([Block(D,H) for _ in range(L)]); self.ln_f=RMSNorm(D)
|
| 34 |
+
def forward(self,x,targets=None):
|
| 35 |
+
h=self.tok(x)+self.pos(torch.arange(x.shape[1],device=x.device))
|
| 36 |
+
for b in self.blocks: h=b(h)
|
| 37 |
+
logits=self.ln_f(h)@self.tok.weight.t()
|
| 38 |
+
if targets is not None: return F.cross_entropy(logits.view(-1,VOCAB),targets.view(-1))
|
| 39 |
+
return logits
|
| 40 |
+
|
| 41 |
+
def load(repo_dir=".",device="cuda" if torch.cuda.is_available() else "cpu"):
|
| 42 |
+
m=GPT().to(device); m.load_state_dict(load_file(f"{repo_dir}/model.safetensors"),strict=True)
|
| 43 |
+
m.eval()
|
| 44 |
+
return m,Tokenizer.from_file(f"{repo_dir}/tokenizer.json")
|
| 45 |
+
|
| 46 |
+
def generate(model,tok,prompt,max_new=120,seed=0,temperature=0.8,device="cuda"):
|
| 47 |
+
g=torch.Generator(device=device).manual_seed(seed)
|
| 48 |
+
ids=torch.tensor([tok.encode(prompt).ids],device=device)
|
| 49 |
+
if ids.shape[1]>SEQ-4: ids=ids[:,-(SEQ-4):]
|
| 50 |
+
with torch.no_grad():
|
| 51 |
+
for _ in range(max_new):
|
| 52 |
+
with torch.autocast(device_type=device,dtype=torch.bfloat16):
|
| 53 |
+
logits=model(ids)
|
| 54 |
+
p=torch.softmax(logits[:,-1].float()/temperature,dim=-1)
|
| 55 |
+
nxt=torch.multinomial(p,1,generator=g); ids=torch.cat([ids,nxt],1)
|
| 56 |
+
if nxt.item()==1: break
|
| 57 |
+
return tok.decode(ids[0].tolist())
|
| 58 |
+
|
| 59 |
+
if __name__=="__main__":
|
| 60 |
+
m,t=load(".")
|
| 61 |
+
print("params:",sum(p.numel() for p in m.parameters()))
|
| 62 |
+
print(generate(m,t,"Once upon a time, there was a little girl named Lily.",seed=0))
|