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 and plotted locally.
• Deliverables: src/train.py, configs/v1.yaml (or .py), checkpoints in checkpoints/ (gitignored), loss curves in reports/.

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
• Sampler produces stories from prompts; gallery committed to reports/
• Every phase re-runnable from a clean clone with documented commands

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