Cormac Shannon
29e323b2fc
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
237 lines
8.9 KiB
Python
237 lines
8.9 KiB
Python
#!/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:])
|