From 29e323b2fcf3fadaea73563ffd2a28a19e7d6ed9 Mon Sep 17 00:00:00 2001 From: Cormac Shannon <> Date: Wed, 28 Oct 2026 18:05:31 +0000 Subject: [PATCH] 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 | 2 +- 06-experiments/configs/spike.py | 35 +++ 06-experiments/data/prepare_tinystories.py | 101 +++++++++ 06-experiments/nanogpt/README-thesis.md | 20 ++ 06-experiments/nanogpt/model.py | 12 ++ 06-experiments/nanogpt/train.py | 17 +- 06-experiments/requirements.txt | 19 ++ 06-experiments/run.py | 236 +++++++++++++++++++++ 8 files changed, 435 insertions(+), 7 deletions(-) create mode 100644 06-experiments/configs/spike.py create mode 100644 06-experiments/data/prepare_tinystories.py create mode 100644 06-experiments/nanogpt/README-thesis.md create mode 100644 06-experiments/requirements.txt create mode 100644 06-experiments/run.py diff --git a/.gitignore b/.gitignore index 0f6a8e9..b1c5073 100644 --- a/.gitignore +++ b/.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 diff --git a/06-experiments/configs/spike.py b/06-experiments/configs/spike.py new file mode 100644 index 0000000..3d279c3 --- /dev/null +++ b/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 diff --git a/06-experiments/data/prepare_tinystories.py b/06-experiments/data/prepare_tinystories.py new file mode 100644 index 0000000..c9e16aa --- /dev/null +++ b/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//{train,val}.bin (uint16 GPT-2 BPE tokens, nanoGPT format) and +data//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() diff --git a/06-experiments/nanogpt/README-thesis.md b/06-experiments/nanogpt/README-thesis.md new file mode 100644 index 0000000..5aa902c --- /dev/null +++ b/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.]` diff --git a/06-experiments/nanogpt/model.py b/06-experiments/nanogpt/model.py index c698f8b..f660f96 100644 --- a/06-experiments/nanogpt/model.py +++ b/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): diff --git a/06-experiments/nanogpt/train.py b/06-experiments/nanogpt/train.py index de57850..6db5158 100644 --- a/06-experiments/nanogpt/train.py +++ b/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']: - model_args[k] = checkpoint_model_args[k] + for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size', 'qk_norm']: + model_args[k] = checkpoint_model_args.get(k, False) # create the model gptconf = GPTConfig(**model_args) model = GPT(gptconf) diff --git a/06-experiments/requirements.txt b/06-experiments/requirements.txt new file mode 100644 index 0000000..123720a --- /dev/null +++ b/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 diff --git a/06-experiments/run.py b/06-experiments/run.py new file mode 100644 index 0000000..04f7529 --- /dev/null +++ b/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/_/ 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:])