tiny-llama

A small Llama-style decoder written in PyTorch for learning and experimentation.

Implemented from scratch:

• RMSNorm, RoPE, grouped-query attention, and SwiGLU
• causal and padded attention masks
• rolling KV cache with absolute RoPE positions
• byte-level BPE tokenizer with BOS/EOS tokens
• greedy, temperature, and top-k generation
• BF16 training, cosine learning-rate decay, early stopping, and checkpoints

Setup

The project uses Python 3.13 and uv.

uv sync
uv run pytest -q

You may need to install a platform-specific PyTorch build for CUDA or ROCm.

TinyStories

The main experiment uses a 14.36M-parameter model, a 1,026-token vocabulary, and a 256-token context window.

Prepare a local dataset subset and tokenizer:

uv run python -m experiments.tinystories.prepare

Train or resume training:

uv run python -m experiments.tinystories.train

Generate from the best checkpoint:

uv run python -m experiments.tinystories.inference

The reference run used 30.44M training tokens. Its best checkpoint was step 19,000 with validation loss 1.6616 and perplexity 5.27.

Layout

tiny_llama/                model, tokenizer, generation, and data utilities
experiments/tinystories/   dataset preparation and training scripts
tests/                     unit and integration tests

Experiment data and checkpoints live beside their scripts and are ignored by Git. This is an educational implementation, not a loader for official Llama weights or a replacement for optimized training frameworks.