mirror of
https://github.com/wassname/evil_MoE.git
synced 2026-08-20 06:00:25 +08:00
wassnameandClaudypoo
07363f1ede
Dropped dead job-ID narrative (job 60/64) on rollout_ablate_frac, the 'vanilla step 17' dead-run ref in eval.py, the 'old signed sum' dead-code ref in proj.py, and the conversational 'current experiment line' lead. Removed the dead probe-traj justfile recipe. Kept all TODO/FIXME and the 'why' memory-tuning comments. Smoke green (cout->0). Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
103 lines
4.6 KiB
Python
103 lines
4.6 KiB
Python
"""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) --
|
|
a long prompt can spike the full-logits lm_head ~4 GiB and OOM without this.
|
|
"""
|
|
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)
|