alex_llm

A small LLM built from scratch — tokenizer to supervised fine-tuning — to understand every layer of the stack, including the experiments that killed its original hypothesis. Trained on an RTX 2080 (8GB) on TinyStoriesV2 (~2.7M GPT-4-written children's stories, ~590M BPE tokens).

If you're here to understand how it works, start with docs/how-it-works.md — an end-to-end walkthrough from raw text to sampled story, keyed to this codebase.

Two models, both 4-layer pre-norm transformers with tied 128-dim embeddings and an 8k byte-level BPE (~1.85M params):

Base model, given a story opening:

The alien came down from the sky and it had a very special gift. He had found out what was inside and he was so proud of his discovery. [eos]

Instruct model, given the directive Words: dragon, cake, brave:

Once upon a time, there was a brave little boy. He had a big sister called a dragon. The dragon was very big and powerful. [...] The dragon wanted to eat the cake. So he jumped up and ran to the cake. But the cake was too high for the dragon. [...] He had just eaten the cake! He was so happy. He had saved the day. [eos]

All required words used, story arc closed, generation self-terminated.

The negative result this repo is built on: it began as a hyperbolic-geometry experiment — Poincaré-ball embeddings and distance-based attention, on the theory that hyperbolic space embeds linguistic hierarchy better at low dimension. The controlled comparison (identical params/tokenizer/data/seed, geometry the only variable) said otherwise: euclidean 5.43 vs hyperbolic 6.72 val perplexity. The hyperbolic path was removed; it lives in git history before commit ac9e011.

pip install -r requirements.txt

Download TinyStoriesV2-GPT4-train.txt and TinyStoriesV2-GPT4-valid.txt from the TinyStories dataset page into the repo root.

python -m alex_llm --train --data TinyStoriesV2-GPT4-train.txt ^
    --val-data TinyStoriesV2-GPT4-valid.txt ^
    --model ts_v2.pt --epochs 5 --seq-len 256 --stride 256 ^
    --batch 64 --vocab 8000 --seed 1337

First run tokenizes and caches the corpus next to the data file; later runs memory-map the cache. With --val-data, validation runs after each epoch, best-model selection uses validation loss, and per-epoch metrics are appended to <model>.metrics.jsonl.

On an RTX 2080 (8GB): batch 64 fits comfortably; ~14 min/epoch. torch.compile is used automatically when Triton is available.

python -m alex_llm --sft --base ts_v2.pt ^
    --data TinyStories-Instruct-train.txt ^
    --val-data TinyStories-Instruct-valid.txt ^
    --model ts_instruct.pt --epochs 2 --batch 32

Fine-tunes a pretrained base on TinyStories-Instruct: prompt = the record's directive fields (Summary/Words/Features/Random sentence) + Story:, response = the story + <eos>. Loss is masked to response tokens, so the model learns to follow the instruction rather than imitate it. SFT runs at seq 512 (the position table is extended from the base's 256; new rows are learned during fine-tuning), with LR warmup + cosine decay and weight decay excluded from biases, norms, and embeddings.

Chatting with an instruct model, your input becomes its Summary: field (or write Summary:/Words:/Features: fields yourself):

python chat.py -m ts_instruct.pt
you> a boy and a dragon become friends

python -m alex_llm --eval -m ts_v2.pt --data TinyStoriesV2-GPT4-valid.txt

Reports held-out loss, perplexity, and bits-per-byte. Perplexity is only comparable between models sharing a tokenizer and vocab; bits-per-byte is comparable across tokenizers (the legacy word tokenizer's bpb is optimistic — it lowercases and drops some characters).

python -m alex_llm -m ts_v2.pt              # canned sample prompts
python chat.py -m ts_v2.pt --temp 0.8       # interactive REPL

These are story-completion models, not instruction-tuned: give them a story opening, not a question. /help in the REPL lists sampling controls (temperature, top-k, top-p).

alex_llm/
  tokenizer.py   byte-level BPE (default) + legacy word tokenizer
  data.py        corpus loading, token cache, dataset
  model.py       TransformerLM (fused SDPA attention, tied embeddings)
  train.py       training loop, AMP, checkpointing, resume
  eval.py        held-out loss / perplexity / bits-per-byte
  cli.py         argument parsing and entry point
chat.py          interactive chat harness
tests/           model and dataset tests (pytest)

python -m pytest tests/ -q