---
license: mit
library_name: pytorch
pipeline_tag: text-generation
tags:
  - transformer
  - tinystories
  - from-scratch
datasets:
  - roneneldan/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.

```python
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.