| """Matched GPT-2 small eval for Stasis step 36,420 endpoint. Reuses the frozen package code |
| (stasis_gpt2_eval.py, METRIC_DEFINITIONS 579ab8c1...) unchanged; only the checkpoint loader differs |
| (accepts step-36420 endpoint by sha). Items, tokenization, tok_budget 8192, BOS, bootstrap identical. |
| Internal row key 'STASIS_12269' is the frozen Stasis-tokenizer key; outputs relabelled STASIS_36420.""" |
| import json, os, sys, time, platform, argparse |
| import torch |
| PKG = os.environ.get("GPT2EVAL_PKG", "package") |
| sys.path.insert(0, os.path.join(PKG, "code")); sys.path.insert(0, os.path.join(PKG, "code")) |
| import stasis_gpt2_eval as E |
| SHA = "5b2dc755ed6e38817be84aa982e4591f64c8cc94ce1fa0316c2ac0336e0bcf85" |
| def load(ckpt, device): |
| from frozen_stasis.config import ModelConfig |
| from frozen_stasis.model import StasisDecoder |
| E.need(E.file_sha(ckpt) == SHA, "step-36420 sha") |
| p = torch.load(ckpt, map_location="cpu", weights_only=True, mmap=True) |
| E.need(p["cursor"]["optimizer_step"] == 36420, "step 36420") |
| sd = p["model"]; m = StasisDecoder(ModelConfig()) |
| E.need(list(sd) == list(m.state_dict()), "keys"); m.load_state_dict(sd, strict=True) |
| E.need(sum(x.numel() for x in m.parameters()) == 109529856, "params") |
| return E.StasisWrap(m).to(device).eval() |
| ap = argparse.ArgumentParser(); ap.add_argument("--ckpt"); ap.add_argument("--out"); ap.add_argument("--limit", type=int) |
| ap.add_argument("--gpt2", action="store_true", help="also run GPT-2 on this device"); ap.add_argument("--threads", type=int, default=8) |
| a = ap.parse_args(); torch.set_num_threads(a.threads); os.makedirs(a.out, exist_ok=True) |
| t0 = time.time(); lf = open(os.path.join(a.out, "run_log.txt"), "a") |
| def log(m): |
| s = f"[{time.time()-t0:7.1f}s] {m}"; print(s, flush=True); lf.write(s + "\n"); lf.flush() |
| torch.manual_seed(20261016); sums = E.verify_package(PKG); log(f"package verified {len(sums)} files metric {E.METRIC_SHA[:16]}") |
| if a.gpt2: |
| g = E.load_gpt2(PKG, "cpu"); rg = E.run_model("GPT2_SMALL", g, PKG, "cpu", 8192, a.limit, a.out, log); del g |
| s = load(a.ckpt, "cpu"); log("Stasis step-36420 loaded strict weights_only") |
| rs = E.run_model("STASIS_12269", s, PKG, "cpu", 8192, a.limit, a.out, log) |
| for b in E.BENCHES: |
| os.replace(os.path.join(a.out, f"ROWS_{b}_STASIS_12269.jsonl"), os.path.join(a.out, f"ROWS_{b}_STASIS_36420.jsonl")) |
| log("done") |
|
|
