d0mzw/pvml-tinystories

PvML TinyStories

Six small transformers trained from scratch on TinyStories, on one machine.

The architecture is a from-scratch GPT-2 reimplementation written while working through the ARENA 3.0 curriculum. Every module was checked against GPT-2's own weights before any training happened, so the code is known correct independently of these runs.

Code: https://github.com/d0mzw/PvML

Runs

Run	d_model	layers	heads	n_ctx	Steps	Loss	Accuracy	Minutes	Parameters
tinystories-d128-l6-h4-ctx512-40k	128	6	4	512	40,000	1.732	0.573	641	14,171,473
tinystories-d128-l6-h4-ctx128-20k	128	6	4	128	20,000	2.080	0.517	67	14,122,321
tinystories-d128-l6-h4-ctx128-5k	128	6	4	128	5,000	2.369	0.476	17	14,122,321
tinystories-d256-l8-h8-ctx128-5k	256	8	8	128	5,000	2.497	0.461	37	32,133,201
tinystories-d32-l4-h16-ctx128-20k	32	4	16	128	20,000	2.782	0.421	35	3,321,681
tinystories-d32-l4-h16-ctx128-5k	32	4	16	128	5,000	2.996	0.395	8	3,321,681

Loss is cross entropy in nats, averaged over the last 500 steps. Accuracy is next-token accuracy on 1000 held-out chunks. Uniform guessing over the 50,257 GPT-2 vocabulary is 10.825 nats.

Only the five at n_ctx 128 are comparable with each other. Predicting a token from 511 tokens of context is an easier problem than from 127, so part of the 512 run's lead is the task rather than the model.

Loading

Each folder holds model.safetensors, the config.json it was trained with, and a summary.json. A run is self-contained: the config rebuilds the architecture and the weights fill it.

import json
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from pvml.config import Config
from pvml.modules.transformer import Transformer

run = "tinystories-d128-l6-h4-ctx512-40k"
cfg = Config(**json.load(open(hf_hub_download("d0mzw/pvml-tinystories", f"{run}/config.json")))["model"])
model = Transformer(cfg)
model.load_state_dict(load_file(hf_hub_download("d0mzw/pvml-tinystories", f"{run}/model.safetensors")))
model.eval()

Sampling is easiest through experiments/sample.py in the GitHub repo, which loads a run directory and gives you a prompt.

Training

• All runs use batch size 32, AdamW at lr=1e-3, weight decay 0.01, seed 0
• Data is the full TinyStories train split, concatenated into one token stream and chopped into fixed n_ctx chunks, so no padding is needed
• Hardware is an ASUS ROG Flow Z13, Radeon 8060S integrated graphics, ROCm

Limitations

These are teaching artifacts, not useful language models. The largest is 32M parameters trained for 10.7 hours on children's stories, and it writes children's stories with visible repetition and referents that drift. Nothing here is suitable for any downstream use.

Disclaimer

The code comes from working through the ARENA 3.0 curriculum. I claim no credit for the original material, and this is not affiliated with or endorsed by ARENA.