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
This commit is contained in:
Cormac Shannon committed
1 parent
971caa808c
commit
29e323b2fc
8 files changed
+435
-7
No files matched your search
@@ -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:])
|
||||
Reference in new issue
Block a user