whiteboard microscopy on hard(795)/drop(430) specimens: FF slips at the 3-term sum (155), carry re-expands and defers (165 ✓); boards show live continuation-plan vs premature digit-commit at identical prefixes

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-16 19:48:58 +02:00
parent f143c67b3d
commit 46b5e1b802
9 changed files with 90 additions and 3 deletions
@@ -0,0 +1,79 @@
+"""Companion to probe_discount.py: lens table over the PROMPT positions
+after k=2 prefill settling (arm A adapter) — what the workspace holds
+about the question before any pause/answer compute begins."""
+import os
+import sys
+from pathlib import Path
+
+import torch
+
+from carry_common import _pos_ids, prompt_prefill
+from loop_common import BandLooper, MergeAdapter, chat_prompt, DIRECT_SUFFIX
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from jlens.core import JLens, ResidualCapture, load_model  # noqa: E402
+
+OUT = Path(os.environ.get("LOOP_OUT",
+    Path(__file__).resolve().parent.parent / "results-loop"))
+Q = "A shirt costs $80 and is on sale for 15% off. How much does it cost?"
+K = 2
+TAG = "discount"
+if "--idx" in sys.argv:
+    import json
+    _i = int(sys.argv[sys.argv.index("--idx") + 1])
+    _star = {it["idx"]: it for it in json.load(open(OUT / "star_data.json"))
+             if it["split"] == "test"}
+    Q = _star[_i]["question"]
+    TAG = f"gsm{_i}"
+
+model, tok = load_model(dtype=torch.bfloat16)
+looper = BandLooper(model)
+band = list(range(looper.l0, looper.l1 + 1))
+adapter = MergeAdapter(d=model.config.get_text_config().hidden_size).cuda()
+adapter.load_state_dict(torch.load(OUT / "adapter_carrycot_e400.pt",
+                                   map_location="cuda"))
+jbar = torch.load(Path(__file__).resolve().parent.parent / "results/jbar.pt",
+                  map_location="cuda")["Jbar"]
+jl = JLens(model, tok, jbar)
+
+prompt = chat_prompt(tok, Q, DIRECT_SUFFIX)
+ids = tok(prompt, add_special_tokens=False,
+          return_tensors="pt")["input_ids"].cuda()
+mask = torch.ones_like(ids)
+
+with torch.no_grad():
+    calls, _ = looper.capture(ids, mask, logits_to_keep=1,
+                              position_ids=_pos_ids(mask))
+    e = looper._hin[looper.l0]
+    pmask = torch.ones_like(ids, dtype=torch.bool)
+    S, X = prompt_prefill(looper, adapter, e, calls, pmask, K)
+    with ResidualCapture(model, layers=band) as rc:
+        looper.band(X, calls)
+    acts = {l: rc.acts[l][0] for l in band}
+    # k=0 reference: plain forward, no adapter anywhere
+    with ResidualCapture(model, layers=band) as rc0:
+        model(input_ids=ids, attention_mask=mask, use_cache=False)
+    acts0 = {l: rc0.acts[l][0] for l in band}
+
+table = {}
+for l in band:
+    idx, p = jl.read(acts[l], l, topk=5)
+    table[l] = (idx.cpu(), p.cpu())
+idx0, p0 = jl.read(acts0[30], 30, topk=5)
+toks = [repr(tok.decode([t])) for t in ids[0].tolist()]
+torch.save({"question": Q, "k": K, "poslab": toks, "band": band,
+            "table": table, "base30": (idx0.cpu(), p0.cpu())},
+           OUT / f"probe_{TAG}_prompt.pt")
+show = [14, 18, 22, 26, 30]
+print(f"=== prompt board after k={K} settle (arm A)  ||  k=0 base ===")
+for i, lab in enumerate(toks):
+    row, row0 = [], []
+    for l in show:
+        idx, _ = jl.read(acts[l][i][None], l, topk=3)
+        row.append("/".join(repr(tok.decode([t]))[1:-1]
+                            for t in idx[0].tolist()))
+    idx0, _ = jl.read(acts0[30][i][None], 30, topk=3)
+    base30 = "/".join(repr(tok.decode([t]))[1:-1] for t in idx0[0].tolist())
+    print(f"{i:3d} {lab:14s} " + "  ".join(
+        f"L{l}:{r}" for l, r in zip(show, row)) + f"   || base-L30:{base30}",
+        flush=True)