File size: 821 Bytes
4397e12 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | """Compare val losses of runs at their last common eval step (matched tokens) and report tok/s.
Usage: $TA_PY scripts/compare_runs.py runA runB ..."""
import json, sys
runs = {}
for d in sys.argv[1:]:
recs = [json.loads(l) for l in open(f"{d}/log.jsonl")]
runs[d] = recs
common = set.intersection(*[{r["step"] for r in recs if "val" in r} for recs in runs.values()])
if not common:
sys.exit("no common eval step yet")
s = max(common)
for d, recs in runs.items():
v = next(r for r in recs if r["step"] == s and "val" in r)
toks = [r["tok_s"] for r in recs[-20:]]
mean = sum(v["val"].values()) / len(v["val"])
print(f"{d}: step {s} tokens {v['tokens']/1e6:.0f}M mean_val {mean:.4f} tok/s {sum(toks)/len(toks):.0f} "
+ " ".join(f"{k}={x:.3f}" for k, x in sorted(v["val"].items())))
|