| --- |
| 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). |
|
|
