| """Stasis 1 step 12,269 vs GPT-2 small matched eval. Implements METRIC_DEFINITIONS.json v1 |
| (sha256 579ab8c12aeb2cb01ed6f8d32028f97065139217f1a7e702bf0d3ff1afa24137). Offline; no hub, no AutoModel.""" |
| import hashlib, json, math, os, sys, time |
| import numpy as np, torch |
| import torch.nn.functional as F |
|
|
| METRIC_SHA = "579ab8c12aeb2cb01ed6f8d32028f97065139217f1a7e702bf0d3ff1afa24137" |
| CKPT_SHAS = {"2139ca7dd8b0c21e7c4b3adce14970fceadb7e542aa61c222e375b5a13b2656b": "endpoint.model_only.pt", |
| "44e66daf5e0eee079711ac96d2bdd10fff7751a0573d56eb5c7b422f958e9645": "recovery-final.pt"} |
| GPT2_SAFETENSORS_SHA = None |
| BENCHES = ["WT2_FULL_BPB", "WT2_PANEL256_CONTINUITY", "LAMBADA_OPENAI", "HELLASWAG", "PIQA", "ARC_EASY"] |
| LN2 = math.log(2) |
|
|
| def need(ok, msg): |
| if not ok: raise RuntimeError("GATE FAILED: " + msg) |
|
|
| def file_sha(p): |
| h = hashlib.sha256() |
| with open(p, "rb") as f: |
| for b in iter(lambda: f.read(1 << 24), b""): h.update(b) |
| return h.hexdigest() |
|
|
| def verify_package(root): |
| need(file_sha(os.path.join(root, "METRIC_DEFINITIONS.json")) == METRIC_SHA, "metric definitions sha (frozen)") |
| sums = {} |
| for line in open(os.path.join(root, "SHA256SUMS")): |
| s, p = line.strip().split(" ", 1); sums[p] = s |
| need(file_sha(os.path.join(root, p)) == s, "package file sha " + p) |
| return sums |
|
|
| |
| class GPT2(torch.nn.Module): |
| """Minimal exact GPT-2 small forward (gelu_new, LN eps 1e-5, tied head) over bundled safetensors.""" |
| def __init__(self, sd): |
| super().__init__() |
| self.p = torch.nn.ParameterDict({k.replace(".", "__"): torch.nn.Parameter(v, requires_grad=False) for k, v in sd.items()}) |
| def g(self, k): return self.p[k.replace(".", "__")] |
| def forward(self, ids): |
| b, t = ids.shape |
| x = self.g("wte.weight")[ids] + self.g("wpe.weight")[:t] |
| for i in range(12): |
| h = F.layer_norm(x, (768,), self.g(f"h.{i}.ln_1.weight"), self.g(f"h.{i}.ln_1.bias"), 1e-5) |
| qkv = h @ self.g(f"h.{i}.attn.c_attn.weight") + self.g(f"h.{i}.attn.c_attn.bias") |
| q, k, v = qkv.view(b, t, 3, 12, 64).permute(2, 0, 3, 1, 4).unbind(0) |
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True).transpose(1, 2).reshape(b, t, 768) |
| x = x + y @ self.g(f"h.{i}.attn.c_proj.weight") + self.g(f"h.{i}.attn.c_proj.bias") |
| h = F.layer_norm(x, (768,), self.g(f"h.{i}.ln_2.weight"), self.g(f"h.{i}.ln_2.bias"), 1e-5) |
| h = F.gelu(h @ self.g(f"h.{i}.mlp.c_fc.weight") + self.g(f"h.{i}.mlp.c_fc.bias"), approximate="tanh") |
| x = x + h @ self.g(f"h.{i}.mlp.c_proj.weight") + self.g(f"h.{i}.mlp.c_proj.bias") |
| x = F.layer_norm(x, (768,), self.g("ln_f.weight"), self.g("ln_f.bias"), 1e-5) |
| return x @ self.g("wte.weight").T |
|
|
| def read_safetensors(path): |
| """Dependency-free safetensors reader (F32 only), so no safetensors/transformers install is needed.""" |
| import struct |
| raw = open(path, "rb").read(); n = struct.unpack("<Q", raw[:8])[0]; hdr = json.loads(raw[8:8 + n]); base = 8 + n; out = {} |
| for k, v in hdr.items(): |
| if k == "__metadata__": continue |
| need(v["dtype"] == "F32", "F32 tensor " + k); a, b = v["data_offsets"] |
| out[k] = torch.from_numpy(np.frombuffer(raw, dtype=np.float32, count=(b - a) // 4, offset=base + a).copy()).view(v["shape"]) |
| return out |
|
|
| def load_gpt2(root, device): |
| sd = read_safetensors(os.path.join(root, "gpt2", "model.safetensors")) |
| keys = {"wte.weight", "wpe.weight", "ln_f.weight", "ln_f.bias"} |
| for i in range(12): |
| for m in ["ln_1", "attn.c_attn", "attn.c_proj", "ln_2", "mlp.c_fc", "mlp.c_proj"]: |
| keys |= {f"h.{i}.{m}.weight", f"h.{i}.{m}.bias"} |
| masks = {k for k in sd if k.endswith(".attn.bias") or k.endswith(".attn.masked_bias")} |
| need(set(sd) - masks == keys, "exact GPT-2 tensor inventory (strict)") |
| sd = {k: sd[k].float() for k in keys} |
| need(sum(v.numel() for v in sd.values()) == 124439808, "GPT-2 param count") |
| need(all(bool(torch.isfinite(v).all()) for v in sd.values()), "finite GPT-2") |
| return GPT2(sd).to(device).eval() |
|
|
| class StasisWrap(torch.nn.Module): |
| def __init__(self, m): super().__init__(); self.m = m |
| def forward(self, ids): return self.m(ids)[0] |
|
|
| def load_stasis(root, ckpt_path, device, ckpt_sha=None): |
| sys.path.insert(0, os.path.join(root, "code")) |
| from frozen_stasis.config import ModelConfig |
| from frozen_stasis.model import StasisDecoder |
| sha = ckpt_sha or file_sha(ckpt_path) |
| need(sha in CKPT_SHAS, "step-12269 checkpoint sha256 (got %s)" % sha) |
| payload = torch.load(ckpt_path, map_location="cpu", weights_only=True, mmap=True) |
| need(payload["cursor"]["optimizer_step"] == 12269, "optimizer_step 12269") |
| sd = payload["model"] |
| m = StasisDecoder(ModelConfig()) |
| need(list(sd) == list(m.state_dict()), "Stasis key set/order") |
| need(all(v.dtype == torch.float32 and bool(torch.isfinite(v).all()) for v in sd.values()), "finite fp32 Stasis") |
| m.load_state_dict(sd, strict=True) |
| need(sum(p.numel() for p in m.parameters()) == 109529856, "Stasis param count") |
| del payload, sd |
| return StasisWrap(m).to(device).eval(), sha |
|
|
| |
| def score_seqs(model, seqs, device, tok_budget): |
| """seqs: list of (ctx_ids_with_bos, tgt_ids). Returns list of (target logprobs np.float64, greedy_all_match).""" |
| order = sorted(range(len(seqs)), key=lambda i: -(len(seqs[i][0]) + len(seqs[i][1]))) |
| out = [None] * len(seqs); i = 0 |
| while i < len(order): |
| L = len(seqs[order[i]][0]) + len(seqs[order[i]][1]) - 1 |
| need(L <= 1024, "sequence > 1024 tokens") |
| nb = max(1, tok_budget // L); batch = order[i:i + nb]; i += nb |
| x = torch.zeros((len(batch), L), dtype=torch.long) |
| for r, j in enumerate(batch): |
| full = seqs[j][0] + seqs[j][1]; x[r, :len(full) - 1] = torch.tensor(full[:-1]) |
| with torch.inference_mode(): |
| logits = model(x.to(device)) |
| for r, j in enumerate(batch): |
| c, t = seqs[j]; s = len(c) - 1 |
| lg = logits[r, s:s + len(t)].float() |
| need(bool(torch.isfinite(lg).all()), "finite logits") |
| lp = torch.log_softmax(lg, -1) |
| tt = torch.tensor(t, device=lg.device) |
| out[j] = (lp.gather(1, tt[:, None])[:, 0].double().cpu().numpy(), bool((lg.argmax(-1) == tt).all())) |
| del logits |
| return out |
|
|
| def load_items(root, bench, limit=None): |
| rows = [json.loads(l) for l in open(os.path.join(root, "data", bench + ".jsonl"))] |
| if limit: |
| if bench.startswith("WT2"): |
| units = list(dict.fromkeys(r["unit"] for r in rows))[:limit]; rows = [r for r in rows if r["unit"] in set(units)] |
| else: rows = rows[:limit] |
| return rows |
|
|
| def run_model(name, model, root, device, tok_budget, limit, outdir, log): |
| bos = 2 if name == "STASIS_12269" else 50256 |
| res = {} |
| for bench in BENCHES: |
| t0 = time.time(); rows = load_items(root, bench, limit) |
| if bench in ("HELLASWAG", "PIQA", "ARC_EASY"): |
| seqs = [([bos] + c["ids"][name]["ctx"], c["ids"][name]["tgt"]) for r in rows for c in r["choices"]] |
| sc = score_seqs(model, seqs, device, tok_budget); k = 0; per = [] |
| for r in rows: |
| lls, nb = [], [] |
| for c in r["choices"]: lls.append(float(sc[k][0].sum())); nb.append(c["tgt_bytes"]); k += 1 |
| a = int(np.argmax(lls)); an = int(np.argmax(np.array(lls) / np.array(nb))) |
| per.append({"unit": r["unit"], "gold": r["gold"], "ll": lls, "acc": int(a == r["gold"]), "acc_norm": int(an == r["gold"])}) |
| else: |
| live = [r for r in rows if r["ids"] is not None] |
| sc = score_seqs(model, [([bos] + r["ids"][name]["ctx"], r["ids"][name]["tgt"]) for r in live], device, tok_budget) |
| per = [] |
| for r, (lp, g) in zip(live, sc): |
| per.append({"unit": r["unit"], "start": r.get("start", 0), "bytes": r["tgt_bytes"], "nll": float(-lp.sum()), "tokens": len(lp), "greedy": int(g)}) |
| per += [{"unit": r["unit"], "start": 0, "bytes": 0, "nll": 0.0, "tokens": 0, "greedy": 0} for r in rows if r["ids"] is None] |
| with open(os.path.join(outdir, f"ROWS_{bench}_{name}.jsonl"), "w") as f: |
| for p in per: f.write(json.dumps(p) + "\n") |
| res[bench] = per; log(f"{name} {bench}: {len(rows)} items in {time.time()-t0:.1f}s") |
| return res |
|
|
| |
| def boot_idx(n, draws=10000, seed=20261016, chunk=500): |
| rng = np.random.default_rng(seed) |
| for s in range(0, draws, chunk): yield rng.integers(0, n, (min(chunk, draws - s), n)) |
|
|
| def bpb_units(per): |
| u = {} |
| for p in per: |
| a = u.setdefault(p["unit"], [0.0, 0]); a[0] += p["nll"]; a[1] += p["bytes"] |
| return u |
|
|
| def compare_bpb(S, G): |
| us, ug = bpb_units(S), bpb_units(G); keys = sorted(us) |
| need(keys == sorted(ug) and all(us[k][1] == ug[k][1] for k in keys), "paired identical bytes") |
| ns = np.array([us[k][0] for k in keys]); ng = np.array([ug[k][0] for k in keys]); b = np.array([us[k][1] for k in keys], float) |
| bs, bg = ns.sum() / (LN2 * b.sum()), ng.sum() / (LN2 * b.sum()) |
| d, rel = [], [] |
| for idx in boot_idx(len(keys)): |
| B = b[idx].sum(1); x = ns[idx].sum(1) / (LN2 * B); y = ng[idx].sum(1) / (LN2 * B) |
| d.append(x - y); rel.append(100 * (x / y - 1)) |
| d, rel = np.concatenate(d), np.concatenate(rel) |
| ci = [float(np.percentile(d, 2.5)), float(np.percentile(d, 97.5))] |
| return {"STASIS_12269_bpb": bs, "GPT2_SMALL_bpb": bg, "delta_bpb": bs - bg, "delta_ci95": ci, |
| "relative_deficit_pct": 100 * (bs / bg - 1), "relative_deficit_ci95": [float(np.percentile(rel, 2.5)), float(np.percentile(rel, 97.5))], |
| "bytes": int(b.sum()), "units": len(keys), "stasis_not_worse": ci[1] < 0} |
|
|
| def compare_acc(S, G, key): |
| need([p["unit"] for p in S] == [p["unit"] for p in G], "paired items") |
| a = np.array([p[key] for p in S], float); g = np.array([p[key] for p in G], float); dd = a - g |
| d = np.concatenate([dd[idx].mean(1) for idx in boot_idx(len(dd))]) |
| ci = [float(np.percentile(d, 2.5)), float(np.percentile(d, 97.5))] |
| return {"STASIS_12269": float(a.mean()), "GPT2_SMALL": float(g.mean()), "delta": float(dd.mean()), "delta_ci95": ci, "n": len(dd), "stasis_not_worse": ci[0] > 0} |
|
|
| def summarize(rs, rg): |
| out = {} |
| for bench in ["WT2_FULL_BPB", "WT2_PANEL256_CONTINUITY"]: out[bench] = {"bpb": compare_bpb(rs[bench], rg[bench])} |
| out["LAMBADA_OPENAI"] = {"acc": compare_acc(rs["LAMBADA_OPENAI"], rg["LAMBADA_OPENAI"], "greedy"), |
| "bpb": compare_bpb(rs["LAMBADA_OPENAI"], rg["LAMBADA_OPENAI"])} |
| for bench in ["HELLASWAG", "PIQA", "ARC_EASY"]: |
| out[bench] = {"acc_norm": compare_acc(rs[bench], rg[bench], "acc_norm"), "acc": compare_acc(rs[bench], rg[bench], "acc")} |
| return out |
|
|
