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

.gitignore

@@ -0,0 +1,6 @@
+.venv/
+__pycache__/
+*.pyc
+data/
+checkpoints/
+.DS_Store

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

requirements.txt

@@ -0,0 +1,2 @@
+torch==2.13.0
+numpy==2.5.1

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()

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()