experiment-llm

Training a small language model completely from scratch on a MacBook Pro (M1): own tokenizer, own transformer implementation, own training loop. No pretrained weights, no model libraries, no cloud compute.

The goal of iteration 1 is completeness and clarity, not model quality: a small model (~4–5M parameters) taken through the entire pipeline end to end, with every component written to be read and understood.

Principles

• From scratch. The tokenizer, model, training loop, and sampler are implemented in this repo. Dependencies are limited to PyTorch (tensor ops + autograd + MPS backend) and NumPy. No HuggingFace libraries, no nanoGPT imports — reference implementations are for reading, not importing.
• Self-sufficient. After a one-time dataset download, everything runs offline on the local machine. No API calls, no cloud GPUs, no external evaluation services.
• Clarity over speed. Straightforward code beats clever code. Every phase produces an inspectable artifact (tokenizer vocab, token dumps, loss curves, sample outputs).
• Config-driven. One experiment = one config. Model size, context length, and training schedule live in a config file, so iteration 2 is a config change, not a rewrite.

Setup

Requires Python 3.13 (Homebrew) on Apple Silicon.

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

python scripts/check_mps.py       # verify the Apple GPU is usable
python scripts/download_data.py   # one-time, ~2.2 GB into data/raw/

After the download, everything runs offline.

Why TinyStories

General web text is hopeless at this scale — a 5M-parameter model trained on it produces word salad. TinyStories (Eldan & Li, 2023) is a ~2GB corpus of synthetic children's stories with a deliberately restricted vocabulary, built precisely so that very small models can learn coherent English. It is the smallest known regime where a from-scratch model produces output that makes sense — grammatical sentences, consistent characters, simple plots. Our v1 model sits below the paper's coherence sweet spot (~30M params); some incoherence is expected and fine.

Pipeline

raw text ──► tokenizer training ──► tokenized corpus ──► pretraining ──► sampling / eval
 (phase 1)       (phase 2)              (phase 2)          (phase 3)        (phase 4)

Phase 0 — Setup

• Standard venv + pip environment; dependencies (torch, numpy) pinned to exact versions in requirements.txt.
• One-time download of TinyStoriesV2 (plain .txt files fetched over HTTPS into data/raw/, which is gitignored). This is the only network access in the whole project.
• Deliverable: scripts/download_data.py, pinned environment, smoke test that MPS is available.

Phase 1 — Tokenizer

• Byte-pair encoding implemented from scratch: trainer, encoder, decoder.
• Trained on a subsample of the corpus (BPE training is quadratic-ish in pure Python; a ~50MB sample is plenty for a 4k vocab on this data).
• Vocab size 4,096 — TinyStories' restricted vocabulary doesn't need more, and a small vocab keeps the embedding table (and the model) small.
• Deliverables: src/tokenizer.py, saved vocab/merges file, round-trip tests (decode(encode(s)) == s), a short report on compression ratio (bytes per token).

Phase 2 — Data pipeline

• Encode the full corpus once with the trained tokenizer into flat binary token files (data/tokens/train.bin, val.bin) — memory-mapped at training time, so the 2GB corpus never has to fit in RAM as Python objects.
• Batching: random crops of context_length + 1 tokens; inputs/targets are shifted views of the same crop.
• Deliverables: src/dataset.py, scripts/encode_corpus.py, token counts logged for train/val.

Phase 3 — Model

• Decoder-only transformer, GPT-style, implemented module by module: token embeddings, causal self-attention, MLP block, RMSNorm, residual stream, weight-tied output head.

• v1 configuration (~4.3M params):

  knob	value
  layers	4
  d_model	256
  heads	4
  context length	256
  vocab	4,096
  positions	learned embeddings (v1 simplicity; RoPE is an iteration-2 candidate)

• Deliverables: src/model.py, a parameter-count printout that matches the hand-computed estimate, a forward-pass shape test.

Phase 4 — Training

• Plain training loop: AdamW, cosine LR schedule with warmup, gradient clipping, periodic validation loss, checkpointing (resumable).
• Runs on MPS; expected wall-clock for v1 is a few hours for ~100–200M tokens (roughly Chinchilla-optimal for this size — we don't need to see the whole corpus).
• Loss curve logged to a plain CSV (reports/loss-v1.csv) and rendered as an ASCII chart by scripts/plot_loss.py — no plotting dependency; the chart is also saved to reports/loss-v1.txt as a committable artifact.
• Deliverables: src/train.py, configs/v1.py, scripts/plot_loss.py, checkpoints in checkpoints/ (gitignored), loss log + chart in reports/.

Commands:

python src/train.py --config v1 --overfit   # sanity check (loss -> ~0)
python src/train.py --config v1             # real run (resumable: --resume)
python scripts/plot_loss.py --config v1     # ASCII loss curve

Phase 5 — Sampling & evaluation

• Autoregressive sampler with temperature and top-k.
• Evaluation is fully local:
  • validation loss / perplexity,
  • a fixed prompt suite (story openings) sampled at every checkpoint so progress is visible qualitatively over training,
  • a sample gallery in reports/ for the final model.
• Deliverables: src/sample.py, reports/samples-v1.md.

Project layout

experiment-llm/
├── README.md
├── requirements.txt
├── configs/            # one file per experiment (v1, v2, ...)
├── scripts/            # one-shot entry points: download, encode, plots
├── artifacts/          # trained tokenizer.json (small, committed)
├── src/
│   ├── tokenizer.py    # BPE: train / encode / decode
│   ├── dataset.py      # memmapped token files, batch sampling
│   ├── model.py        # the transformer
│   ├── train.py        # training loop
│   └── sample.py       # generation
├── tests/              # round-trip + shape + overfit-one-batch tests
├── data/               # gitignored: raw/ (txt), tokens/ (bin)
├── checkpoints/        # gitignored
└── reports/            # loss curves, sample galleries (committed)

Definition of done (iteration 1)

• Environment reproducible with pip install -r requirements.txt; dataset downloaded and hashed
• BPE tokenizer trains, round-trips, and is saved to disk
• Corpus encoded to binary token files
• Model builds; parameter count matches the estimate
• Overfit-one-batch sanity check passes (loss → ~0 on a single batch)
• Full training run completes on MPS with checkpoints + loss curve (v1: 20k steps, ~58 min on M1, val loss 1.52 / perplexity ~4.6)
• Sampler produces stories from prompts; gallery committed to reports/ (v1: grammatical, story-arc coherent; plot logic wobbles as expected)
• Every phase re-runnable from a clean clone with documented commands

Degradation experiments (iteration 1.5)

A controlled study of how a model gets worse: start from the v1 baseline and change exactly one variable per run, so any difference in the output is attributable to that variable. Each variant config is literally v1 with one field replaced (see configs/*.py); tests/test_configs.py enforces that only the intended field changes.

config	one change from v1	isolates	expected failure mode
v1	— (baseline)	—	grammatical, coherent story arcs
small	2L/64d, ~0.38M params	capacity	generic, repetitive, forgets rare words
short	max_iters 20k→2k	training time	half-learned grammar, weaker structure
ctx	context_length 256→64	long-range coherence	fine locally, plot threads break sooner
data	max_train_tokens→1M	data quantity	overfits: memorizes, novel prompts collapse

Run each (they all train faster than v1; short is ~minutes):

python src/train.py --config small     # then short, ctx, data
python scripts/plot_loss.py --config small

Then compare generated text across every trained variant, side by side (same prompt, same RNG seed, so differences reflect the model):

python scripts/compare.py                       # default prompt, all configs
python scripts/compare.py --suite               # every suite prompt -> reports/compare.md
python scripts/compare.py --configs v1,small    # a subset

Iteration 2 candidates (out of scope for now)

• Scale toward the coherence sweet spot (~15–30M params, longer training)
• RoPE positions, SwiGLU MLP, other modern architecture upgrades
• Instruction tuning on TinyStories-Instruct (pretrain → SFT arc)
• MLX port for faster training on Apple Silicon