# 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.

```sh
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](https://arxiv.org/abs/2305.07759)
(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:

```sh
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)

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

```sh
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):

```sh
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