Mieaz/gpt22m-chat

🤖 GPT-22M — trained from scratch, fully transparent

A 23.1M-parameter GPT trained from scratch on 600M tokens (2 epochs) on free cloud GPU (Tesla P100, 16GB). Built to be small, honest, and completely documented — every number below is verified from the training logs.

Try it in your browser: GPT-22M Chat Space (runs 100% locally via Transformers.js — nothing leaves your browser).

---

🧮 Model Card — the short version

Property	Value
Parameters	23.1M (21.5M trainable + tied-embedding accounting)
Architecture	24 layers × 256 hidden · GQA (8 heads / 4 KV) · ffn 1024 · tied embeddings
Context	256 tokens (seq_len 256)
Vocabulary	16,000 (ByteLevel BPE, trained on stories + chat mixed)
Tokenizer quality	2.40 tokens/word (vs 3.62 for a naive 4095-vocab; GPT-2's 50K ≈ 1.47)
Training data	600M tokens = 2 epochs over 300M (200M TinyStories + 100M ultrachat dialogues)
Schedule	WSD (warmup → stable 5e-4 → final 10% anneals into chat at 50/50 mix)
Final held-out val loss	1.52 (TinyStories-valid; 18 checkpoints, monotone 2.20 → 1.52)
GPU time	~6 hours on Tesla P100 (fp16 AMP, batch 64 × seq 256)

---

📊 Training trace (audited, all real)

Loss curve (training loss, sampled every ~3,000 steps):

step   200  loss 4.39   (init ≈ ln(16000) = 9.68)
step  3200  loss 2.24
step  9200  loss 1.98
step 15200  loss 1.88
step 21200  loss 1.88
step 27200  loss 1.83
step 33200  loss 2.55   ← annealing phase: LR decay + chat upweighting
step 36600  loss 2.18

Held-out validation loss (TinyStories-valid), the generalization proof — 18 checkpoints, monotone decrease:

2.20 → 1.91 → 1.80 → 1.74 → 1.70 → 1.68 → 1.66 → 1.64 → 1.63
→ 1.62 → 1.61 → 1.60 → 1.60 → 1.59 → 1.59 → 1.58 → 1.58 → 1.53

Final val loss 1.52 (independently re-measured at benchmark: 1.52). Zero OOM/errors; GPU consumption verified at exactly 1:1 with wall clock throughout.

---

💬 What it actually generates (real, from the benchmark kernel)

Prompt	Output (first ~70 chars)
User: hi	"Yes, the key to a memorable family holiday destination is the opportunity to explore the city's history and culture..."
what is your name?	"I don't have access to real-time information about the specific services..."
Once upon a time there was a little fox	"Visit your friend's home and have a surprising experience together..."
what is 7 times 8?	"1. Measure the length of your..." (wrong answer — see Limitations)

What improved vs the 15M predecessor: val loss 1.75 → 1.52; output went from invented words ("equesterngerationalusearchieving") to real English, grammatical sentences, appropriate response forms.

---

⚖️ Honest limitations

• 21.5–23M params is tiny. It produces recognizable, grammatical language and correct response forms, but not fluent conversation or reasoning.
• Arithmetic is wrong (7×8 ≠ 56). Small LMs memorize patterns, they don't compute.
• Repetition occurs under sampling; lowering temperature / adding repetition penalty helps.
• Domain: strongest on simple narrative (TinyStories-style); general chat is the hard mode at this size.
• This is a research-scale toy, not a production chatbot. Use it to study small-model training, scaling, and deployment — not as a product.

---

🔬 How it was trained (reproduce it)

1. Data: roneneldan/TinyStories (train + held-out valid) + HuggingFaceH4/ultrachat_200k (train_sft). Chat tokenized per-dialogue (real \n between turns, single <|endoftext|>) so turn adjacency is preserved.
2. Tokenizer: ByteLevelBPE, vocab 16,000, trained on stories+chat mixed (2.40 tok/word).
3. Arch: 24L×256, GQA-4, tied embeddings, seq 256 (the architecture lessons from MobileLLM: deep-and-thin + GQA + weight-sharing beat shallow-wide at the same size).
4. Training: fp16 AMP, AdamW 5e-4 (β=0.9/0.95, wd 0.1), warmup 200, grad clip 1.0, batch 64 × seq 256. WSD schedule: stable 5e-4 for 90% of steps, then linear decay into a 50/50 chat mixture in the final 10% (the SmolLM2 annealing method: broad → specialized).
5. 2 epochs / 600M tokens (supported by Muennighoff et al.: up to 4 epochs of repeats ≈ free).
6. Hardware: free Tesla P100 (sm_60), torch 2.4.1+cu118, ~6h.

---

🚀 Deployment (browser inference)

• ONNX: hand-rolled with-past KV-cache export (49 inputs / 49 outputs: input_ids, attention_mask, past_key_values.{i}.key/value for 24 layers). Validated in onnxruntime for every (seq, past) shape including the generation pattern (seq=1, past=7).
• Space: static Hugging Face Space running Transformers.js v4.2.0 (WebGPU, WASM fallback). Model, tokenizer, and ONNX all ship with the repo; inference is fully client-side.

---

📁 Files

File	Purpose
model.pt	PyTorch checkpoint (step 36,614, val 1.52)
onnx/model.onnx	With-past ONNX for browser inference (86 MB)
config.json	GPT2-style config (use_cache: true)
tokenizer.json, vocab.json, merges.txt	ByteLevel BPE (16,000 vocab)
tokenizer_config.json	add_bos/eos: false
training_meta.json	Step, val loss, schedule, data
README.md	This card

---

🧭 Why this exists

This model is the output of a deliberate, documented study of small language models: what they can and cannot do, how to train them well on a budget, and how to deploy them to real users for free. The full methodology (13+ battle-tested lessons, with citations) lives in the companion playbook. Questions, reproductions, and improvements are welcome.