frame-A campaign: A1 front-fusion PASSED (0.87-0.91 both variants/seeds, bit-exact off), A2 consumption REFUTED (0.009/0.010 vs 0.15 line, controls 0.000/0.0003), A3 gated; all arms incl. controls shipped

proto_frame/campaign_a/campaign_a.py

@@ -0,0 +1,508 @@
+"""FRAME PROGRAM Phase A campaign (plan v2, Phil go: "run the campaign").
+
+Serialized stages on one GPU writer:
+  P   prep: probes (L6-first bridge, L6-last), teacher/self frame
+      targets, per-site token positions. From existing dumps.
+  A0a smoke: 1-seed 1500-step self-target run; gauge must move up.
+  A1  4 arms: (self|bert) x seeds(0,1). Bar: digit L6-first refit
+      0.28 -> >=0.75; refute <0.50; dead zone named. Side gauges +
+      bit-exact off + comparator rows cited.
+  A0b entry oracle: inject TRUE embedding sequences -> ceiling.
+  A2  4 arms: byte s0,s1 + random-codebook ctrl + shuffled-row ctrl
+      (trained, matched budget). Bar: unseen flip >=0.50; refute
+      <0.15; controls <=0.02.
+All adapters saved (weights-completeness: every arm incl. controls).
+Pure Adam wd=0. Env pinned in ledger.
+"""
+import json
+import os
+import sys
+import time
+import zlib
+
+sys.path.insert(0, r"E:\mirel\geolip-bytelex")
+sys.path.insert(0, r"E:\mirel\amoe-lora\src")
+sys.stdout.reconfigure(encoding="utf-8", errors="replace")
+import numpy as np
+import torch
+import torch.nn.functional as F
+import transformers
+from transformers import AutoTokenizer, T5EncoderModel
+import transformers.utils.logging as hlog
+
+hlog.set_verbosity_error()
+from amoe.core.adapter import AdapterSpec, RelayPatchwork, BlockWithAdapter
+from geolip.bytelex.frame import (apply_whitening, fit_whitening,
+                                  procrustes, split_sites,
+                                  top_k_accuracy)
+
+D = r"E:\mirel\data\bytelex\proto_frame"
+OUT = rf"{D}\campaign_a"
+os.makedirs(OUT, exist_ok=True)
+DEV = "cuda"
+SEED = zlib.crc32(b"frame-proto-v01") & 0xFFFFFFFF
+SPEC = AdapterSpec(n_slots=16, K=64, D=4, hidden=178)
+L_T5 = 6
+STEPS, SMOKE_STEPS, BS_LINES = 6000, 1500, 32
+torch.manual_seed(SEED)
+
+dump = np.load(rf"{D}\frame_dump_v2.npz")
+sweep = np.load(rf"{D}\sweep_dump_v02.npz")
+anch = np.load(rf"{D}\frame_anchors_v01.npz")
+states = json.load(open(r"E:\mirel\data\bytelex\words_of_C.json",
+                        encoding="utf-8"))
+walk = json.load(open(rf"{D}\t5_walk_of_C.json", encoding="utf-8"))
+CLS = np.array(["digit" if s["text"].isdigit() else
+                ("Name" if s["text"][0].isupper() else "word")
+                for s in states])
+sid = dump["sid"].astype(np.int64)
+KK = dump["KK"]
+ok = KK[:, 0] > 0
+sp = split_sites(sid[ok], seed=SEED)
+gix = {k: np.flatnonzero(ok)[v] for k, v in sp.items()}
+site_cls = CLS[sid]
+lines = open(r"E:\mirel\data\bytelex\codex_v1.txt",
+             "rb").read().decode("ascii").split("\n")
+tkA = AutoTokenizer.from_pretrained("google/flan-t5-small")
+E_Cn = F.normalize(torch.tensor(anch["E_C"], dtype=torch.float32,
+                                device=DEV), dim=-1)
+b_prior = torch.tensor(anch["b_prior"], dtype=torch.float32,
+                       device=DEV)
+sT = torch.tensor(float(anch["s"]), device=DEV).clamp(1, 100)
+LT = sweep["LT"].tolist()
+J6 = LT.index(6)
+
+# ---------------- P: prep -------------------------------------------
+def train_probe(H, tag):
+    mu, w = fit_whitening(H[gix["fit"]].astype(np.float64))
+    Z = apply_whitening(H.astype(np.float64), mu, w)
+    m = np.zeros((999, Z.shape[1]))
+    for s in range(999):
+        r = gix["train"][sid[gix["train"]] == s]
+        if len(r):
+            m[s] = Z[r].mean(0)
+    W = torch.nn.Parameter(torch.tensor(
+        procrustes(m, anch["E_C"].astype(np.float64)),
+        dtype=torch.float32, device=DEV))
+    tZ = torch.tensor(Z, dtype=torch.float32, device=DEV)
+    tsid_ = torch.tensor(sid, device=DEV)
+    opt = torch.optim.Adam([W], lr=1e-3, weight_decay=0.0)
+    rng = np.random.default_rng(SEED + 7)
+    for ep in range(6):
+        order = rng.permutation(gix["train"])
+        for i in range(0, len(order), 512):
+            b = torch.tensor(order[i:i + 512], device=DEV)
+            opt.zero_grad(set_to_none=True)
+            zf = F.normalize(tZ[b] @ W, dim=-1)
+            F.cross_entropy(sT * (zf @ E_Cn.T), tsid_[b]).backward()
+            opt.step()
+    print(f"[P] probe {tag} trained", flush=True)
+    return mu, w, W.detach()
+
+
+PREP = rf"{OUT}\prep.pt"
+if not os.path.exists(PREP):
+    muF, wF, W_first = train_probe(sweep["ST"][:, J6, 0], "L6-first")
+    muL, wL, W_last = train_probe(sweep["ST"][:, J6, 1], "L6-last")
+    # teacher frame vectors (bert L8 first through v0.1 maps)
+    H_B = dump["H_B"][:, 0].astype(np.float64)
+    muB, wB = fit_whitening(H_B[gix["fit"]])
+    ZB = apply_whitening(H_B, muB, wB)
+    zT = F.normalize(torch.tensor(ZB, dtype=torch.float32)
+                     @ torch.tensor(anch["W_T"], dtype=torch.float32),
+                     dim=-1)
+    ZL = apply_whitening(sweep["ST"][:, J6, 1].astype(np.float64),
+                         muL, wL)
+    zSelf = F.normalize(torch.tensor(ZL, dtype=torch.float32)
+                        @ W_last.cpu(), dim=-1)
+    # per-site t5 token anchor positions
+    pos = np.full((len(sid), 2), -1, dtype=np.int32)   # first_ix, kT
+    by_line = {}
+    for k in range(len(sid)):
+        by_line.setdefault(int(dump["line"][k]), []).append(k)
+    for li, ks in by_line.items():
+        e = tkA(lines[li], add_special_tokens=False,
+                return_offsets_mapping=True)
+        off = e["offset_mapping"]
+        for k in ks:
+            lo, hi = int(dump["lo"][k]), int(dump["hi"][k])
+            ix = [i for i, (s, t) in enumerate(off)
+                  if t > s and s < hi and t > lo]
+            if ix:
+                pos[k] = (ix[0], len(ix))
+    torch.save({"muF": muF, "wF": wF, "W_first": W_first.cpu(),
+                "muL": muL, "wL": wL, "W_last": W_last.cpu(),
+                "zT": zT, "zSelf": zSelf, "pos": pos}, PREP)
+    print("[P] prep saved", flush=True)
+prep = torch.load(PREP, weights_only=False)
+pos = prep["pos"]
+W_first_frozen = prep["W_first"].to(DEV)
+tmuF = torch.tensor(prep["muF"], dtype=torch.float32, device=DEV)
+twF = torch.tensor(prep["wF"], dtype=torch.float32, device=DEV)
+zT_all = prep["zT"].to(DEV)
+zSelf_all = prep["zSelf"].to(DEV)
+
+by_line_train = {}
+for k in gix["train"]:
+    if pos[k, 0] >= 0:
+        by_line_train.setdefault(int(dump["line"][k]), []).append(int(k))
+train_lines = sorted(by_line_train)
+
+
+def fresh_adapted():
+    m = T5EncoderModel.from_pretrained("google/flan-t5-small").to(DEV)
+    m.eval()
+    for p in m.parameters():
+        p.requires_grad_(False)
+    wraps = []
+    for i, blk in enumerate(m.encoder.block):
+        w = BlockWithAdapter(blk, RelayPatchwork(512, SPEC).to(DEV))
+        m.encoder.block[i] = w
+        wraps.append(w)
+    params = [p for w in wraps for p in w.adapter.parameters()]
+    return m, wraps, params
+
+
+def frame_first(h):
+    return F.normalize(((h - tmuF) @ twF) @ W_first_frozen, dim=-1)
+
+
+def batch_lines(chosen, inject=None):
+    """Tokenize chosen lines; returns ids/att + per-site (row, pos)."""
+    seqs = [tkA(lines[li], add_special_tokens=False)["input_ids"]
+            for li in chosen]
+    mx = max(len(s) for s in seqs)
+    ids = torch.full((len(seqs), mx), tkA.pad_token_id,
+                     dtype=torch.long)
+    att = torch.zeros((len(seqs), mx), dtype=torch.long)
+    for j, s in enumerate(seqs):
+        ids[j, :len(s)] = torch.tensor(s)
+        att[j, :len(s)] = 1
+    return ids.to(DEV), att.to(DEV)
+
+
+def train_a1(variant, seed, steps, tag):
+    torch.manual_seed(SEED + seed)
+    m, wraps, params = fresh_adapted()
+    tgt = zSelf_all if variant == "self" else zT_all
+    opt = torch.optim.Adam(params, lr=1e-3, weight_decay=0.0)
+    rng = np.random.default_rng(SEED + 100 + seed)
+    t0 = time.time()
+    for st in range(1, steps + 1):
+        chosen = [train_lines[int(i)] for i in
+                  rng.integers(0, len(train_lines), BS_LINES)]
+        ids, att = batch_lines(chosen)
+        h = m(input_ids=ids, attention_mask=att,
+              output_hidden_states=True).hidden_states[L_T5]
+        rows, cols, ks = [], [], []
+        for j, li in enumerate(chosen):
+            for k in by_line_train[li]:
+                rows.append(j)
+                cols.append(int(pos[k, 0]))
+                ks.append(k)
+        z = frame_first(h[rows, cols])
+        kt = torch.tensor(ks, device=DEV)
+        lid = F.cross_entropy(sT * (z @ E_Cn.T) + b_prior,
+                              torch.tensor(sid[ks], device=DEV))
+        lmse = ((z - tgt[kt]) ** 2).sum(-1).mean()
+        loss = lid + lmse
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        opt.step()
+        if st % 500 == 0:
+            print(f"[{tag}] step {st} id={float(lid):.4f} "
+                  f"mse={float(lmse):.4f}", flush=True)
+    print(f"[{tag}] trained in {time.time()-t0:.0f}s", flush=True)
+    return m, wraps
+
+
+@torch.no_grad()
+def extract_L6(m, readout=0):
+    H = np.zeros((len(sid), 512), dtype=np.float32)
+    by_line = {}
+    for k in range(len(sid)):
+        if pos[k, 0] >= 0:
+            by_line.setdefault(int(dump["line"][k]), []).append(k)
+    lis = sorted(by_line)
+    for bs in range(0, len(lis), 128):
+        chosen = lis[bs:bs + 128]
+        ids, att = batch_lines(chosen)
+        h = m(input_ids=ids, attention_mask=att,
+              output_hidden_states=True).hidden_states[L_T5].cpu()
+        for j, li in enumerate(chosen):
+            for k in by_line[li]:
+                p = int(pos[k, 0]) if readout == 0 else \
+                    int(pos[k, 0] + pos[k, 1] - 1)
+                H[k] = h[j, p].numpy()
+    return H
+
+
+def refit_gauge(H, tag):
+    """The pinned eval protocol: refit probe on (possibly adapted)
+    states, same splits/epochs as the v0.2 sweep."""
+    mu, w = fit_whitening(H[gix["fit"]].astype(np.float64))
+    Z = apply_whitening(H.astype(np.float64), mu, w)
+    m_ = np.zeros((999, 512))
+    for s in range(999):
+        r = gix["train"][sid[gix["train"]] == s]
+        if len(r):
+            m_[s] = Z[r].mean(0)
+    W = torch.nn.Parameter(torch.tensor(
+        procrustes(m_, anch["E_C"].astype(np.float64)),
+        dtype=torch.float32, device=DEV))
+    tZ = torch.tensor(Z, dtype=torch.float32, device=DEV)
+    opt = torch.optim.Adam([W], lr=1e-3, weight_decay=0.0)
+    rng = np.random.default_rng(SEED + 7)
+    for ep in range(6):
+        order = rng.permutation(gix["train"])
+        for i in range(0, len(order), 512):
+            b = torch.tensor(order[i:i + 512], device=DEV)
+            opt.zero_grad(set_to_none=True)
+            zf = F.normalize(tZ[b] @ W, dim=-1)
+            F.cross_entropy(sT * (zf @ E_Cn.T),
+                            torch.tensor(sid, device=DEV)[b]).backward()
+            opt.step()
+    ev = gix["eval"]
+    with torch.no_grad():
+        zf = F.normalize(tZ[torch.tensor(ev, device=DEV)] @ W, -1)
+        lg = (sT * zf @ E_Cn.T).cpu().numpy()
+    out = {"overall": round(top_k_accuracy(lg, sid[ev]), 4)}
+    for c in ("word", "Name", "digit"):
+        mm = site_cls[ev] == c
+        out[c] = round(top_k_accuracy(lg[mm], sid[ev][mm]), 4)
+    print(f"[gauge {tag}] {json.dumps(out)}", flush=True)
+    return out
+
+
+def save_arm(wraps, name):
+    st = {}
+    for i, w in enumerate(wraps):
+        for k, v in w.adapter.state_dict().items():
+            st[f"{i}.{k}"] = v.detach().cpu()
+    torch.save(st, rf"{OUT}\{name}.adapters.pt")
+
+
+PART = rf"{OUT}\ledger_partial.json"
+LED = (json.load(open(PART, encoding="utf-8"))
+       if os.path.exists(PART) else {})
+if LED:
+    print(f"[resume] partial ledger: {sorted(LED)}", flush=True)
+LED.setdefault("_env", {"transformers": transformers.__version__,
+                        "torch": torch.__version__, "seed": int(SEED),
+                        "spec": "n16 K64 D4 h178", "steps": STEPS})
+LED.setdefault("parity_unadapted_L6_first", {"digit": 0.2803,
+                                             "overall": 0.9442})
+LED.setdefault("comparator_zero_param", {"L6_last_digit": 0.8298,
+                                         "L2_concat_digit": 0.9537})
+
+# ---------------- A0a smoke ----------------------------------------
+if "A0a_smoke" in LED:
+    print("=== A0a SMOKE (cached, skip) ===", flush=True)
+else:
+    print("=== A0a SMOKE ===", flush=True)
+    m, wraps = train_a1("self", 99, SMOKE_STEPS, "A0a")
+    H = extract_L6(m)
+    g = refit_gauge(H, "A0a-smoke")
+    LED["A0a_smoke"] = g
+    del m
+    torch.cuda.empty_cache()
+    assert g["digit"] > 0.33, f"A0a GATE FAILED: digit {g['digit']}"
+    print(f"[A0a] GATE PASS (digit {g['digit']})", flush=True)
+    with open(rf"{OUT}\ledger_partial.json", "w") as f:
+        json.dump(LED, f, indent=1)
+
+# ---------------- A1 arms ------------------------------------------
+for variant in ("self", "bert"):
+    for seed in (0, 1):
+        tag = f"A1-{variant}-s{seed}"
+        if tag in LED and os.path.exists(
+                rf"{OUT}\{tag}.adapters.pt"):
+            print(f"=== {tag} (cached, skip) ===", flush=True)
+            continue
+        print(f"=== {tag} ===", flush=True)
+        m, wraps = train_a1(variant, seed, STEPS, tag)
+        H = extract_L6(m)
+        g = refit_gauge(H, tag)
+        # side: bit-exact off
+        for w in wraps:
+            w.enabled = False
+        ids, att = batch_lines(train_lines[:8])
+        with torch.no_grad():
+            h_off = m(input_ids=ids, attention_mask=att,
+                      output_hidden_states=True).hidden_states[L_T5]
+        clean = T5EncoderModel.from_pretrained(
+            "google/flan-t5-small").to(DEV).eval()
+        with torch.no_grad():
+            h_cl = clean(input_ids=ids, attention_mask=att,
+                         output_hidden_states=True).hidden_states[L_T5]
+        bit = bool(torch.equal(h_off, h_cl))
+        del clean
+        for w in wraps:
+            w.enabled = True
+        LED[tag] = {"gauge": g, "bitexact_off": bit}
+        save_arm(wraps, tag)
+        del m
+        torch.cuda.empty_cache()
+        with open(rf"{OUT}\ledger_partial.json", "w") as f:
+            json.dump(LED, f, indent=1)
+
+# ---------------- A0b entry oracle + A2 ----------------------------
+print("=== A0b/A2 ===", flush=True)
+kk_w = np.array([w["k"] for w in walk])
+ids_of = [w["ids"] for w in walk]
+rngw = np.random.default_rng(SEED + 41)
+perm = rngw.permutation(999)
+sfit = set(int(x) for x in perm[:599])
+sprobe = [int(x) for x in perm[599:]]
+probe_by_k = {}
+for u in sprobe:
+    probe_by_k.setdefault(int(kk_w[u]), []).append(u)
+E_lift = torch.zeros(999, 512, device=DEV)
+E_lift[:, :256] = F.normalize(torch.tensor(
+    anch["E_C"], dtype=torch.float32, device=DEV), dim=-1)
+E_rand_l = torch.zeros(999, 512, device=DEV)
+E_rand_l[:, :256] = F.normalize(torch.randn(999, 256, device=DEV), -1)
+prm = torch.tensor(np.random.default_rng(SEED + 43).permutation(999),
+                   device=DEV)
+E_shuf_l = E_lift[prm].clone()
+
+ev_sites = [int(k) for k in gix["eval"]
+            if int(kk_w[sid[k]]) in probe_by_k
+            and len(probe_by_k[int(kk_w[sid[k]])]) > 1
+            and pos[k, 1] == kk_w[sid[k]] and pos[k, 0] >= 0]
+tr_sites = [int(k) for k in gix["train"]
+            if int(sid[k]) in sfit and pos[k, 0] >= 0
+            and pos[k, 1] == kk_w[sid[k]]]
+print(f"[A2] train sites {len(tr_sites)} probe sites {len(ev_sites)}",
+      flush=True)
+
+
+def inject_forward(m, ks, alt_map, table):
+    chosen = sorted({int(dump["line"][k]) for k in ks})
+    li_row = {li: j for j, li in enumerate(chosen)}
+    ids, att = batch_lines(chosen)
+    x = m.get_input_embeddings()(ids).clone()
+    slots = []
+    for k in ks:
+        j = li_row[int(dump["line"][k])]
+        u = alt_map.get(k, int(sid[k]))
+        p0, kt = int(pos[k, 0]), int(pos[k, 1])
+        x[j, p0:p0 + kt] = table[u]
+        slots.append((k, j, p0, u))
+    h = m(inputs_embeds=x, attention_mask=att,
+          output_hidden_states=True).hidden_states[L_T5]
+    return h, slots
+
+
+@torch.no_grad()
+def a2_probe(m, table, tag):
+    rngp = np.random.default_rng(SEED + 37)
+    flips = n = 0
+    for i in range(0, len(ev_sites), 64):
+        ks = ev_sites[i:i + 64]
+        alt = {}
+        for k in ks:
+            cand = probe_by_k[int(kk_w[sid[k]])]
+            up = cand[int(rngp.integers(0, len(cand)))]
+            while up == int(sid[k]):
+                up = cand[int(rngp.integers(0, len(cand)))]
+            alt[k] = up
+        h, slots = inject_forward(m, ks, alt, table)
+        z = frame_first(h[[j for _, j, _, _ in slots],
+                          [p for _, _, p, _ in slots]])
+        pred = (z @ E_Cn.T).argmax(-1).cpu().numpy()
+        flips += int((pred == np.array([u for *_, u in slots])).sum())
+        n += len(slots)
+    r = round(flips / max(n, 1), 4)
+    print(f"[A2-probe {tag}] flip={r} n={n}", flush=True)
+    return r
+
+
+# oracle: true embedding sequences of alt states
+emb0 = T5EncoderModel.from_pretrained("google/flan-t5-small").to(DEV)
+emb0.eval()
+true_seq = torch.zeros(999, 4, 512, device=DEV)
+with torch.no_grad():
+    for u in range(999):
+        for i, tid in enumerate(ids_of[u][:4]):
+            true_seq[u, i] = emb0.get_input_embeddings().weight[tid]
+
+
+@torch.no_grad()
+def oracle_probe():
+    rngp = np.random.default_rng(SEED + 37)
+    flips = n = 0
+    for i in range(0, len(ev_sites), 64):
+        ks = ev_sites[i:i + 64]
+        chosen = sorted({int(dump["line"][k]) for k in ks})
+        li_row = {li: j for j, li in enumerate(chosen)}
+        ids, att = batch_lines(chosen)
+        x = emb0.get_input_embeddings()(ids).clone()
+        slots = []
+        for k in ks:
+            cand = probe_by_k[int(kk_w[sid[k]])]
+            up = cand[int(rngp.integers(0, len(cand)))]
+            while up == int(sid[k]):
+                up = cand[int(rngp.integers(0, len(cand)))]
+            j = li_row[int(dump["line"][k])]
+            p0, kt = int(pos[k, 0]), int(pos[k, 1])
+            x[j, p0:p0 + kt] = true_seq[up, :kt]
+            slots.append((j, p0, up))
+        h = emb0(inputs_embeds=x, attention_mask=att,
+                 output_hidden_states=True).hidden_states[L_T5]
+        z = frame_first(h[[j for j, _, _ in slots],
+                          [p for _, p, _ in slots]])
+        pred = (z @ E_Cn.T).argmax(-1).cpu().numpy()
+        flips += int((pred == np.array([u for *_, u in slots])).sum())
+        n += len(slots)
+    return round(flips / max(n, 1), 4)
+
+
+orc = oracle_probe()
+LED["A0b_entry_oracle_true_emb"] = orc
+print(f"[A0b] oracle ceiling flip={orc}", flush=True)
+del emb0
+torch.cuda.empty_cache()
+
+
+def train_a2(table, seed, tag):
+    torch.manual_seed(SEED + seed)
+    m, wraps, params = fresh_adapted()
+    opt = torch.optim.Adam(params, lr=1e-3, weight_decay=0.0)
+    rng = np.random.default_rng(SEED + 200 + seed)
+    for st in range(1, STEPS + 1):
+        ks = [tr_sites[int(i)] for i in
+              rng.integers(0, len(tr_sites), 48)]
+        h, slots = inject_forward(m, ks, {}, table)
+        z = frame_first(h[[j for _, j, _, _ in slots],
+                          [p for _, _, p, _ in slots]])
+        tgt = torch.tensor([u for *_, u in slots], device=DEV)
+        loss = F.cross_entropy(sT * (z @ E_Cn.T) + b_prior, tgt)
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        opt.step()
+        if st % 500 == 0:
+            print(f"[{tag}] step {st} id={float(loss):.4f}", flush=True)
+    return m, wraps
+
+
+for tag, table, seed in (("A2-byte-s0", E_lift, 0),
+                         ("A2-byte-s1", E_lift, 1),
+                         ("A2-randctrl-s0", E_rand_l, 0),
+                         ("A2-shufctrl-s0", E_shuf_l, 0)):
+    if tag in LED and os.path.exists(rf"{OUT}\{tag}.adapters.pt"):
+        print(f"=== {tag} (cached, skip) ===", flush=True)
+        continue
+    print(f"=== {tag} ===", flush=True)
+    m, wraps = train_a2(table, seed, tag)
+    LED[tag] = {"flip_unseen": a2_probe(m, table, tag)}
+    save_arm(wraps, tag)
+    del m
+    torch.cuda.empty_cache()
+    with open(rf"{OUT}\ledger_partial.json", "w") as f:
+        json.dump(LED, f, indent=1)
+
+with open(rf"{OUT}\campaign_a_ledger.json", "w", encoding="utf-8") as f:
+    json.dump(LED, f, indent=1)
+print("[CAMPAIGN A] COMPLETE", flush=True)
+print(json.dumps(LED, indent=1), flush=True)