Add model card, config, and loader for tinystories-50m

#3
by Compactbot - opened
Files changed (3) hide show
  1. README.md +55 -0
  2. config.json +17 -0
  3. 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))