E1 halting-gate machinery (adapter, trainer, gated eval) + pre-registration item 18
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
Claude Fable 5
@@ -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()
|
||||
@@ -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<i}(1 - p_j), q_kmax += remainder (ACT-style)
|
||||
|
||||
Training (soft, no RL): s_mix = sum_i q_i * s_i -> 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
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user