TinyStories Transformer (from scratch)
A ~23M-parameter decoder-only Transformer, trained from scratch on
TinyStories, with
every component implemented directly on raw PyTorch
tensor ops, no nn.Transformer, no HuggingFace transformers classes.
For details, check out the GitHub Repo.
Model details
| Parameters | ~22.7M |
| Architecture | Decoder-only Transformer, pre-norm, RoPE, SwiGLU |
| d_model | 512 |
| Layers | 4 |
| Heads | 16 |
| d_ff | 1344 |
| Context length | 256 |
| Vocab size | 10,000 |
| Tokenizer | Byte-pair encoding, trained from scratch on TinyStories (zehl/tinystories-tokenized-10k) |
| Training steps | 5,000 |
| Validation loss | 1.6206 |
| Optimizer | AdamW (alpha_max=0.001), cosine LR schedule, global-norm gradient clipping |
Usage
uv venv --python 3.12
uv pip install "cs336_basics @ git+https://github.com/zhL-d/llm-from-scratch.git@2b1593c1bc049dde93d332b9f5a6cf9d57575744"
uv pip install huggingface_hub
uv run generate.py "Once upon a time"
Example output for "Once upon a time":
Once upon a time, there was a little boy named Tim. Tim was an ordinary boy who loved to play outside. One day, he went to the park to play. At the park, Tim saw a big tree. He wanted to climb it. He started to climb the tree. He climbed higher and higher. When he was high, he saw a big bird. The bird was singing a happy song. Tim was surprised. He did not know birds could sing. The bird sang a happy song. Tim listened to the bird sing and listened. He was not scared anymore. He played with the bird and they became best friends.
generate.py loads model.pt + config.json from here,
pulls the tokenizer's vocab.json/merges.json from the
dataset repo,
and runs top-p sampling (temperature=0.7, top_p=0.9) via the
Decoding
function from the zhL-d/llm-from-scratch repo, pinned to the exact commit this
checkpoint was produced under.
Source
- Code:
zhL-d/llm-from-scratch - Tokenizer + tokenized dataset:
zehl/tinystories-tokenized-10k - Training data: TinyStories (Eldan & Li, 2023)
- Downloads last month
- 124