File size: 8,445 Bytes
f930dac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | #!/usr/bin/env python3
"""E3 Stage A: can protected learned bindings persist at all?
Semantic tasks only (SEM_LATEST + SEM_2HOP), 4.39M, three seeds, two arms:
--mode bind exact-address writes, unrelated slots bitwise untouched
--mode unbound 192 slots, unstructured writes (E3-A3 capacity control)
Arms are PAIRED: same seed gives identical initial parameters, identical
lifetime and batch ordering, identical schedule and identical eval examples.
Gates are frozen in preregistration/PREREGISTRATION_E3_BIND.md and evaluated by
eval/stage_a_eval.py, not here.
"""
import argparse
import json
import sys
import time
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
sys.path.insert(0, str(Path(__file__).resolve().parent))
from pns.common import ckpt_root, logs_root, shards_root # noqa: E402
from pns.model.bind import BindConfig, PNSBind # noqa: E402
from pns.model.modules import enum_legal_mask # noqa: E402
from pns.train.loader import LifetimeBatcher, iter_eval_batches # noqa: E402
from pns.world.schema import ET, Fam # noqa: E402
SEM = (int(Fam.SEM_LATEST), int(Fam.SEM_2HOP))
TB = 32
def to_dev(b, dev):
return {k: torch.from_numpy(np.ascontiguousarray(
v.view(np.int64) if v.dtype == np.uint64 else v.astype(np.int64))).to(dev)
for k, v in b.items()}
@torch.no_grad()
def evaluate(model, dev, split, n_lifetimes, intervention="none"):
"""Returns rows: (family, ok, delay, n_unrelated_writes_since_evidence)."""
model.eval()
rows = []
for b in iter_eval_batches(split, shards_root(), 24, n_lifetimes):
g = to_dev(b, dev)
B, L = g["etype"].shape
state = model.initial_state(B, dev)
if intervention == "payload_zero":
state = torch.zeros_like(state)
swap_at = L // 2
# count unrelated binding writes per event, for the interference split
wcount = torch.zeros(B, device=dev)
wcount_hist = torch.zeros(B, L, device=dev)
with torch.autocast("cuda", dtype=torch.bfloat16):
for t in range(L):
if intervention == "payload_swap" and t == swap_at:
state = state.roll(1, dims=0)
state, out = model.step(
state, g["tok"][:, t], g["etype"][:, t], g["dt"][:, t],
g["bind_write"][:, t], g["bind_read"][:, t],
g["bind_slot_ent"], g["bind_slot_attr"],
freeze_writes=(intervention == "payload_zero"))
wcount = wcount + (g["bind_write"][:, t] >= 0).float()
wcount_hist[:, t] = wcount
sel = torch.isin(g["family"][:, t],
torch.tensor(SEM, device=dev))
if sel.any():
legal = enum_legal_mask(g["enum_legal"][sel, t])
pred = out["enum"][sel].masked_fill(~legal, -1e9).argmax(-1)
ok = (pred == g["enum_gold"][sel, t]).float().cpu().numpy()
idx = torch.nonzero(sel).flatten()
for j, i in enumerate(idx.tolist()):
d = int(g["delay"][i, t])
ev_t = max(0, t - d)
n_unrel = float(wcount_hist[i, t] - wcount_hist[i, ev_t])
rows.append((int(g["family"][i, t]), float(ok[j]), d, n_unrel))
model.train()
return rows
def summarise(rows):
a = np.array([r[1] for r in rows]) if rows else np.zeros(0)
d = np.array([r[2] for r in rows]) if rows else np.zeros(0)
u = np.array([r[3] for r in rows]) if rows else np.zeros(0)
fam = np.array([r[0] for r in rows]) if rows else np.zeros(0)
out = {"n": len(rows), "acc": float(a.mean()) if len(a) else float("nan")}
for lo, hi, name in ((17, 10**9, "gt16"), (129, 10**9, "gt128"),
(1, 16, "le16")):
m = (d >= lo) & (d <= hi)
if m.sum() > 20:
out[f"acc_{name}"] = round(float(a[m].mean()), 4)
out[f"n_{name}"] = int(m.sum())
for f, nm in ((int(Fam.SEM_LATEST), "latest"), (int(Fam.SEM_2HOP), "2hop")):
m = fam == f
if m.sum() > 20:
out[f"acc_{nm}"] = round(float(a[m].mean()), 4)
mm = m & (d > 16)
if mm.sum() > 20:
out[f"acc_{nm}_gt16"] = round(float(a[mm].mean()), 4)
# interference split: accuracy vs unrelated writes since the evidence
if len(u):
q = np.quantile(u, [0.33, 0.66])
for i, (lo, hi) in enumerate([(-1, q[0]), (q[0], q[1]), (q[1], 1e9)]):
m = (u > lo) & (u <= hi)
if m.sum() > 20:
out[f"acc_unrel_t{i}"] = round(float(a[m].mean()), 4)
out[f"unrel_t{i}_mean"] = round(float(u[m].mean()), 1)
return out
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--mode", choices=["bind", "unbound"], required=True)
ap.add_argument("--run", required=True)
ap.add_argument("--seed", type=int, default=1)
ap.add_argument("--updates", type=int, default=2600)
ap.add_argument("--batch", type=int, default=128)
ap.add_argument("--lr", type=float, default=3e-4)
ap.add_argument("--eval-lifetimes", type=int, default=400)
args = ap.parse_args()
dev = "cuda"
torch.manual_seed(args.seed) # identical init across arms
model = PNSBind(BindConfig(mode=args.mode)).to(dev)
opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01)
rd = ckpt_root() / args.run
rd.mkdir(parents=True, exist_ok=True)
log = open(logs_root() / f"{args.run}.jsonl", "a")
print(f"{args.run}: mode={args.mode} seed={args.seed} "
f"params={sum(p.numel() for p in model.parameters())/1e6:.2f}M", flush=True)
batcher = LifetimeBatcher("e3_train", shards_root(), args.batch,
seed=args.seed) # identical ordering across arms
u, t0 = 0, time.time()
for lt in batcher:
if u >= args.updates:
break
g = to_dev(lt, dev)
B, L = g["etype"].shape
state = model.initial_state(B, dev)
for c0 in range(0, L, TB):
state = state.detach()
opt.zero_grad(set_to_none=True)
losses = []
with torch.autocast("cuda", dtype=torch.bfloat16):
for t in range(c0, min(c0 + TB, L)):
state, out = model.step(
state, g["tok"][:, t], g["etype"][:, t], g["dt"][:, t],
g["bind_write"][:, t], g["bind_read"][:, t],
g["bind_slot_ent"], g["bind_slot_attr"])
sel = torch.isin(g["family"][:, t], torch.tensor(SEM, device=dev))
if sel.any():
legal = enum_legal_mask(g["enum_legal"][sel, t])
lg = out["enum"][sel].masked_fill(~legal, -1e9)
losses.append(F.cross_entropy(lg, g["enum_gold"][sel, t]))
if losses:
loss = torch.stack(losses).mean()
loss.backward()
gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
if torch.isfinite(loss):
opt.step()
u += 1
if u % 200 == 0:
s = {"u": u, "loss": round(float(loss), 4), "gn": round(float(gn), 2),
"s": round(time.time() - t0)}
log.write(json.dumps(s) + "\n"); log.flush()
print(json.dumps(s), flush=True)
if u >= args.updates:
break
torch.save({"model": model.state_dict(), "cfg": vars(model.cfg),
"args": vars(args)}, rd / "final.pt")
res = {"run": args.run, "mode": args.mode, "seed": args.seed}
for iv in ("none", "payload_swap", "payload_zero"):
res[iv] = summarise(evaluate(model, dev, "e3_dev", args.eval_lifetimes, iv))
print(iv, json.dumps(res[iv]), flush=True)
res["delta_payload_swap"] = round(res["none"]["acc"] - res["payload_swap"]["acc"], 4)
res["delta_payload_zero"] = round(res["none"]["acc"] - res["payload_zero"]["acc"], 4)
from pns.common import atomic_write_json, eval_root
atomic_write_json(eval_root() / f"STAGEA_{args.run}.json", res)
print("delta_swap", res["delta_payload_swap"],
"delta_zero", res["delta_payload_zero"], flush=True)
if __name__ == "__main__":
main()
|