tiny-lm-stories-12m

A 12M-parameter GPT-2 shaped model that writes short children's stories. It is from Part 2 of tiny-lm, a project that builds a language model step by step and then looks inside it.

Size	6 layers, width 384, 6 heads, 12.3M parameters
Context	256 tokens
Tokeniser	byte-level BPE, 4,096 tokens
Training	one pass over TinyStories (553M tokens) on a rented GPU
Validation loss	1.28 nats per token

There are two models in this repository:

• The root is the base model. It continues a story.
• instruct/ is the base model after fine-tuning on TinyStoriesInstruct. It writes a story that uses three given words.

The weights are a standard GPT2LMHeadModel. The q, k and v biases that GPT-2 has are zero, because the model was trained without them. Start the text after <|endoftext|>, as in training.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "lillycham/tiny-lm-stories-12m"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo)

ids = tok("Once upon a time, a little fox", return_tensors="pt").input_ids
ids = torch.cat([torch.tensor([[tok.eos_token_id]]), ids], dim=1)
out = model.generate(ids, max_new_tokens=100, do_sample=True, temperature=0.8)
print(tok.decode(out[0, 1:]))

For the instruct model, load it with subfolder="instruct" and use this prompt:

User: Tell me a story that uses the words dragon, soup and happy.
Assistant:

What it is for

The model is small and meant for study. It has an induction head (layer 5, head 3) that copies names from its context. The tiny-lm repository has tools to look at it. Do not expect facts from it.

Data and licence

Trained on TinyStories and TinyStoriesInstruct (Eldan and Li, 2023, CDLA-Sharing-1.0). The code and weights are MIT.