diff --git a/src/projected_grpo/eval.py b/src/projected_grpo/eval.py new file mode 100644 index 0000000..e9aba94 --- /dev/null +++ b/src/projected_grpo/eval.py @@ -0,0 +1,102 @@ +"""Evaluation and reference-model helpers for the training loop. + +Three read-only helpers that touch the model but never train it: a reference +log-prob pass (the AntiPaSTO adapter zeroed = the base model), the deploy-time +quarantine ablation, and a hack/solve eval on a fixed prompt subset. +""" +from __future__ import annotations + +from contextlib import contextmanager + +import torch + +from .proj import per_token_logps +from .rewards import compute_reward + + +def ref_logprobs_via_zero_delta( + model, merged: torch.Tensor, wrappers: dict, plen: int, +) -> torch.Tensor: + """Compute pi_ref logprobs on completion tokens only. + + AntiPaSTO: W' = W + U diag(delta_S) Vh. At delta_S=0, W' = W exactly + (verified bit-exact in step 1). Save -> zero -> forward -> restore. + Zero extra VRAM vs a separately loaded ref_model. + + Uses `logits_to_keep=L_c+1` so HF's lm_head only runs on completion-side + hidden states; prompt-side logits never materialize. Saves + ~plen/(plen+L_c) memory at the lm_head call (~33% at plen=500, L_c=1024). + That was the OOM site at vanilla step 17 (long prompt -> 4 GiB lm_head spike). + """ + saved = {n: info["delta_S"].data.clone() for n, info in wrappers.items()} + try: + for info in wrappers.values(): + info["delta_S"].data.zero_() + L_c = merged.shape[1] - plen + logits = model(merged, logits_to_keep=L_c + 1).logits[:, :-1] + return per_token_logps(logits, merged[:, plen:]) + finally: + for n, info in wrappers.items(): + info["delta_S"].data.copy_(saved[n]) + + +@contextmanager +def ablate_quarantine(wrappers: dict): + """Zero the routing quarantine (delta_S_hack) for the duration -- the + eval-time ablation of the routed hack capability. Save -> zero -> (eval) -> + restore. The route/route2 arms' deployment model IS this ablated state. + + TODO(post-deploy-finetune): SGTM's ablate(trainable=True) reinits the forget + weights to the retain-dims' std instead of zeroing, so the model stays + finetunable after the quarantine is removed (no dead hole). We zero because + we only eval after deploy; add the reinit path if we ever retrain post-ablate. + See docs/grad_routing/sgtm_vs_ours.md.""" + saved = {n: info["delta_S_hack"].data.clone() for n, info in wrappers.items()} + for info in wrappers.values(): + info["delta_S_hack"].data.zero_() + try: + yield + finally: + for n, info in wrappers.items(): + info["delta_S_hack"].data.copy_(saved[n]) + + +@torch.no_grad() +def eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg, device, max_new) -> dict: + """Student-only generate + grade on a FIXED prompt subset (no teacher, no + backward) -- a clean read of what the current adapter does. Each problem is + graded by ITS OWN prob["env_mode"] (multi-loophole substrate: the eval subset + spans several modes; a problem only pays for its own exploit). + + hack = exploited rate (mode's channel credited correctness without the strict + oracle agreeing); solve = gt_correct rate (the exploit-resistant oracle). Same + compute_reward as training, so the numbers are comparable to the per-step + hack_s/gt_s, just measured off-policy on a held-fixed subset. + """ + model.config.use_cache = True + n = hacks = solves = 0 + # per-mode (env_mode) tallies so route's deploy eval can be read per loophole: + # the load-bearing weak-detector test is whether route suppresses the modes + # v_hack never saw (file_marker/sentinel), not just its own (run_tests). + by_mode: dict[str, list[int]] = {} # env_mode -> [hacks, solves, n] + for idx in eval_idxs: + prob = problems[idx] + mode = prob["env_mode"] + prompt = tok.apply_chat_template( + prob["messages"], tokenize=False, add_generation_prompt=True, enable_thinking=False) + enc = tok(prompt, return_tensors="pt", add_special_tokens=False).to(device) + if enc.input_ids.shape[1] + max_new > 2048: + continue + out = model.generate(**enc, generation_config=gen_cfg) + comps = out[:, enc.input_ids.shape[1]:] + tally = by_mode.setdefault(mode, [0, 0, 0]) + for t in tok.batch_decode(comps, skip_special_tokens=True): + r = compute_reward( + t, canonical_solution=prob["canonical"], gt_tests=prob["gt_tests"], + setup_code=prob["setup_code"], func_name_hint=prob["func_name"], + env_mode=mode) + hacks += int(r.exploited); tally[0] += int(r.exploited) + solves += int(r.gt_correct); tally[1] += int(r.gt_correct) + n += 1; tally[2] += 1 + model.config.use_cache = False + return dict(hack=hacks / max(1, n), solve=solves / max(1, n), n=n, by_mode=by_mode) diff --git a/src/projected_grpo/train.py b/src/projected_grpo/train.py index 44ba705..37a8d01 100644 --- a/src/projected_grpo/train.py +++ b/src/projected_grpo/train.py @@ -85,6 +85,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig from .antipasto import wrap_model_with_antipasto from .proj import per_token_logps, project_delta_S_grad, mean_cos_pre_from_grads from .rewards import EnvMode, compute_reward +from .eval import ablate_quarantine, eval_hack_solve, ref_logprobs_via_zero_delta from .tablelog import setup_logging, StepLogger CACHE_ROOT = Path("svd_cache") @@ -490,95 +491,6 @@ def postprocess_v_hack( return v_hack -@torch.no_grad() -def ref_logprobs_via_zero_delta( - model, merged: torch.Tensor, wrappers: dict, plen: int, -) -> torch.Tensor: - """Compute pi_ref logprobs on completion tokens only. - - AntiPaSTO: W' = W + U diag(delta_S) Vh. At delta_S=0, W' = W exactly - (verified bit-exact in step 1). Save -> zero -> forward -> restore. - Zero extra VRAM vs a separately loaded ref_model. - - Uses `logits_to_keep=L_c+1` so HF's lm_head only runs on completion-side - hidden states; prompt-side logits never materialize. Saves - ~plen/(plen+L_c) memory at the lm_head call (~33% at plen=500, L_c=1024). - That was the OOM site at vanilla step 17 (long prompt -> 4 GiB lm_head spike). - """ - saved = {n: info["delta_S"].data.clone() for n, info in wrappers.items()} - try: - for info in wrappers.values(): - info["delta_S"].data.zero_() - L_c = merged.shape[1] - plen - logits = model(merged, logits_to_keep=L_c + 1).logits[:, :-1] - return per_token_logps(logits, merged[:, plen:]) - finally: - for n, info in wrappers.items(): - info["delta_S"].data.copy_(saved[n]) - - -@contextmanager -def ablate_quarantine(wrappers: dict): - """Zero the routing quarantine (delta_S_hack) for the duration -- the - eval-time ablation of the routed hack capability. Save -> zero -> (eval) -> - restore. The route/route2 arms' deployment model IS this ablated state. - - TODO(post-deploy-finetune): SGTM's ablate(trainable=True) reinits the forget - weights to the retain-dims' std instead of zeroing, so the model stays - finetunable after the quarantine is removed (no dead hole). We zero because - we only eval after deploy; add the reinit path if we ever retrain post-ablate. - See docs/grad_routing/sgtm_vs_ours.md.""" - saved = {n: info["delta_S_hack"].data.clone() for n, info in wrappers.items()} - for info in wrappers.values(): - info["delta_S_hack"].data.zero_() - try: - yield - finally: - for n, info in wrappers.items(): - info["delta_S_hack"].data.copy_(saved[n]) - - -@torch.no_grad() -def eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg, device, max_new) -> dict: - """Student-only generate + grade on a FIXED prompt subset (no teacher, no - backward) -- a clean read of what the current adapter does. Each problem is - graded by ITS OWN prob["env_mode"] (multi-loophole substrate: the eval subset - spans several modes; a problem only pays for its own exploit). - - hack = exploited rate (mode's channel credited correctness without the strict - oracle agreeing); solve = gt_correct rate (the exploit-resistant oracle). Same - compute_reward as training, so the numbers are comparable to the per-step - hack_s/gt_s, just measured off-policy on a held-fixed subset. - """ - model.config.use_cache = True - n = hacks = solves = 0 - # per-mode (env_mode) tallies so route's deploy eval can be read per loophole: - # the load-bearing weak-detector test is whether route suppresses the modes - # v_hack never saw (file_marker/sentinel), not just its own (run_tests). - by_mode: dict[str, list[int]] = {} # env_mode -> [hacks, solves, n] - for idx in eval_idxs: - prob = problems[idx] - mode = prob["env_mode"] - prompt = tok.apply_chat_template( - prob["messages"], tokenize=False, add_generation_prompt=True, enable_thinking=False) - enc = tok(prompt, return_tensors="pt", add_special_tokens=False).to(device) - if enc.input_ids.shape[1] + max_new > 2048: - continue - out = model.generate(**enc, generation_config=gen_cfg) - comps = out[:, enc.input_ids.shape[1]:] - tally = by_mode.setdefault(mode, [0, 0, 0]) - for t in tok.batch_decode(comps, skip_special_tokens=True): - r = compute_reward( - t, canonical_solution=prob["canonical"], gt_tests=prob["gt_tests"], - setup_code=prob["setup_code"], func_name_hint=prob["func_name"], - env_mode=mode) - hacks += int(r.exploited); tally[0] += int(r.exploited) - solves += int(r.gt_correct); tally[1] += int(r.gt_correct) - n += 1; tally[2] += 1 - model.config.use_cache = False - return dict(hack=hacks / max(1, n), solve=solves / max(1, n), n=n, by_mode=by_mode) - - # 2-char env_mode codes for compact per-mode hack columns (hk_rt, hk_xc, ...). MODE_CODE: dict[str, str] = { "run_tests": "rt", "eq_override": "eq", "exit_code": "xc",