rzen
9fd56b6063
Controlled study of how the model degrades. Four variants, each = v1 with a single field replaced (dataclasses.replace), so differences are attributable: - small: capacity down (2L/64d, ~0.38M params) - short: training time down (max_iters 20k->2k) - ctx: context window down (256->64) - data: data quantity down (max_train_tokens->1M; new TrainConfig knob + train.py slice) - scripts/compare.py: sample same prompt across all trained configs with a shared seed, print side by side, write reports/compare.md - tests/test_configs.py: enforces one-parameter-at-a-time (only intended fields differ from v1) + small param count (3 tests). Full suite: 29 passing.
207 lines
9.0 KiB
Markdown
207 lines
9.0 KiB
Markdown
# 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
|