diff --git a/results-loop/probe_discount_prompt.pt b/results-loop/probe_discount_prompt.pt new file mode 100644 index 0000000..19de156 Binary files /dev/null and b/results-loop/probe_discount_prompt.pt differ diff --git a/results-loop/probe_gsm430.pt b/results-loop/probe_gsm430.pt new file mode 100644 index 0000000..d7a299d Binary files /dev/null and b/results-loop/probe_gsm430.pt differ diff --git a/results-loop/probe_gsm430_ff.pt b/results-loop/probe_gsm430_ff.pt new file mode 100644 index 0000000..6ee764c Binary files /dev/null and b/results-loop/probe_gsm430_ff.pt differ diff --git a/results-loop/probe_gsm430_prompt.pt b/results-loop/probe_gsm430_prompt.pt new file mode 100644 index 0000000..501a7af Binary files /dev/null and b/results-loop/probe_gsm430_prompt.pt differ diff --git a/results-loop/probe_gsm795.pt b/results-loop/probe_gsm795.pt new file mode 100644 index 0000000..013e2fb Binary files /dev/null and b/results-loop/probe_gsm795.pt differ diff --git a/results-loop/probe_gsm795_ff.pt b/results-loop/probe_gsm795_ff.pt new file mode 100644 index 0000000..0b2e38b Binary files /dev/null and b/results-loop/probe_gsm795_ff.pt differ diff --git a/results-loop/probe_gsm795_prompt.pt b/results-loop/probe_gsm795_prompt.pt new file mode 100644 index 0000000..cd0dd76 Binary files /dev/null and b/results-loop/probe_gsm795_prompt.pt differ diff --git a/scripts/probe_discount.py b/scripts/probe_discount.py index c355db6..67efe0a 100644 --- a/scripts/probe_discount.py +++ b/scripts/probe_discount.py @@ -27,6 +27,14 @@ 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, P = 2, 6 +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}" FF = "--ff" in sys.argv ADAPTER = OUT / ("adapter_carrycot_ff_e400.pt" if FF else "adapter_carrycot_e400.pt") @@ -45,7 +53,7 @@ mask = torch.ones_like(ids) with torch.no_grad(): gen = generate_carry_c(looper, adapter, tok, ids, mask, K, P, - max_new_tokens=48, feedforward=FF) + max_new_tokens=72, feedforward=FF) text = tok.decode(gen[0, ids.shape[1] + P:], skip_special_tokens=False) print("GENERATED:", repr(text), flush=True) @@ -111,5 +119,5 @@ for i, lab in enumerate(poslab): torch.save({"question": Q, "k": K, "p": P, "generated": text, "band": band, "poslab": poslab, "table": table, "digit_probs": probs}, - OUT / ("probe_discount_ff.pt" if FF else "probe_discount.pt")) -print("\nsaved", OUT / ("probe_discount_ff.pt" if FF else "probe_discount.pt")) + OUT / (f"probe_{TAG}_ff.pt" if FF else f"probe_{TAG}.pt")) +print("\nsaved", OUT / (f"probe_{TAG}_ff.pt" if FF else f"probe_{TAG}.pt")) diff --git a/scripts/probe_discount_prompt.py b/scripts/probe_discount_prompt.py new file mode 100644 index 0000000..a8f2712 --- /dev/null +++ b/scripts/probe_discount_prompt.py @@ -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)