commit 5a7f2177bc485722303de68108033e75bd8ae18a Author: rzen Date: Sun Jul 12 10:24:52 2026 -0400 Phase 0: project plan, pinned environment, data download + MPS smoke test - README with full pipeline plan (tokenizer -> data -> model -> training -> eval) - pip/venv setup pinned to torch 2.13.0, numpy 2.5.1 (Python 3.13) - scripts/download_data.py: one-time TinyStoriesV2 fetch with SHA-256 recording - scripts/check_mps.py: verifies MPS backend, backward pass, GPU speedup diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..249d86d --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +.venv/ +__pycache__/ +*.pyc +data/ +checkpoints/ +.DS_Store diff --git a/README.md b/README.md new file mode 100644 index 0000000..6a0a478 --- /dev/null +++ b/README.md @@ -0,0 +1,161 @@ +# 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 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 +├── 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 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..01c134d --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +torch==2.13.0 +numpy==2.5.1 diff --git a/scripts/check_mps.py b/scripts/check_mps.py new file mode 100644 index 0000000..e5ccbca --- /dev/null +++ b/scripts/check_mps.py @@ -0,0 +1,50 @@ +"""Smoke test: PyTorch can see the Apple GPU and run a training step on it. + +Checks three things: + 1. the MPS backend is available, + 2. forward + backward passes produce finite gradients on MPS, + 3. MPS matmul is actually faster than CPU (i.e., the GPU is really used). + +Usage: python scripts/check_mps.py +""" + +import time + +import torch + + +def bench_matmul(device: str, n: int = 2048, iters: int = 20) -> float: + """Mean seconds per (n x n) matmul; .item() forces device sync.""" + x = torch.randn(n, n, device=device) + y = torch.randn(n, n, device=device) + for _ in range(3): # warmup + (x @ y).sum().item() + start = time.perf_counter() + for _ in range(iters): + (x @ y).sum().item() + return (time.perf_counter() - start) / iters + + +def main() -> None: + print(f"torch {torch.__version__}") + assert torch.backends.mps.is_available(), "MPS backend not available" + assert torch.backends.mps.is_built(), "torch was built without MPS support" + print("MPS backend: available") + + # A miniature training step: forward, loss, backward, finite grads. + w = torch.randn(64, 64, device="mps", requires_grad=True) + x = torch.randn(128, 64, device="mps") + loss = ((x @ w) ** 2).mean() + loss.backward() + assert w.grad is not None and torch.isfinite(w.grad).all() + print(f"backward pass on MPS: ok (loss={loss.item():.4f})") + + cpu = bench_matmul("cpu") + mps = bench_matmul("mps") + print(f"2048x2048 matmul: cpu {cpu * 1e3:.1f} ms | mps {mps * 1e3:.1f} ms" + f" | speedup {cpu / mps:.1f}x") + print("smoke test passed") + + +if __name__ == "__main__": + main() diff --git a/scripts/download_data.py b/scripts/download_data.py new file mode 100644 index 0000000..a898edb --- /dev/null +++ b/scripts/download_data.py @@ -0,0 +1,71 @@ +"""One-time download of the TinyStoriesV2 corpus. + +This is the only network access in the entire project. Files land in +data/raw/ (gitignored). SHA-256 hashes are recorded in data/raw/SHA256SUMS +so the corpus can be verified later without re-downloading. + +Usage: python scripts/download_data.py +""" + +import hashlib +import urllib.request +from pathlib import Path + +BASE_URL = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main" +FILES = [ + "TinyStoriesV2-GPT4-train.txt", # ~2.2 GB + "TinyStoriesV2-GPT4-valid.txt", # ~22 MB +] +RAW_DIR = Path(__file__).resolve().parent.parent / "data" / "raw" +CHUNK = 1 << 20 # 1 MB + + +def download(name: str) -> Path: + dest = RAW_DIR / name + if dest.exists(): + print(f"{name}: already present, skipping download") + return dest + tmp = dest.with_suffix(dest.suffix + ".part") + url = f"{BASE_URL}/{name}" + print(f"{name}: downloading") + with urllib.request.urlopen(url) as response, open(tmp, "wb") as f: + total = int(response.headers.get("Content-Length") or 0) + done = 0 + while chunk := response.read(CHUNK): + f.write(chunk) + done += len(chunk) + if total: + print( + f"\r {done / 1e6:,.0f} / {total / 1e6:,.0f} MB" + f" ({done * 100 // total}%)", + end="", + flush=True, + ) + print() + tmp.rename(dest) # only a fully written file ever gets the real name + return dest + + +def sha256(path: Path) -> str: + h = hashlib.sha256() + with open(path, "rb") as f: + while chunk := f.read(CHUNK): + h.update(chunk) + return h.hexdigest() + + +def main() -> None: + RAW_DIR.mkdir(parents=True, exist_ok=True) + lines = [] + for name in FILES: + path = download(name) + digest = sha256(path) + print(f"{name}: {path.stat().st_size / 1e6:,.0f} MB sha256={digest}") + lines.append(f"{digest} {name}\n") + sums = RAW_DIR / "SHA256SUMS" + sums.write_text("".join(lines)) + print(f"hashes recorded in {sums}") + + +if __name__ == "__main__": + main()