File size: 3,782 Bytes
8443a50 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | ---
license: mit
language: en
tags:
- gpt
- from-scratch
- tinystories
- small-language-model
- pytorch
datasets:
- roneneldan/TinyStories
pipeline_tag: text-generation
---
# TinyStories SLM
A ~12.4M-parameter decoder-only transformer, **written from scratch** — no
`transformers`, no `tokenizers`, no `lightning`. Tokenizer, model, training
loop, and sampler are all hand-written PyTorch. Full source, training logs,
and design writeup: **[github.com/sur950/tinystories-slm](https://github.com/sur950/tinystories-slm)**.
It trains on [TinyStories](https://arxiv.org/abs/2305.07759) — Ronen Eldan
and Yuanzhi Li's dataset of GPT-4-generated children's stories — and asks how
small a language model can be and still write coherent English.
## Not a `transformers` model
This architecture (RMSNorm, tied embeddings, hand-rolled KV cache) isn't a
`transformers`-compatible class, so `AutoModel.from_pretrained(...)` won't
work here. To actually run it:
```bash
git clone https://github.com/sur950/tinystories-slm.git
cd tinystories-slm
./setup.sh
```
Download `stage2.pt`/`stage3.pt`/`tokenizer.json` from this repo into
`./checkpoints` and `./data` respectively, then:
```bash
python chat.py --stage 2 # base model: plain text continuation, no chat format
python chat.py --stage 3 # instruction-tuned: greetings, "tell me a story about X",
# and a redirect for anything outside TinyStories' domain
```
## Two checkpoints
| file | what it is |
|---|---|
| `stage2.pt` | Base language model. Trained in two passes over the ~560M-token corpus (Chinchilla-sized for this model). No chat format — continues whatever text precedes it, same as the original TinyStories models. |
| `stage3.pt` | `stage2.pt` fine-tuned on a small (~4k example) instruction set — see `instruct_raw.jsonl` in this repo for the exact data. Responds to greetings, takes "tell me a story about X" requests, and gives a fixed redirect for anything outside the TinyStories domain instead of hallucinating an answer. |
`stage3.pt` is the one to use for anything chat-like — **not** the
lower-val-loss checkpoint from the same run (not included here). Val loss on
this fine-tune is a single number blended across three very unevenly sized
categories, so "lowest aggregate val" mostly tracks the largest category
(story requests) and says little about whether the smaller, demo-critical
categories (greetings, redirects) are reliable. Tested directly: the final
checkpoint answered "Hi" correctly on-template ~4/5 times; the lowest-val
checkpoint from partway through the same run only ~2/5. Full writeup in the
GitHub repo's `DEVELOPMENT.md`.
## Architecture
```
vocab 4096 d_model 384 layers 6 heads 6 context 512 tied embeddings
12,391,296 parameters total
```
Byte-level BPE tokenizer, vocab size 4096, fit on this same corpus — not the
standard GPT-2 vocabulary, so token IDs from other tokenizers won't map
correctly onto this model.
## Limitations
- English-only, children's-story register — no world knowledge, no code, no
reasoning beyond what a simple short story requires.
- `stage3.pt`'s instruction-following is narrow by design: greetings, story
requests within a small set of trained keywords, and off-domain redirects.
It is not a general chatbot.
- 512-token context window.
## Citation
```bibtex
@misc{eldan2023tinystories,
title = {TinyStories: How Small Can Language Models Be and Still
Speak Coherent English?},
author = {Ronen Eldan and Yuanzhi Li},
year = {2023},
eprint = {2305.07759},
archivePrefix = {arXiv},
}
```
Built by [Suresh](https://github.com/sur950). [MIT licensed](https://github.com/sur950/tinystories-slm/blob/main/LICENSE).
|