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

.gitignore

@@ -29,7 +29,7 @@ __pycache__/
 .venv/
 
 # Experiment data, checkpoints and raw outputs (run metadata in runs/ IS tracked)
-06-experiments/data/
+06-experiments/data/*/
 06-experiments/vendor/
 06-experiments/runs/*/ckpt.pt
 06-experiments/runs/*/stdout.log

06-experiments/configs/spike.py

@@ -0,0 +1,35 @@
+# Tech-spike config: a tiny model on the small TinyStories subset, to check that the
+# pipeline (data -> train on MPS -> run log) works end to end. NOT an experiment config:
+# sizes, steps and LR here are placeholders, not decisions.
+
+dataset = 'tinystories_spike'
+device = 'mps'
+dtype = 'float32'          # train.py only uses autocast on CUDA; MPS runs in fp32
+compile = False            # torch.compile on MPS is not worth it for a spike
+
+# model (~3.2M non-embedding params; the GPT-2 vocab adds ~12.9M embedding params)
+n_layer = 4
+n_head = 4
+n_embd = 256
+block_size = 256
+bias = False
+dropout = 0.0
+qk_norm = False
+
+# optimisation (AdamW, nanoGPT defaults otherwise)
+batch_size = 32
+gradient_accumulation_steps = 1
+learning_rate = 1e-3
+max_iters = 300
+lr_decay_iters = 300
+min_lr = 1e-4
+warmup_iters = 30
+weight_decay = 0.1
+grad_clip = 1.0
+
+# logging / eval
+eval_interval = 100
+eval_iters = 20
+log_interval = 10
+always_save_checkpoint = False
+seed = 1337

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

06-experiments/nanogpt/README-thesis.md

@@ -0,0 +1,20 @@
+# nanoGPT (vendored)
+
+`model.py`, `train.py`, `configurator.py` and `LICENSE` were copied unmodified from
+[karpathy/nanoGPT](https://github.com/karpathy/nanoGPT) at commit
+`3adf61e154c3fe3fca428ad6bc3818b27a3b8291` (MIT licence), in their own commit.
+
+Every change since then is marked `[thesis]` in the code; `git log -p -- 06-experiments/nanogpt`
+shows them all. Summary:
+
+| File | Change | Why |
+|---|---|---|
+| `model.py` | `qk_norm` config flag; when on, LayerNorm (no bias by default) over each head's queries and keys before the dot product. `1/sqrt(d)` scaling kept | the independent variable. Off = exactly upstream |
+| `train.py` | `qk_norm`, `seed`, `data_root` config options | QK-norm toggle; seeds for repeat runs (upstream hard-codes 1337); data path |
+| `train.py` | `configurator.py` loaded relative to the script | so runs can launch from any directory |
+| `train.py` | writes resolved `config.json` into `out_dir` | run logging |
+| `train.py` | resume reads `qk_norm` from checkpoint | so resumed runs keep the same architecture |
+
+**Design choice to confirm:** QK-norm here is LayerNorm-style, as in Dehghani et al. (2023) and
+Wortsman et al. (2023). Henry et al. (2020) used L2-normalisation with a learnable scale instead,
+and OLMo 2 uses RMSNorm. `[Your call: which variant the thesis studies.]`

06-experiments/nanogpt/model.py

@@ -41,6 +41,14 @@ class CausalSelfAttention(nn.Module):
         self.n_head = config.n_head
         self.n_embd = config.n_embd
         self.dropout = config.dropout
+        # [thesis] QK-norm: LayerNorm over each head's query and key vectors before the
+        # dot product (the "qk-layernorm" of Dehghani et al. 2023 / Wortsman et al. 2023).
+        # Off by default, so qk_norm=False is exactly upstream nanoGPT.
+        self.qk_norm = config.qk_norm
+        if self.qk_norm:
+            head_dim = config.n_embd // config.n_head
+            self.q_norm = LayerNorm(head_dim, bias=config.bias)
+            self.k_norm = LayerNorm(head_dim, bias=config.bias)
         # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
         self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
         if not self.flash:
@@ -57,6 +65,9 @@ class CausalSelfAttention(nn.Module):
         k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
         q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
         v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        if self.qk_norm:  # [thesis] normalise over the head dimension
+            q = self.q_norm(q)
+            k = self.k_norm(k)
 
         # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
         if self.flash:
@@ -114,6 +125,7 @@ class GPTConfig:
     n_embd: int = 768
     dropout: float = 0.0
     bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
+    qk_norm: bool = False # [thesis] LayerNorm on per-head queries and keys
 
 class GPT(nn.Module):
 

06-experiments/nanogpt/train.py

@@ -54,6 +54,9 @@ n_head = 12
 n_embd = 768
 dropout = 0.0 # for pretraining 0 is good, for finetuning try 0.1+
 bias = False # do we use bias inside LayerNorm and Linear layers?
+qk_norm = False # [thesis] QK-norm on/off: the independent variable
+seed = 1337 # [thesis] base RNG seed (upstream hard-codes 1337)
+data_root = 'data' # [thesis] where dataset folders live
 # adamw optimizer
 learning_rate = 6e-4 # max learning rate
 max_iters = 600000 # total number of training iterations
@@ -74,7 +77,7 @@ dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported
 compile = True # use PyTorch 2.0 to compile the model to be faster
 # -----------------------------------------------------------------------------
 config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
-exec(open('configurator.py').read()) # overrides from command line or config file
+exec(open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'configurator.py')).read()) # overrides from command line or config file
 config = {k: globals()[k] for k in config_keys} # will be useful for logging
 # -----------------------------------------------------------------------------
 
@@ -103,7 +106,9 @@ print(f"tokens per iteration will be: {tokens_per_iter:,}")
 
 if master_process:
     os.makedirs(out_dir, exist_ok=True)
-torch.manual_seed(1337 + seed_offset)
+    import json  # [thesis] record the fully resolved config for the run log
+    with open(os.path.join(out_dir, 'config.json'), 'w') as f: json.dump(config, f, indent=2)
+torch.manual_seed(seed + seed_offset)
 torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
 torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
 device_type = 'cuda' if 'cuda' in device else 'cpu' # for later use in torch.autocast
@@ -112,7 +117,7 @@ ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torc
 ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
 
 # poor man's data loader
-data_dir = os.path.join('data', dataset)
+data_dir = os.path.join(data_root, dataset)
 def get_batch(split):
     # We recreate np.memmap every batch to avoid a memory leak, as per
     # https://stackoverflow.com/questions/45132940/numpy-memmap-memory-usage-want-to-iterate-once/61472122#61472122
@@ -145,7 +150,7 @@ if os.path.exists(meta_path):
 
 # model init
 model_args = dict(n_layer=n_layer, n_head=n_head, n_embd=n_embd, block_size=block_size,
-                  bias=bias, vocab_size=None, dropout=dropout) # start with model_args from command line
+                  bias=bias, vocab_size=None, dropout=dropout, qk_norm=qk_norm) # start with model_args from command line
 if init_from == 'scratch':
     # init a new model from scratch
     print("Initializing a new model from scratch")
@@ -163,8 +168,8 @@ elif init_from == 'resume':
     checkpoint_model_args = checkpoint['model_args']
     # force these config attributes to be equal otherwise we can't even resume training
     # the rest of the attributes (e.g. dropout) can stay as desired from command line
-    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
+    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size', 'qk_norm']:
-        model_args[k] = checkpoint_model_args[k]
+        model_args[k] = checkpoint_model_args.get(k, False)
     # create the model
     gptconf = GPTConfig(**model_args)
     model = GPT(gptconf)

06-experiments/requirements.txt

@@ -0,0 +1,19 @@
+# exact versions used for the tech spike (28/10/2026); install with: uv venv --python 3.12 .venv && uv pip install -r requirements.txt
+certifi==2026.10.23
+charset-normalizer==3.5.2
+filelock==4.2.4
+fsspec==2026.9.0
+idna==3.20
+jinja2==3.1.6
+markupsafe==3.0.4
+mpmath==1.3.0
+networkx==3.7
+numpy==2.5.4
+regex==2026.9.29
+requests==2.34.2
+setuptools==84.2.1
+sympy==1.14.0
+tiktoken==0.14.0
+torch==2.15.0
+typing-extensions==4.16.0
+urllib3==2.8.0

06-experiments/run.py

@@ -0,0 +1,236 @@
+#!/usr/bin/env python3
+"""Run one nanoGPT training run and log everything needed to reproduce it.
+
+Usage (from 06-experiments/):
+    .venv/bin/python run.py NAME CONFIG.py [key=value ...] [--note "why this run"]
+
+Example:
+    .venv/bin/python run.py spike-qk-on configs/spike.py qk_norm=True learning_rate=3e-3
+
+Creates runs/<YYYYmmdd-HHMMSS>_<NAME>/ containing:
+    run.json       what was run and how: git commit (+ dirty flag), nanoGPT upstream commit,
+                   full resolved config, dataset checksums, hardware/software, outcome, summary
+    config.json    every hyperparameter nanoGPT actually used (written by train.py)
+    metrics.jsonl  one line per logged iteration / eval
+    diff.patch     uncommitted changes at launch time (only if the repo was dirty)
+    stdout.log     raw output (git-ignored)
+    ckpt.pt        checkpoint (git-ignored)
+and appends one row to runs/index.csv — the at-a-glance table of every run.
+
+A run is stopped and marked "diverged" if the training loss becomes NaN/inf or
+exceeds --diverge-loss (default 20; the initial loss is ~10.8 for a 50k vocab).
+"""
+import csv
+import datetime as dt
+import json
+import math
+import os
+import platform
+import re
+import shlex
+import socket
+import subprocess
+import sys
+import time
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+RUNS = os.path.join(HERE, "runs")
+UPSTREAM_NANOGPT = "3adf61e154c3fe3fca428ad6bc3818b27a3b8291"  # see nanogpt/README-thesis.md
+
+ITER_RE = re.compile(r"^iter (\d+): loss ([-\w.]+), time ([\d.]+)ms")
+EVAL_RE = re.compile(r"^step (\d+): train loss ([-\w.]+), val loss ([-\w.]+)")
+PARAMS_RE = re.compile(r"number of parameters: ([\d.]+)M")
+
+
+def sh(*cmd):
+    try:
+        return subprocess.run(cmd, cwd=HERE, capture_output=True, text=True, check=True).stdout.strip()
+    except Exception:
+        return ""
+
+
+def git_info(run_dir):
+    sha = sh("git", "rev-parse", "HEAD")
+    dirty = sh("git", "status", "--porcelain", "--untracked-files=no")
+    info = {"commit": sha, "branch": sh("git", "rev-parse", "--abbrev-ref", "HEAD"), "dirty": bool(dirty)}
+    if dirty:
+        open(os.path.join(run_dir, "diff.patch"), "w").write(sh("git", "diff", "HEAD"))
+        info["diff_file"] = "diff.patch"
+    return info
+
+
+def env_info():
+    import torch  # imported here so `--help` works without torch
+
+    return {
+        "host": socket.gethostname(),
+        "platform": platform.platform(),
+        "machine": platform.machine(),
+        "cpu": sh("sysctl", "-n", "machdep.cpu.brand_string") or platform.processor(),
+        "python": platform.python_version(),
+        "torch": torch.__version__,
+        "mps_available": torch.backends.mps.is_available(),
+        "cuda_available": torch.cuda.is_available(),
+    }
+
+
+def data_info(dataset):
+    path = os.path.join(HERE, "data", dataset, "prep.json")
+    info = json.load(open(path)) if os.path.exists(path) else {"warning": "no prep.json found"}
+    for split in ("train", "val"):
+        p = os.path.join(HERE, "data", dataset, f"{split}.bin")
+        if os.path.exists(p):
+            info[f"{split}_bin_bytes"] = os.path.getsize(p)
+    return info
+
+
+def parse_value(v):
+    for cast in (int, float):
+        try:
+            return cast(v)
+        except ValueError:
+            pass
+    if v in ("True", "False"):
+        return v == "True"
+    return v
+
+
+def num(x):
+    try:
+        return float(x)
+    except ValueError:
+        return float("nan")
+
+
+def main(argv):
+    if len(argv) < 2 or argv[0] in ("-h", "--help"):
+        sys.exit(__doc__)
+    note = ""
+    if "--note" in argv:
+        i = argv.index("--note")
+        note = argv[i + 1]
+        del argv[i:i + 2]
+    diverge_loss = 20.0
+    if "--diverge-loss" in argv:
+        i = argv.index("--diverge-loss")
+        diverge_loss = float(argv[i + 1])
+        del argv[i:i + 2]
+    name, config_file, overrides = argv[0], argv[1], argv[2:]
+
+    # Resolve the config the same way nanoGPT will: config file, then overrides.
+    cfg = {}
+    exec(open(config_file).read(), {}, cfg)
+    cfg = {k: v for k, v in cfg.items() if not k.startswith("_")}
+    for o in overrides:
+        k, v = o.split("=", 1)
+        cfg[k] = parse_value(v)
+
+    run_id = dt.datetime.now().strftime("%Y%m%d-%H%M%S") + "_" + re.sub(r"[^\w.-]", "-", name)
+    run_dir = os.path.join(RUNS, run_id)
+    os.makedirs(run_dir)
+    cmd = [sys.executable, os.path.join(HERE, "nanogpt", "train.py"), config_file,
+           f"--out_dir={run_dir}", f"--data_root={os.path.join(HERE, 'data')}"]
+    cmd += [f"--{o}" for o in overrides]
+
+    record = {
+        "run_id": run_id,
+        "name": name,
+        "note": note,
+        "status": "running",
+        "started": dt.datetime.now().astimezone().isoformat(timespec="seconds"),
+        "command": " ".join(shlex.quote(c) for c in cmd),
+        "config_file": os.path.relpath(config_file, HERE),
+        "overrides": overrides,
+        "config_requested": cfg,
+        "git": git_info(run_dir),
+        "nanogpt_upstream_commit": UPSTREAM_NANOGPT,
+        "data": data_info(cfg.get("dataset", "?")),
+        "environment": env_info(),
+    }
+    rec_path = os.path.join(run_dir, "run.json")
+
+    def save():
+        json.dump(record, open(rec_path, "w"), indent=2)
+
+    save()
+    print(f"run {run_id}\n  {record['command']}")
+
+    metrics = open(os.path.join(run_dir, "metrics.jsonl"), "w")
+    log = open(os.path.join(run_dir, "stdout.log"), "w")
+    t0 = time.time()
+    status, diverged_at, params_m = "completed", None, None
+    first_loss = last_loss = None
+    best_val, last_eval = math.inf, None
+    proc = subprocess.Popen(cmd, cwd=HERE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1,
+                            env={**os.environ, "PYTHONUNBUFFERED": "1"})  # stream lines so divergence stops the run early
+    try:
+        for line in proc.stdout:
+            log.write(line)
+            line = line.strip()
+            if m := PARAMS_RE.search(line):
+                params_m = float(m.group(1))
+            elif m := ITER_RE.match(line):
+                it, loss = int(m.group(1)), num(m.group(2))
+                metrics.write(json.dumps({"type": "train", "iter": it, "loss": loss, "ms": float(m.group(3)),
+                                          "t": round(time.time() - t0, 2)}) + "\n")
+                first_loss = loss if first_loss is None else first_loss
+                last_loss = loss
+                if not math.isfinite(loss) or loss > diverge_loss:
+                    status, diverged_at = "diverged", it
+                    print(f"  diverged at iter {it} (loss {loss}); stopping")
+                    proc.terminate()
+                    break
+            elif m := EVAL_RE.match(line):
+                it, tr, va = int(m.group(1)), num(m.group(2)), num(m.group(3))
+                metrics.write(json.dumps({"type": "eval", "iter": it, "train_loss": tr, "val_loss": va,
+                                          "t": round(time.time() - t0, 2)}) + "\n")
+                last_eval = {"iter": it, "train_loss": tr, "val_loss": va}
+                if math.isfinite(va):
+                    best_val = min(best_val, va)
+                print(f"  step {it}: train {tr:.4f}  val {va:.4f}")
+            metrics.flush()
+        proc.wait()
+        if status == "completed" and proc.returncode != 0:
+            status = "failed"
+    except KeyboardInterrupt:
+        proc.terminate()
+        status = "interrupted"
+    finally:
+        metrics.close()
+        log.close()
+
+    resolved = os.path.join(run_dir, "config.json")
+    cfg_used = json.load(open(resolved)) if os.path.exists(resolved) else cfg
+    record.update({
+        "status": status,
+        "finished": dt.datetime.now().astimezone().isoformat(timespec="seconds"),
+        "wall_seconds": round(time.time() - t0, 1),
+        "summary": {
+            "params_millions_nanogpt_count": params_m,
+            "first_train_loss": first_loss,
+            "last_train_loss": last_loss,
+            "last_eval": last_eval,
+            "best_val_loss": best_val if math.isfinite(best_val) else None,
+            "diverged_at_iter": diverged_at,
+            "returncode": proc.returncode,
+        },
+    })
+    save()
+
+    index = os.path.join(RUNS, "index.csv")
+    cols = ["run_id", "status", "qk_norm", "learning_rate", "n_layer", "n_embd", "seed", "max_iters",
+            "params_M", "best_val_loss", "diverged_at", "wall_s", "git_commit", "dirty", "note"]
+    new = not os.path.exists(index)
+    with open(index, "a", newline="") as f:
+        w = csv.writer(f)
+        if new:
+            w.writerow(cols)
+        w.writerow([run_id, status, cfg_used.get("qk_norm"), cfg_used.get("learning_rate"), cfg_used.get("n_layer"),
+                    cfg_used.get("n_embd"), cfg_used.get("seed"), cfg_used.get("max_iters"), params_m,
+                    record["summary"]["best_val_loss"], diverged_at, record["wall_seconds"],
+                    record["git"]["commit"][:10], record["git"]["dirty"], note])
+    print(f"  {status}; best val {record['summary']['best_val_loss']}; {record['wall_seconds']}s -> {os.path.relpath(run_dir, HERE)}")
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])