Add QK-norm toggle, TinyStories prep and automatic run logging

nanoGPT gets a qk_norm flag (off = upstream behaviour), a seed option and
a resolved-config dump. run.py logs each run's git commit, config, data
checksums, environment and metrics, and stops diverged runs.

Claude-Session: https://claude.ai/code/session_01S8KN7NuUCm8nXLmy6eDuVv

06-experiments/data/prepare_tinystories.py

@@ -0,0 +1,101 @@
+#!/usr/bin/env python3
+"""Download TinyStories (Eldan & Li 2023, V2 / GPT-4 split) and tokenise it for nanoGPT.
+
+Usage (from 06-experiments/):
+    .venv/bin/python data/prepare_tinystories.py            # full: ~2.2 GB train text
+    .venv/bin/python data/prepare_tinystories.py --spike    # small: valid file only, split 90/10
+
+Writes data/<name>/{train,val}.bin (uint16 GPT-2 BPE tokens, nanoGPT format) and
+data/<name>/prep.json (source URLs, SHA-256 of the raw files, token counts), so every
+run can record exactly which data it used. Raw text and .bin files are git-ignored.
+"""
+import hashlib
+import json
+import os
+import sys
+import urllib.request
+
+import numpy as np
+import tiktoken
+
+BASE = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/"
+TRAIN, VALID = "TinyStoriesV2-GPT4-train.txt", "TinyStoriesV2-GPT4-valid.txt"
+HERE = os.path.dirname(os.path.abspath(__file__))
+SEP = "<|endoftext|>"
+
+
+def download(fname, raw_dir):
+    path = os.path.join(raw_dir, fname)
+    if not os.path.exists(path):
+        print(f"downloading {fname} ...")
+        tmp = path + ".part"
+        urllib.request.urlretrieve(BASE + fname, tmp)
+        os.rename(tmp, path)
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            h.update(chunk)
+    return path, h.hexdigest()
+
+
+def stories(path):
+    """Yield one story at a time (stories are separated by <|endoftext|>)."""
+    buf = []
+    with open(path, encoding="utf-8") as f:
+        for line in f:
+            if SEP in line:
+                head, _, tail = line.partition(SEP)
+                buf.append(head)
+                s = "".join(buf).strip()
+                if s:
+                    yield s
+                buf = [tail]
+            else:
+                buf.append(line)
+    s = "".join(buf).strip()
+    if s:
+        yield s
+
+
+def write_bin(texts, out_path, enc):
+    eot = enc.eot_token
+    ids = []
+    for t in texts:
+        ids.extend(enc.encode_ordinary(t))
+        ids.append(eot)
+    arr = np.array(ids, dtype=np.uint16)
+    arr.tofile(out_path)
+    return len(arr)
+
+
+def main():
+    spike = "--spike" in sys.argv
+    name = "tinystories_spike" if spike else "tinystories"
+    out_dir = os.path.join(HERE, name)
+    raw_dir = os.path.join(HERE, "raw")
+    os.makedirs(out_dir, exist_ok=True)
+    os.makedirs(raw_dir, exist_ok=True)
+    enc = tiktoken.get_encoding("gpt2")
+    meta = {"dataset": "TinyStories V2 (GPT-4)", "source": BASE, "tokenizer": "tiktoken gpt2", "files": {}}
+
+    vpath, vsha = download(VALID, raw_dir)
+    meta["files"][VALID] = vsha
+    if spike:
+        allv = list(stories(vpath))
+        cut = int(len(allv) * 0.9)
+        meta["split"] = f"valid file split 90/10 by story order: {cut} train / {len(allv) - cut} val stories"
+        n_tr = write_bin(allv[:cut], os.path.join(out_dir, "train.bin"), enc)
+        n_va = write_bin(allv[cut:], os.path.join(out_dir, "val.bin"), enc)
+    else:
+        tpath, tsha = download(TRAIN, raw_dir)
+        meta["files"][TRAIN] = tsha
+        meta["split"] = "official train / valid files"
+        n_tr = write_bin(stories(tpath), os.path.join(out_dir, "train.bin"), enc)
+        n_va = write_bin(stories(vpath), os.path.join(out_dir, "val.bin"), enc)
+    meta["tokens"] = {"train": n_tr, "val": n_va}
+    json.dump(meta, open(os.path.join(out_dir, "prep.json"), "w"), indent=2)
+    print(json.dumps(meta, indent=2))
+
+
+if __name__ == "__main__":
+    main()