diff --git a/results-loop/PROTOCOL_UNIFIED.md b/results-loop/PROTOCOL_UNIFIED.md index 7cd5118..a0169b2 100644 --- a/results-loop/PROTOCOL_UNIFIED.md +++ b/results-loop/PROTOCOL_UNIFIED.md @@ -359,3 +359,21 @@ on GSM alone in the correct regime, delivers no usable gain and the standard fidelity damage; combined with the scope note, the boundary claim (structural: supervision density + state-evolution bottleneck) is fully supported. + +18. **E1: learned per-prompt halting gate (pre-registered 2026-07-16 + ~00:20, before any arm runs; PLAN_SELFPACED.md).** HaltingMergeAdapter: + frozen-recipe merge + ACT-style halting head on the last prompt + position's workspace state; soft state-mixture training, CE + lambda * + E[iters], penalty warmup at step 100; NO difficulty curriculum (mixed + batches โ€” the gate must discover the allocation). k_max=4, e400/e600 + checkpoints, deploy = sequential halting at 0.5 cumulative mass, + generation via frozen-prompt at per-item k*. Arms: lambda in + {0, 1e-3, 1e-2}, seed 0. Eval: 250 items, vs anchors k=0 (0.488), + uniform merge k=4 (0.512/0.885/0.464), probe-gate E0 (0.520/0.975/0.286). + Predictions: (a) some lambda gives overall >= 0.512 at mean E[k] <= + 2.4 (60% of uniform-4); (b) easy >= 0.95 at that lambda; (c) k*-vs-hard + point-biserial r > 0.3; (d) hard >= 0.286 (beats E0's frozen probe). + Collapse (E[k] pinned at 1 or 4 for all lambda) falsifies E1 and + triggers the plan's kill criterion. lambda=0 control isolates whether + the CE gradient alone moves the gate (expected: barely โ€” penalty + provides the pressure). diff --git a/scripts/eval_gate_code.py b/scripts/eval_gate_code.py new file mode 100644 index 0000000..5195a35 --- /dev/null +++ b/scripts/eval_gate_code.py @@ -0,0 +1,114 @@ +"""E1 eval (PLAN_SELFPACED): halting-gate MBPP eval. + +Per item: deploy-time halting picks k* (1..kmax); generation then uses the +frozen-prompt path at that k (items grouped by k* for batching). Reports +pass@1 by label, the k* distribution by label, mean compute, and the +gate-difficulty point-biserial correlation. +""" + +import argparse +import json +import os +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import torch + +from halting_common import HaltingMergeAdapter, halted_k_per_item +from loop_common import BandLooper +from prep_mbpp import DIRECT_SUFFIX, extract_code, mbpp_prompt, run_tests + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--adapter", required=True) + ap.add_argument("--tag", required=True) + ap.add_argument("--kmax", type=int, default=4) + ap.add_argument("--n", type=int, default=250) + ap.add_argument("--batch", type=int, default=8) + args = ap.parse_args() + + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + looper = BandLooper(model) + adapter = HaltingMergeAdapter( + d=model.config.get_text_config().hidden_size).cuda() + adapter.load_state_dict(torch.load(args.adapter, map_location="cuda")) + adapter.eval() + + items = [it for it in json.load(open(OUT / "mbpp_data.json")) + if it["split"] == "test"][: args.n] + print(f"[{args.tag}] gated MBPP eval on {len(items)}, kmax={args.kmax}", + flush=True) + + # phase 1: per-item k* + kstars = [] + with torch.no_grad(): + for i in range(0, len(items), args.batch): + chunk = items[i : i + args.batch] + enc = tok([mbpp_prompt(tok, it, DIRECT_SUFFIX) for it in chunk], + return_tensors="pt", padding=True, + add_special_tokens=False).to("cuda") + pl = enc["attention_mask"].sum(-1) - 1 # left-pad: last position + pl = torch.full_like(pl, enc["input_ids"].shape[1] - 1) + ks = halted_k_per_item(looper, adapter, enc["input_ids"], + args.kmax, enc["attention_mask"], pl) + kstars.extend(ks.tolist()) + t0 = time.time() + + # phase 2: generate grouped by k* + codes = [None] * len(items) + for kval in sorted(set(kstars)): + idxs = [i for i, kk in enumerate(kstars) if kk == kval] + for j in range(0, len(idxs), args.batch): + grp = idxs[j : j + args.batch] + enc = tok([mbpp_prompt(tok, items[i], DIRECT_SUFFIX) for i in grp], + return_tensors="pt", padding=True, + add_special_tokens=False).to("cuda") + gen = looper.generate_frozen_prompt( + adapter, tok, enc["input_ids"], kval, max_new_tokens=220, + attention_mask=enc["attention_mask"]) + for gi, i in enumerate(grp): + txt = tok.decode(gen[gi, enc["input_ids"].shape[1]:], + skip_special_tokens=True) + codes[i] = extract_code(txt) + with ThreadPoolExecutor(8) as ex: + oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]), + zip(codes, items))) + + per_label, kdist = {}, {} + for it, ok, kk in zip(items, oks, kstars): + d = per_label.setdefault(it["label"], [0, 0, 0.0]) + d[0] += ok; d[1] += 1; d[2] += kk + kdist.setdefault(it["label"], []).append(kk) + acc = sum(oks) / len(items) + by_label = {l: c / n for l, (c, n, _) in per_label.items()} + mean_k = {l: sum(v) / len(v) for l, v in kdist.items()} + hard = torch.tensor([it["label"] == "hard" for it in items], dtype=torch.float) + kk = torch.tensor(kstars, dtype=torch.float) + r = ((kk - kk.mean()) * (hard - hard.mean())).mean() / (kk.std() * hard.std() + 1e-9) + print(f"pass@1={acc:.3f} by_label={ {l: round(v,3) for l,v in by_label.items()} }") + print(f"mean k* by label: { {l: round(v,2) for l,v in mean_k.items()} } " + f"overall E[k]={sum(kstars)/len(kstars):.2f} " + f"gate-difficulty r={r:.3f} ({time.time()-t0:.0f}s)", flush=True) + + json.dump({"tag": args.tag, "acc": acc, "by_label": by_label, + "mean_kstar": mean_k, "corr_hard": r.item(), + "kstars": kstars, + "per_item": [{"task_id": it["task_id"], "ok": bool(o), + "kstar": kk_} for it, o, kk_ in + zip(items, oks, kstars)]}, + open(OUT / f"eval_code_{args.tag}.json", "w"), indent=1) + print("wrote", OUT / f"eval_code_{args.tag}.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/halting_common.py b/scripts/halting_common.py new file mode 100644 index 0000000..6c0f74f --- /dev/null +++ b/scripts/halting_common.py @@ -0,0 +1,101 @@ +"""E1 (PLAN_SELFPACED): learned per-prompt halting over the workspace loop. + +HaltingMergeAdapter = frozen-recipe MergeAdapter + a halting head read off +the workspace state at the last prompt position after each iteration: + + p_i = sigmoid(w ยท [e_last ; s_hat_i,last] + b) (halt after iter i) + q_i = p_i * prod_{j suffix -> CE, plus a +compute penalty lambda * E[iters] = lambda * sum_i q_i * i. +The halting head's weight is zero-init and its bias starts at -6 +(p ~ 0.0025), so at init >=98% of the mass sits on k_max: the model is +approximately the fixed-k merge (documented tolerance, not bit-exact). + +Deploy: iterate until cumulative halt mass crosses 0.5; per-item k*. +""" + +import torch +import torch.nn as nn +from torch.utils.checkpoint import checkpoint + +from loop_common import ALPHA, MergeAdapter + + +class HaltingMergeAdapter(MergeAdapter): + + def __init__(self, d=1536, hidden=512, alpha=ALPHA, bias0=-6.0): + super().__init__(d=d, hidden=hidden, alpha=alpha) + self.halt = nn.Linear(2 * d, 1) + nn.init.zeros_(self.halt.weight) + nn.init.constant_(self.halt.bias, bias0) + + def halt_prob(self, e_last, s_last): + """p(halt | state) from the last prompt position. (B,) in (0,1).""" + e32, s32 = e_last.float(), s_last.float() + s_hat = s32 * (e32.norm(dim=-1, keepdim=True) + / (s32.norm(dim=-1, keepdim=True) + 1e-6)) + return torch.sigmoid(self.halt(torch.cat([e32, s_hat], -1))).squeeze(-1) + + +def halting_loop_logits(looper, adapter, input_ids, k_max, attention_mask, + prompt_last_idx, loop_mask=None, use_checkpoint=True): + """Soft-halting teacher-forced logits. + + prompt_last_idx: (B,) index of each item's last prompt position. + Returns (logits_of_mixed_state, q, exp_iters): + q (B, k_max) halting distribution over iterations 1..k_max. + """ + calls, _ = looper.capture(input_ids, attention_mask, logits_to_keep=1) + e = looper._hin[looper.l0].detach() + B = e.shape[0] + bidx = torch.arange(B, device=e.device) + with torch.no_grad(): + s = looper.band(e, calls) + s_mix = torch.zeros_like(e) + keep = torch.ones(B, device=e.device) # prob of not-yet-halted + qs = [] + for i in range(k_max): + x = adapter(e, s) + if loop_mask is not None: + x = torch.where(loop_mask[..., None], x, e) + if use_checkpoint: + s = checkpoint(lambda x_: looper.band(x_, calls), x, + use_reentrant=False) + else: + s = looper.band(x, calls) + p = adapter.halt_prob(e[bidx, prompt_last_idx], + s[bidx, prompt_last_idx]) + q_i = keep * p if i < k_max - 1 else keep # remainder to k_max + keep = keep * (1 - p) + s_mix = s_mix + q_i[:, None, None] * s.float() + qs.append(q_i) + q = torch.stack(qs, -1) # (B, k_max) + exp_iters = (q * torch.arange(1, k_max + 1, device=q.device)).sum(-1) + logits = looper.suffix_logits(s_mix.to(e.dtype), calls) + return logits, q, exp_iters + + +@torch.no_grad() +def halted_k_per_item(looper, adapter, input_ids, k_max, attention_mask, + prompt_last_idx, thresh=0.5): + """Deploy-time halting: smallest k where cumulative halt mass >= thresh + (k_max if never). Returns (B,) ints in 1..k_max.""" + calls, _ = looper.capture(input_ids, attention_mask, logits_to_keep=1) + e = looper._hin[looper.l0].detach() + B = e.shape[0] + bidx = torch.arange(B, device=e.device) + s = looper.band(e, calls) + keep = torch.ones(B, device=e.device) + cum = torch.zeros(B, device=e.device) + kstar = torch.full((B,), k_max, dtype=torch.long, device=e.device) + for i in range(k_max): + x = adapter(e, s) + s = looper.band(x, calls) + p = adapter.halt_prob(e[bidx, prompt_last_idx], + s[bidx, prompt_last_idx]) + cum = cum + keep * p + keep = keep * (1 - p) + newly = (cum >= thresh) & (kstar == k_max) & (i < k_max - 1) + kstar[newly] = i + 1 + return kstar diff --git a/scripts/train_gate_code.py b/scripts/train_gate_code.py new file mode 100644 index 0000000..0961420 --- /dev/null +++ b/scripts/train_gate_code.py @@ -0,0 +1,120 @@ +"""E1 trainer (PLAN_SELFPACED): halting-gate merge on MBPP prompt-only loop. + +Same data/recipe as train_merge_code.py (STaR-verified code supervision, +prompt-only loop_mask), but NO difficulty->depth curriculum: every batch +mixes difficulties and runs the soft-halting loop to k_max โ€” the GATE must +learn the allocation the curriculum used to hand-code. + +Loss: CE(code tokens | s_mix) + lam * mean(E[iters]). +""" + +import argparse +import json +import os +import math +import random +import sys +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +from halting_common import HaltingMergeAdapter, halting_loop_logits +from loop_common import BandLooper +from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt +from train_merge_code import build_code_batch, MAX_TOK + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) +STEPS = 600 +BATCH = 4 +WARMUP = 20 + +ap = argparse.ArgumentParser() +ap.add_argument("--seed", type=int, default=0) +ap.add_argument("--lam", type=float, default=1e-3, + help="compute penalty weight on E[iterations]") +ap.add_argument("--kmax", type=int, default=4) +ap.add_argument("--lr", type=float, default=1e-3) +ap.add_argument("--lam-warmup", type=int, default=100, + help="steps before the penalty ramps in (anti-collapse)") +ARGS = ap.parse_args() +LR = ARGS.lr +TAG = (f"gate_l{ARGS.lam:g}" + (f"_s{ARGS.seed}" if ARGS.seed else "")) + + +def lr_at(step): + if step < WARMUP: + return LR * (step + 1) / WARMUP + t = (step - WARMUP) / max(1, STEPS - WARMUP) + return 1e-4 + 0.5 * (LR - 1e-4) * (1 + math.cos(math.pi * t)) + + +def lam_at(step): + return ARGS.lam * min(1.0, max(0.0, (step - ARGS.lam_warmup) / 100)) + + +def main(): + rng = random.Random(ARGS.seed) + torch.manual_seed(ARGS.seed) + data = json.load(open(OUT / "mbpp_data.json")) + model, tok = load_model(dtype=torch.bfloat16) + for p in model.parameters(): + p.requires_grad_(False) + looper = BandLooper(model) + adapter = HaltingMergeAdapter( + d=model.config.get_text_config().hidden_size).cuda() + opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01) + + train = [it for it in data if it["split"] == "train" + and it["label"] != "drop" and it["sol_code"]] + train = [it for it in train + if len(tok(mbpp_prompt(tok, it, DIRECT_SUFFIX))["input_ids"]) + + len(tok(it["sol_code"])["input_ids"]) + 12 <= MAX_TOK] + val = train[:24] + pool = train[24:] + print(f"pool={len(pool)} val={len(val)} lam={ARGS.lam} kmax={ARGS.kmax}", + flush=True) + + log = [] + t0 = time.time() + for step in range(STEPS): + batch = rng.sample(pool, BATCH) + ids, msk, lab, lmask = build_code_batch(tok, batch) + # last prompt position = last position where label == -100 & mask==1 + pl = (lmask.long().cumsum(-1).argmax(-1)) + for g in opt.param_groups: + g["lr"] = lr_at(step) + logits, q, ei = halting_loop_logits( + looper, adapter, ids, ARGS.kmax, msk, pl, loop_mask=lmask) + ce = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(), + lab[:, 1:].flatten(), ignore_index=-100) + loss = ce + lam_at(step) * ei.mean() + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0) + opt.step() + hard_b = [it["label"] == "hard" for it in batch] + log.append({"step": step, "ce": ce.item(), "ei": ei.mean().item(), + "ei_hard": (ei[torch.tensor(hard_b)].mean().item() + if any(hard_b) else None), + "ei_easy": (ei[~torch.tensor(hard_b)].mean().item() + if not all(hard_b) else None)}) + if step % 10 == 0: + r = log[-1] + print(f"step {step:4d} ce={r['ce']:.4f} E[k]={r['ei']:.2f} " + f"(hard {r['ei_hard']} / easy {r['ei_easy']}) " + f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True) + if step % 100 == 99 or step == STEPS - 1: + torch.save(adapter.state_dict(), + OUT / f"adapter_{TAG}_e{step+1}.pt") + json.dump(log, open(OUT / f"train_{TAG}_log.json", "w")) + print("done", flush=True) + + +if __name__ == "__main__": + main()