wassname/evil_MoE
1
0
Fork 0
mirror of https://github.com/wassname/evil_MoE.git synced 2026-07-28 11:19:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(c): extract eval.py (ref_logprobs, ablate_quarantine, eval_hack_solve)

Browse Source 
Relocate the three read-only model helpers out of train.py into eval.py. They use
only torch + per_token_logps (proj) + compute_reward (rewards); no train globals.
Training numbers identical across all 4 smoke arms (resid/qE diagnostic cosines
show last-digit bf16 noise only). MODE_CODE stays in train.py.

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
This commit is contained in:
wassname
2026-06-01 09:22:30 +00:00
co-authored by Claudypoo
parent 692f0ac00a
commit 2001f41e03
2 changed files with 103 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/projected_grpo/eval.py
+102
View File
@@ -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)
src/projected_grpo/train.py
+1 -89
View File
@@ -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",