wassname/evil_MoE
1
0
Fork 0
mirror of https://github.com/wassname/evil_MoE.git synced 2026-07-18 12:20:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

rename route2->routeV; heavy seeded final eval; save delta_S_hack

Browse Source 
route2 (binary-tau) and routeV (banded gate) are different methods -- give the
new one a distinct id so old/new runs can't be confused (see hypothesis doc).
- src/vgrout/* + justfile: route2->routeV, routing2->routingV (figs.py keeps the
  old keys for plotting historical runs).
- Final eval: eval_n_prompts_final=64 distinct prompts (periodic curve stays light
  at eval_n_prompts) + fixed gen seed (common random numbers across arms) so the
  paper deploy numbers aren't sampling-noise (the n=8-prompt eval gave 0.031 vs
  0.125 at the same checkpoint).
- save_ckpt: also write delta_S_hack to sibling _hack.safetensors so runs can be
  re-scored knob-ON at higher n later (train.safetensors stays delta_S-only).

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
This commit is contained in:
wassname
2026-06-06 12:08:28 +00:00
co-authored by Claudypoo
parent 9c76584970
commit 73936c822f
6 changed files with 118 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
justfile
+15 -15
View File
@@ -45,14 +45,14 @@ smoke-route *ARGS:
--teacher-pool-dir=out/pools/teacher_pool --mix-ratio=0.5 \
--eval-ablate-every=10 --eval-n-prompts=2 {{ ARGS }}
# Routing-v2 path (route2): per-rollout calibrated-tau cosine routing into the
# Routing-v2 path (routeV): per-rollout calibrated-tau cosine routing into the
# scale-matched delta_S_hack quarantine. Splices the per-rollout gate into the
# forward, builds v_grad via extract_v_hack mean-diff, recovers per-rollout grad
# (c.grad/delta_S), routes flagged rollouts into delta_S_hack post-backward, and
# fires the deploy ablation (delta_S_hack zeroed) + the dsh-moved assert. Exercises
# tau/hkgap/qE logging too.
smoke-route2 *ARGS:
BEARTYPE=1 {{ TRAIN }} smoke --intervention=route2 \
smoke-routeV *ARGS:
BEARTYPE=1 {{ TRAIN }} smoke --intervention=routeV \
--teacher-pool-dir=out/pools/teacher_pool --mix-ratio=0.5 \
--eval-ablate-every=10 --eval-n-prompts=2 {{ ARGS }}
@@ -109,7 +109,7 @@ fast-vanilla *ARGS:
# Goal 1: same recipe with --intervention=erase. Run only after fast-vanilla passes UAT.
# mix_ratio=0.125 + group=8 are the locked-in fast defaults (config), not flags here.
# No --v-hack-path: erase uses the prog_wide default (auto-extracts v_hack_pairset_prog_wide),
# same pair set as route2, so the arms are apples-to-apples.
# same pair set as routeV, so the arms are apples-to-apples.
fast-projected *ARGS:
{{ TRAIN }} fast --intervention=erase \
--teacher-pool-dir=out/pools/teacher_pool \
@@ -164,7 +164,7 @@ queue-substrate:
pueue add -w "$PWD" -o 5 -l "why: vanilla emergence reference (4-mode substrate); resolve: per-mode deploy-hack baseline for the overlay" -- {{ TRAIN }} fast --intervention=none --out-tag=_sub4_vanilla
pueue add -w "$PWD" -o 5 -l "why: erase arm (one-sided projection); resolve: per-mode deploy hack vs vanilla at matched solve" -- {{ TRAIN }} fast --intervention=erase --out-tag=_sub4_erase
pueue add -w "$PWD" -o 5 -l "why: route arm (subspace-projection quarantine, rf5); resolve: deploy hack on held-out modes vs vanilla at matched solve" -- {{ TRAIN }} fast --intervention=route --out-tag=_sub4_route
pueue add -w "$PWD" -o 5 -l "why: route2 calibrated-tau routing into scale-matched delta_S_hack; resolve: held-out deploy hack suppressed vs vanilla at matched solve" -- {{ TRAIN }} fast --intervention=route2 --out-tag=_sub4_route2
pueue add -w "$PWD" -o 5 -l "why: routeV calibrated-tau routing into scale-matched delta_S_hack; resolve: held-out deploy hack suppressed vs vanilla at matched solve" -- {{ TRAIN }} fast --intervention=routeV --out-tag=_sub4_routeV
# CANONICAL plotting entrypoint for the substrate sweep. One command, four figures
# (per-mode by-method + by-hack, and the aggregate "total hacks per arm" + overlay,
@@ -173,7 +173,7 @@ plot GLOB='logs/*_sub4_*.log' STEM='out/figs/substrate':
uv run python scripts/plot_substrate.py {{ GLOB }} --out-stem {{ STEM }}
# All-arms per-mode DEPLOY overlay (#162) from the per_mode_deploy.json artifacts
# (honest shipped-model numbers; route2-safe -- reads JSON, not logs). Default
# (honest shipped-model numbers; routeV-safe -- reads JSON, not logs). Default
# globs every sub4 run dir. -> out/figs/deploy_overlay.png
plot-deploy GLOB='out/runs/*sub4*/per_mode_deploy.json' OUT='out/figs/deploy_overlay.png':
uv run python scripts/plot_deploy_overlay.py {{ GLOB }} --out {{ OUT }}
@@ -435,31 +435,31 @@ paper-md:
# substrate unless noted. All emit out/runs/<ts>_<tag>/per_mode_deploy.json.
# ─────────────────────────────────────────────────────────────────────────────
# H: route2 deploy-hack stays ~0 to convergence while vanilla saturates (not
# H: routeV deploy-hack stays ~0 to convergence while vanilla saturates (not
# collapses). Long-run A4 figure. Stabilised optimizer: tiny KL beta=1e-5 (anchor
# coherence, too weak to undo the hack reward -- see RESEARCH_JOURNAL 2026-06-02
# job-85 divergence) + normal Adam 0.9/0.99; lr unchanged (SVD adapter tolerates).
# ARM in {none, route2}. UAT: deploy hack/solve trajectory to 200, no lp_s collapse.
# ARM in {none, routeV}. UAT: deploy hack/solve trajectory to 200, no lp_s collapse.
paper-longrun ARM SEED='41':
pueue add -w "$PWD" -o 0 -l "why: {{ ARM }}-200 KL-stabilised (beta=1e-5, Adam 0.9/0.99) long-run A4; resolve: route2 deploy hack~0 to 200 while vanilla saturates w/o collapse" -- \
pueue add -w "$PWD" -o 0 -l "why: {{ ARM }}-200 KL-stabilised (beta=1e-5, Adam 0.9/0.99) long-run A4; resolve: routeV deploy hack~0 to 200 while vanilla saturates w/o collapse" -- \
{{ TRAIN }} fast --intervention={{ ARM }} --seed={{ SEED }} \
--beta=1e-5 --adam-beta1=0.9 --adam-beta2=0.99 \
--steps=200 --eval-ablate-every=20 --out-tag=_{{ ARM }}200_kl5_s{{ SEED }}
# H: route2 suppresses ENDOGENOUSLY-emerged hacks (no teacher mix at all -- pure
# H: routeV suppresses ENDOGENOUSLY-emerged hacks (no teacher mix at all -- pure
# on-policy). mix=0 keeps the pool only for the 4-mode partition + v_grad extraction.
# 800 steps ~= 100 reference-paper steps. ARM in {none, route2}. SLOW (~32h).
# 800 steps ~= 100 reference-paper steps. ARM in {none, routeV}. SLOW (~32h).
paper-noteacher ARM SEED='41' STEPS='800':
pueue add -w "$PWD" -o 0 -l "why: {{ ARM }} NO-TEACHER mix=0 pure on-policy {{ STEPS }}step; resolve: does route2 suppress endogenous hacks vs vanilla" -- \
pueue add -w "$PWD" -o 0 -l "why: {{ ARM }} NO-TEACHER mix=0 pure on-policy {{ STEPS }}step; resolve: does routeV suppress endogenous hacks vs vanilla" -- \
{{ TRAIN }} fast --intervention={{ ARM }} --seed={{ SEED }} \
--mix-ratio=0 --steps={{ STEPS }} --eval-ablate-every=20 \
--out-tag=_{{ ARM }}_noteacher_s{{ SEED }}
# H: route2 holds suppression after the teacher crutch is removed. Teacher-seeds all
# H: routeV holds suppression after the teacher crutch is removed. Teacher-seeds all
# 4 hacks for OFF steps, then cuts to pure on-policy. Smarter no-teacher test (pure
# mix=0 from step 0 may never emerge all modes). ARM in {none, route2}.
# mix=0 from step 0 may never emerge all modes). ARM in {none, routeV}.
paper-teacheroff ARM SEED='41' OFF='40' STEPS='200':
pueue add -w "$PWD" -o 0 -l "why: {{ ARM }} teacher-off@{{ OFF }} curriculum (seed hacks then on-policy); resolve: route2 deploy hack stays ~0 after teacher cut at {{ OFF }}" -- \
pueue add -w "$PWD" -o 0 -l "why: {{ ARM }} teacher-off@{{ OFF }} curriculum (seed hacks then on-policy); resolve: routeV deploy hack stays ~0 after teacher cut at {{ OFF }}" -- \
{{ TRAIN }} fast --intervention={{ ARM }} --seed={{ SEED }} \
--teacher-off-step={{ OFF }} --steps={{ STEPS }} --eval-ablate-every=20 \
--out-tag=_{{ ARM }}_toff{{ OFF }}_s{{ SEED }}
src/vgrout/antipasto.py
+5 -5
View File
@@ -77,10 +77,10 @@ def _delta_hook(layer: nn.Linear, args: tuple, y: Tensor) -> Tensor:
shape [r] in the same basis (capacity-balanced, no sink-bias) and both 0 at
init -> identity. Routing decides per update which gradient lands in which:
erase strips the hack-ward part (proj.py); route parks it in delta_S_hack
by subspace projection (proj.py); route2 parks it by a per-rollout
by subspace projection (proj.py); routeV parks it by a per-rollout
calibrated-tau cosine gate (train.py, post-backward).
For route2's per-rollout routing (layer._antipasto_grad_probe) we splice a
For routeV's per-rollout routing (layer._antipasto_grad_probe) we splice a
per-token gate c (init 1, forward-identity) onto the delta_S path: after
backward c.grad = delta_S * g_b, so train.py recovers the per-rollout delta_S
gradient, flags rollouts by cos(g_b, v_grad) vs tau, and routes the flagged
@@ -122,7 +122,7 @@ def wrap_model_with_antipasto(
Frozen U/Vh stored on the layer as buffers `_antipasto_{U,Vh}` in the
layer's native dtype. delta_S/delta_S_hack kept in fp32 (tiny, ~r per module).
`grad_probe` (route2 only): splice a per-token gate c into the delta_S path so
`grad_probe` (routeV only): splice a per-token gate c into the delta_S path so
train.py can recover the per-rollout delta_S gradient and route flagged
rollouts into delta_S_hack post-backward. Off -> plain forward (none/erase/route).
"""
@@ -154,7 +154,7 @@ def wrap_model_with_antipasto(
linear.register_parameter("_antipasto_delta_S_hack", delta_S_hack)
info = {"layer": linear, "delta_S": delta_S,
"delta_S_hack": delta_S_hack, "handle": None, "r": r}
linear._antipasto_grad_probe = grad_probe # route2: gate the delta_S path
linear._antipasto_grad_probe = grad_probe # routeV: gate the delta_S path
linear._antipasto_gate = None # grad-probe leaf, set per forward
info["handle"] = linear.register_forward_hook(_delta_hook)
out[name] = info
@@ -206,7 +206,7 @@ def ref_logprobs_via_zero_delta(
def ablate_quarantine(wrappers: dict):
"""Zero the routing quarantine (δS_hack) for the duration: the deploy-time
ablation of the routed hack capability. Save -> zero -> (eval) -> restore.
The route/route2 deployment model IS this ablated state.
The route/routeV 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, keeping the model
src/vgrout/eval.py
+1 -1
View File
@@ -44,7 +44,7 @@ def ref_logprobs_via_zero_delta(
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.
restore. The route/routeV 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
src/vgrout/figs.py
+3
View File
@@ -22,6 +22,9 @@ FIGS_DIR = Path("docs/figs")
# falls through to its raw tag, so a new arm shows up loud rather than silently
# mislabelled.
ARM_DISPLAY = {
# routeV is the current banded-gate arm; routing2/route2 are the old binary-tau runs
# (kept so historical run artifacts still plot -- see rename, 2026-06-06).
"routingV": "route", "routeV": "route",
"routing2": "route", "route2": "route",
"routing2_grad": "route", "routing2_act": "route (act)",
"projected": "erase", "route": "route", "erase": "erase", "vanilla": "vanilla",
src/vgrout/tablelog.py
+4 -4
View File
@@ -106,7 +106,7 @@ class StepLogger:
_Col("hack_s", 7, "hack_s?", "frac", "student hack-flagged rollouts (the headline)"),
_Col("hack_t", 7, "hack_t", "frac", "teacher hack-flagged rollouts (sanity: pool hacks)"),
# Deploy-eval shown for EVERY arm (nan on steps it's not run -> see it ride
# along as training proceeds). route/route2: quarantine knob OFF. vanilla/erase:
# along as training proceeds). route/routeV: quarantine knob OFF. vanilla/erase:
# the trained model itself. Apples-to-apples knob-off deploy number, the plot series.
_Col("hack_deploy", 7, "hk_dep", "+.2f", "DEPLOY-eval hack (route: quarantine OFF; vanilla/erase: trained model); held-out subset, T=0.7, every eval_ablate_every steps; nan between"),
_Col("solve_deploy", 7, "slv_dep", "+.2f", "DEPLOY-eval solve (same cadence; nan between)"),
@@ -133,19 +133,19 @@ class StepLogger:
_Col("cos_post", 6, "cout", ".2f", "hack-ward fraction AFTER projection (want ~0: all removed)"),
_Col("fired", 5, "fired", ".2f", "fraction of modules where projection fired"),
]
# route2: the routing gate is the pair-calibrated BAND. Per rollout,
# routeV: the routing gate is the pair-calibrated BAND. Per rollout,
# f = clamp((cos(g_b,v_grad) - lower)/(upper - lower), 0, 1) routes that
# fraction into the quarantine. lower/upper = mean clean/hack pair cosines.
# Surface where live cos sits (tau), the band width (hkgap), the routed
# fraction (frout, the mass gauge), and the post-routing leak (resid).
if arm == "routing2":
if arm == "routingV":
cols += [
_Col("tau", 6, "tau", "+.2f", "median live cos(g_b, v_grad); should sit inside the band [lower, upper]"),
_Col("hkgap", 6, "hkgap", "+.2f", "band width upper-lower (mean hack-pair minus clean-pair cosine); >0 = v_grad separates (else direction dead/random)"),
_Col("frout", 6, "frout", "+.2f", "mean routed fraction f over rollouts (the routed-mass gauge; compare real-vs-random at matched frout)"),
_Col("resid", 6, "resid", "+.2f", "cos(deployed delta_S.grad AFTER routing, v_grad); ~0 = hack stripped cleanly, >0 = leak into deployed knob"),
]
if arm in ("routing", "routing2"):
if arm in ("routing", "routingV"):
cols += [
# Deploy eval (knob-OFF) is hk_dep below. The train-vs-deploy 2x2's
# knob-ON pass runs once post-loop (FINAL EVAL), not per-step; the
src/vgrout/train.py
+90 -69
View File
@@ -17,7 +17,7 @@ Arms (--intervention, one knob):
none measure only; δS.grad untouched (vanilla GRPO)
erase subtract the hack-ward component of δS.grad
route park that component in the δS_hack quarantine, ablated at deploy (Cloud 2024)
route2 route per-rollout by a calibrated-τ cosine gate, cos(g_b, v_grad) > τ
routeV route per-rollout by a calibrated-τ cosine gate, cos(g_b, v_grad) > τ
Hyperparameters from ariahw/rl-rewardhacking config.py (docs/grpo_hyperparams.md);
SmokeConfig / FastConfig / FullConfig below hold the scale knobs.
@@ -116,8 +116,8 @@ class Config:
sub-30-min iteration loops.
"""
# The four arms (see module docstring). `arm` (property below) is the derived
# display name; route2 gate spec: docs/spec/20260601_calibrated_tau_route2grad.md.
intervention: Literal["none", "erase", "route", "route2"] = "erase"
# display name; routeV gate spec: docs/spec/20260601_calibrated_tau_route2grad.md.
intervention: Literal["none", "erase", "route", "routeV"] = "erase"
# ── scale knobs: every preset overrides these ──
model: str = "Qwen/Qwen3-4B"
steps: int = 100
@@ -141,7 +141,7 @@ class Config:
preserve_magnitude: bool = True
gate_mode: Literal["one_sided", "no_gate", "reverse"] = "one_sided"
project_overshoot: float = 1.0 # remove overshoot*c_use@V; 1.0=just remove, 1.1=10% reversal of hack-ward grad
# route/route2 exploration floor: fraction of student rollouts sampled with the
# route/routeV exploration floor: fraction of student rollouts sampled with the
# quarantine (δS_hack) ablated, i.e. from the DEPLOYED model. Intent: guard hack-
# saturation -- if on-policy sampling collapses onto hacking, every rollout routes
# to the quarantine and the deployed δS never sees a solve gradient. Grading these
@@ -184,7 +184,11 @@ class Config:
# one n=64 pass, ~230s, not two). Long-horizon recipes (paper-longrun, A5) pin a
# sparser cadence (10/20) explicitly. See journal 2026-06-04 (a) for the cost audit.
eval_ablate_every: int = 5
eval_n_prompts: int = 8
eval_n_prompts: int = 8 # periodic (per-step) deploy eval: light, for the smoothed curve
# Final (post-loop) eval covers MORE distinct prompts than the periodic curve so the
# paper deploy hack/solve has a tight CI (the periodic n=8-prompts eval is sampling-noisy:
# eval gen is do_sample T=0.7, see EVAL_GEN_SEED). Capped at the available pool size.
eval_n_prompts_final: int = 64
# Save the deploy adapter (δS only, ~2.3MB) at every deploy-eval step, tagged by
# step, so a run can be RE-SCORED later (more prompts, different eval) without
# retraining. Tiny per ckpt; a 200-step run at every-10 is ~46MB. Off for big sweeps.
@@ -194,8 +198,8 @@ class Config:
# it. DEFAULT prog_wide (30 pairs) -- the proven main set; richer than the 18 hand-crafted
# vgrout.pairs.PAIRS, which remain the fallback only if this is set to None explicitly.
vhack_pairs_path: Path | None = Path("out/pairsets/prog_wide.json")
# Directionality control: replace route2's pair-derived v_grad with a per-module
# Haar-random unit vector. Tests whether route2's suppression NEEDS the direction
# Directionality control: replace routeV's pair-derived v_grad with a per-module
# Haar-random unit vector. Tests whether routeV's suppression NEEDS the direction
# (H4: alignment) or is alignment-agnostic quarantine-absorption (H2). Seeded so
# multiple draws give a distribution ("works half the time?"). NOTE on what Haar
# actually controls for: its ~0 cos with the hack dir is concentration of measure
@@ -207,11 +211,11 @@ class Config:
# "must it point at the hack specifically?". Cosine here is correlational; the ablation
# is the causal test. Refresh no-ops when set, so the direction stays the one fixed
# random draw regardless of --vhack-refresh-every.
route2_random_v_seed: int | None = None
# route2 granularity: False = route per ROLLOUT (sum tokens, one cos/f per rollout;
routeV_random_v_seed: int | None = None
# routeV granularity: False = route per ROLLOUT (sum tokens, one cos/f per rollout;
# the preregistered default, denoises the cos sign + matches GRPO per-rollout adv).
# True = route per TOKEN (one cos/f per token; finer but noisier). Ablation arm.
route2_per_token: bool = False
routeV_per_token: bool = False
# Per-source cin diagnostic: split each prompt's backward into student-only
# + teacher-only passes (~2x backward time). 1 = every step (default; full
# signal); N>1 = only every Nth step (combined backward elsewhere, ~halves
@@ -237,7 +241,7 @@ class Config:
# so every arm runs pure on-policy past step 30, keeping deploy numbers apples-to-
# apples. None = never cut. See step-loop use.
teacher_off_step: int | None = 30
# A5 no-cheat generalisation: restrict teacher demos (and thus the route2 tau
# A5 no-cheat generalisation: restrict teacher demos (and thus the routeV tau
# hack-anchor) to these env_modes only. Held-out modes stay in the training set
# but train PURELY ON-POLICY (no teacher rows, never seed the hack-anchor) -- the
# student must emerge them itself, and we measure whether routing on the
@@ -265,7 +269,7 @@ class Config:
"""Display name for run-id / BLUF / logs (results.py + plot_dynamics
classify off this). One-to-one with intervention; not a CLI flag."""
return {"none": "vanilla", "erase": "projected",
"route": "routing", "route2": "routing2"}[self.intervention]
"route": "routing", "routeV": "routingV"}[self.intervention]
@dataclass(kw_only=True)
@@ -274,7 +278,7 @@ class SmokeConfig(Config):
the every-25-step save_ckpt trigger. ~1-2 min wall-clock."""
model: str = "llamafactory/tiny-random-qwen3"
steps: int = 30
group: int = 4 # >=4 so route2 smoke (mix=0.5 -> G_s=2) can split a rollout_ablate_frac slice; G_s=1 couldn't
group: int = 4 # >=4 so routeV smoke (mix=0.5 -> G_s=2) can split a rollout_ablate_frac slice; G_s=1 couldn't
max_new: int = 32
n_problems: int = 100
beta: float = 0.0
@@ -321,9 +325,9 @@ class FullConfig(Config):
def _haar_unit_dirs(v_grad: dict, seed: int, device) -> dict:
"""Per-module Haar-random unit vectors matching v_grad's shapes -- the OUT-OF-SUBSPACE
directionality control for route2 (~0 cos with the hack dir by concentration of measure,
directionality control for routeV (~0 cos with the hack dir by concentration of measure,
not by being a 'cleaner' placebo). Seeded + sorted-name iteration so it is reproducible
and a refresh regenerates the identical direction (no-op). See Config.route2_random_v_seed."""
and a refresh regenerates the identical direction (no-op). See Config.routeV_random_v_seed."""
g = torch.Generator().manual_seed(seed)
out = {}
for name in sorted(v_grad):
@@ -444,12 +448,12 @@ def main(cfg: Config) -> int:
model.config.use_cache = False
# ── AntiPaSTO adapter: δS (kept) + δS_hack (quarantine), same shape r ──
is_route2 = cfg.intervention == "route2"
is_routeV = cfg.intervention == "routeV"
wrappers = wrap_model_with_antipasto(
model, model_name, CACHE_ROOT, device,
grad_probe=is_route2, # route2 needs the per-rollout δS gate probe
grad_probe=is_routeV, # routeV needs the per-rollout δS gate probe
)
# δS_hack only gets a grad under route (proj.py subspace split) or route2
# δS_hack only gets a grad under route (proj.py subspace split) or routeV
# (per-rollout τ routing); under none/erase its grad stays None, so AdamW skips
# it and it stays exactly 0 (forward adds 0 -> identity).
delta_params = [info["delta_S"] for info in wrappers.values()]
@@ -457,25 +461,25 @@ def main(cfg: Config) -> int:
logger.info(f"trainable delta_S: {sum(p.numel() for p in delta_params):,} "
f"(+{sum(p.numel() for p in delta_hack_params):,} delta_S_hack quarantine)")
# ── hack direction: v_hack (erase/route project against it) or v_grad (route2) ──
# ── hack direction: v_hack (erase/route project against it) or v_grad (routeV) ──
# Vanilla (none) is pure GRPO and ignores v_hack entirely (the cin/cout columns
# are hidden, so v_hack=None just means no subspace machinery).
v_grad = None # set only by the route2 grad-mask branch below
if cfg.intervention in ("none", "route2"):
v_grad = None # set only by the routeV grad-mask branch below
if cfg.intervention in ("none", "routeV"):
if cfg.intervention == "none" and cfg.v_hack_path is not None:
logger.info(f"vanilla arm: ignoring --v-hack-path={cfg.v_hack_path} "
"(no projection; cin/cout diagnostics off)")
v_hack = None # route2 routes via the mask, not erase/route grad surgery
if is_route2:
v_hack = None # routeV routes via the mask, not erase/route grad surgery
if is_routeV:
# The persona pairs are the only "detector" (weak, self-supervised). They
# produce the routing direction; no oracle, no gt_pass.
if cfg.vhack_pairs_path is not None:
from .pairs_from_pool import load_pairs_json
MASK_PAIRS = load_pairs_json(cfg.vhack_pairs_path)
logger.info(f"route2 pairs: pool-derived ({cfg.vhack_pairs_path}) -> {len(MASK_PAIRS)} pairs")
logger.info(f"routeV pairs: pool-derived ({cfg.vhack_pairs_path}) -> {len(MASK_PAIRS)} pairs")
else:
from .pairs import PAIRS as MASK_PAIRS
logger.info(f"route2 pairs: hand-crafted PAIRS -> {len(MASK_PAIRS)} pairs")
logger.info(f"routeV pairs: hand-crafted PAIRS -> {len(MASK_PAIRS)} pairs")
model.eval()
# gradient-space mean-diff. extract_v_hack gives per-pair GRPO gradients
# on δS; v_grad = unit(mean(g_hack - g_clean)) per module, oriented
@@ -490,22 +494,22 @@ def main(cfg: Config) -> int:
for name in wrappers:
d = (raw_grads[f"hack/{name}"] - raw_grads[f"clean/{name}"]).mean(0)
v_grad[name] = (d / d.norm().clamp_min(1e-12)).to(device)
logger.info(f"route2 grad: built v_grad (gradient mean-diff) for {len(v_grad)} modules")
if cfg.route2_random_v_seed is not None:
v_grad = _haar_unit_dirs(v_grad, cfg.route2_random_v_seed, device)
logger.info(f"route2 grad: OVERRODE v_grad with Haar-random dirs "
f"(seed={cfg.route2_random_v_seed}) -- directionality control (H2 vs H4)")
logger.info(f"routeV grad: built v_grad (gradient mean-diff) for {len(v_grad)} modules")
if cfg.routeV_random_v_seed is not None:
v_grad = _haar_unit_dirs(v_grad, cfg.routeV_random_v_seed, device)
logger.info(f"routeV grad: OVERRODE v_grad with Haar-random dirs "
f"(seed={cfg.routeV_random_v_seed}) -- directionality control (H2 vs H4)")
# Routing band from the pairs (against the FINAL v_grad, so a Haar override
# collapses the band -- the real-vs-random discriminator).
route_band = route_band_edges(raw_grads, v_grad, device)
_mean_bw = sum(hi - lo for lo, hi in route_band.values()) / len(route_band)
logger.info(f"route2 band: edges from {len(route_band)} modules, "
logger.info(f"routeV band: edges from {len(route_band)} modules, "
f"mean width(upper-lower)={_mean_bw:+.3f} "
f"(>0 = pairs separate; ~0 = random/degenerate)")
# On a REAL v_grad the band must open (hack pairs align more than clean).
# A collapsed/inverted real band = broken extraction silently mimicking the
# random control -> fail loud. The Haar control is allowed to collapse.
if cfg.route2_random_v_seed is None:
if cfg.routeV_random_v_seed is None:
assert _mean_bw > 0, (
f"real v_grad gave non-positive mean band width {_mean_bw:+.3f}: "
"hack pairs do not separate from clean -> extraction broken")
@@ -577,7 +581,7 @@ def main(cfg: Config) -> int:
if cfg.teacher_pool_dir is not None:
# mix=0 is the NO-TEACHER ablation: pure on-policy GRPO (G_t=0, no teacher
# rollouts injected) while the pool is still loaded for the 4-mode partition
# and route2 v_grad extraction. Using the pairs for v_grad is allowed under
# and routeV v_grad extraction. Using the pairs for v_grad is allowed under
# the no-cheat invariant; mixing teacher rollouts into training is the thing
# mix=0 removes. mix in [0,1).
if not (0.0 <= cfg.mix_ratio < 1.0):
@@ -725,7 +729,7 @@ def main(cfg: Config) -> int:
See Config.rollout_ablate_frac for why. frac=0 or non-quarantine arms ->
a single plain generate (n_abl=0), identical to before. Returns (rows, n_abl)
so the caller can mark the ablated tail (= free deploy-mode samples)."""
n_abl = round(n * cfg.rollout_ablate_frac) if cfg.intervention in ("route", "route2") else 0
n_abl = round(n * cfg.rollout_ablate_frac) if cfg.intervention in ("route", "routeV") else 0
parts = []
if n - n_abl > 0:
parts.append(model.generate(**enc, generation_config=gen_cfg,
@@ -773,7 +777,7 @@ def main(cfg: Config) -> int:
rollout_log_path.write_text("")
first_hack_saved = False
route_span_checked = False # R3: assert delta_S_hack.grad in span(V) once
# route2-grad routing band is built from the pairs at v_grad extraction time
# routeV-grad routing band is built from the pairs at v_grad extraction time
# (route_band[name] = (lower, upper)); see route_band_edges. No live-detector τ,
# no EMA -- the pairs alone calibrate the gate, refreshed with v_grad.
last_gen_sample = None # first student rollout of the latest step (for collapse inspection)
@@ -806,16 +810,22 @@ def main(cfg: Config) -> int:
# dropped from the per-step table as redundant; reconstruct here).
hr = sum(r["hack_s"][0] + r["hack_t"][0] for r in rows) / max(1, n_gens)
pr = sum(r["gt_s"][0] + r["gt_t"][0] for r in rows) / max(1, n_gens)
# Save δS only (not δS_hack). For route this is exactly the
# deployment adapter: the quarantine knob is ablated at eval, so dropping
# it here == the model you'd ship.
# train.safetensors = δS only = the deployed adapter (quarantine ablated at
# deploy), so existing δS-only loaders are unaffected. δS_hack (the quarantine
# knob) goes to a sibling _hack.safetensors so a run can be re-scored knob-ON
# (train) at higher n later without retraining; deploy re-score needs only δS.
_ckpt = path or ckpt_path
tensors = {n: info["delta_S"].detach().cpu().contiguous()
for n, info in wrappers.items()}
save_file(tensors, str(path or ckpt_path), metadata={
save_file(tensors, str(_ckpt), metadata={
"model": model_name, "dtype": "bf16", "step": str(len(rows)),
"hack_rate": f"{hr:.6f}", "pass_rate": f"{pr:.6f}",
"rows": json.dumps(rows), "cfg": json.dumps(vars(cfg), default=str),
})
hack_tensors = {n: info["delta_S_hack"].detach().cpu().contiguous()
for n, info in wrappers.items()}
save_file(hack_tensors, str(_ckpt.with_name(_ckpt.stem + "_hack.safetensors")),
metadata={"model": model_name, "step": str(len(rows))})
# disable=None: auto-disable the bar when stdout is NOT a tty (pueue, pipes,
# file redirects). In those contexts every per-step `logger.info(step_logger.row)`
@@ -829,7 +839,7 @@ def main(cfg: Config) -> int:
# ── training loop: generate -> grade -> backward -> project -> step ──
for step in pbar:
# Teacher-off curriculum: seed hacks via the teacher pool for the first N
# steps, then cut to pure on-policy (G_t=0) so we test whether route2 holds
# steps, then cut to pure on-policy (G_t=0) so we test whether routeV holds
# the suppression once the teacher crutch is gone. Monotonic: stays off.
if cfg.teacher_off_step is not None and step >= cfg.teacher_off_step and G_t > 0:
logger.info(f"teacher-off curriculum: step {step} >= {cfg.teacher_off_step} "
@@ -864,12 +874,12 @@ def main(cfg: Config) -> int:
# what the projection + optimizer step ultimately sees.
step_grad_s: dict[str, torch.Tensor] = {}
step_grad_t: dict[str, torch.Tensor] = {}
# route2: the flagged rollouts' δS-grad contribution, accumulated per module
# routeV: the flagged rollouts' δS-grad contribution, accumulated per module
# across prompts, parked into δS_hack.grad at injection (the quarantine,
# deleted at deploy). Mirrors how proj.py parks route's removed component.
step_grad_hack: dict[str, torch.Tensor] = {}
# route2: recover the per-rollout δS grad from the gate (c.grad = δS * g_b),
# routeV: recover the per-rollout δS grad from the gate (c.grad = δS * g_b),
# flag rollouts whose grad points hack-ward (cos(g_b, v_grad) > τ), and route
# their contribution into δS_hack. Only axes where δS has moved (|δS| > GATE_EPS)
# carry a reliable per-rollout split; near-zero axes keep the full grad, so
@@ -880,7 +890,7 @@ def main(cfg: Config) -> int:
step_hkgap: list[float] = [] # band width upper-lower (pair separation; ~0 = random/degenerate)
step_resid: list[float] = [] # cos(δS.grad AFTER routing, v_grad): hack-ward leak into deployed knob
def _route2_grad_filter(info, n_rollouts: int) -> torch.Tensor:
def _routeV_grad_filter(info, n_rollouts: int) -> torch.Tensor:
g = info["delta_S"].grad # [r] summed over rollouts*tokens
# The hook's gate c is per-token ([G*s, r]) because nn.Linear sees a
# flattened batch. reshape [G*s, r] -> [G, s, r]. Pad tokens carry ~0 grad
@@ -897,10 +907,10 @@ def main(cfg: Config) -> int:
# per-rollout (default): sum tokens first -> one cos/f per rollout. Denoises
# the cos sign (a clean rollout's tokens scatter ~50% over cos>0; the
# token-sum points reliably clean-ward) and matches GRPO's per-rollout adv.
# per-token (route2_per_token): one cos/f per token -- finer but noisier.
# per-token (routeV_per_token): one cos/f per token -- finer but noisier.
lower, upper = route_band[name]
band = max(upper - lower, 1e-6)
if cfg.route2_per_token:
if cfg.routeV_per_token:
g_u = torch.where(reliable, cg_full / dS_safe, torch.zeros_like(cg_full)) # [G, s, r]
cos_u = (g_u @ vg) / g_u.norm(dim=2).clamp_min(1e-12) # [G, s]
f = ((cos_u - lower) / band).clamp(0.0, 1.0) # [G, s]
@@ -933,10 +943,10 @@ def main(cfg: Config) -> int:
# On split steps: 2 backwards per prompt, populates step_grad_s/_t.
# On skipped steps: 1 combined backward, step_grad_s/_t stay empty and
# cos_pre_s/cos_pre_t go to NaN (mean_cos_pre_from_grads returns NaN on empty dict).
# route2 has no v_hack so cos_pre is NaN regardless: force the single combined
# routeV has no v_hack so cos_pre is NaN regardless: force the single combined
# backward (the split would just double cost). The grad-mask reads its
# per-rollout gate from that one backward.
split_this_step = (step % cfg.cos_pre_split_every == 0) and not is_route2
split_this_step = (step % cfg.cos_pre_split_every == 0) and not is_routeV
# Phase timers (per-step cumulative, seconds). Each GPU phase ends in a
# CPU-blocking op (decode / .item()), so perf_counter is sync-accurate
# without explicit cuda.synchronize. Tells us whether wall-time is
@@ -1239,11 +1249,11 @@ def main(cfg: Config) -> int:
g = info["delta_S"].grad
if g is None:
continue
# route2 routes here: split each rollout's δS.grad by its cosine to
# routeV routes here: split each rollout's δS.grad by its cosine to
# v_grad against the pair-calibrated band, park the routed fraction in
# δS_hack (via step_grad_hack in the filter).
if is_route2:
g = _route2_grad_filter(info, merged.shape[0])
if is_routeV:
g = _routeV_grad_filter(info, merged.shape[0])
step_grad_s[name] = (step_grad_s[name] + g.detach().clone()
if name in step_grad_s
else g.detach().clone())
@@ -1264,7 +1274,7 @@ def main(cfg: Config) -> int:
info["delta_S"].grad = gs
else:
info["delta_S"].grad = gs + gt
# route2: park the flagged rollouts' contribution into δS_hack.grad (its own
# routeV: park the flagged rollouts' contribution into δS_hack.grad (its own
# forward-path grad was wiped by the per-prompt zero_grad; we impose the routed
# grad here, like proj.py's route).
for name, g in step_grad_hack.items():
@@ -1281,10 +1291,10 @@ def main(cfg: Config) -> int:
diag = {"mean_cos_pre": float("nan"), "mean_cos_post": float("nan"),
"frac_fired": float("nan"), "mean_cos_pre_s": float("nan"),
"mean_cos_pre_t": float("nan")}
# route2: mean routed fraction f (mean over modules*prompts) -- also the
# routeV: mean routed fraction f (mean over modules*prompts) -- also the
# frout streaming column; logged here too for the no-v_hack diag branch.
if is_route2 and step_flagged:
logger.debug(f"route2 routed frac f (mean over modules*prompts): "
if is_routeV and step_flagged:
logger.debug(f"routeV routed frac f (mean over modules*prompts): "
f"{sum(step_flagged)/len(step_flagged):+.3f}")
else:
if split_this_step:
@@ -1349,10 +1359,10 @@ def main(cfg: Config) -> int:
# saved cache and overwrite the in-memory v_hack dict.
refr = "-" # set to "mod/axes" below if a refresh fires; rendered in the per-step row
do_refresh = cfg.vhack_refresh_every > 0 and (step + 1) % cfg.vhack_refresh_every == 0
if do_refresh and is_route2 and cfg.route2_random_v_seed is not None:
if do_refresh and is_routeV and cfg.routeV_random_v_seed is not None:
do_refresh = False # keep the one fixed Haar draw; re-extracting would replace it
if do_refresh and is_route2:
# route2 v_grad refresh: re-extract against the CURRENT model so the
if do_refresh and is_routeV:
# routeV v_grad refresh: re-extract against the CURRENT model so the
# routing direction tracks where hacks separate now, not at step 0.
# Without this the frozen direction goes stale -- cin_t decays to cin_s
# within ~6 steps. Same MASK_PAIRS (the weak
@@ -1370,7 +1380,7 @@ def main(cfg: Config) -> int:
model, tok, wrappers, MASK_PAIRS,
top_k=1, tau_axis=0.0, n_heldout=2, device=device,
)
for name in wrappers: # update in place so _route2_grad_filter's closure sees it
for name in wrappers: # update in place so _routeV_grad_filter's closure sees it
d = (raw_grads[f"hack/{name}"] - raw_grads[f"clean/{name}"]).mean(0)
v_grad[name] = (d / d.norm().clamp_min(1e-12)).to(device)
route_band = route_band_edges(raw_grads, v_grad, device) # rebuild band on the fresh v_grad
@@ -1447,7 +1457,7 @@ def main(cfg: Config) -> int:
# ── periodic DEPLOY-eval (EVERY arm) -- the apples-to-apples curve ──
# Eval the DEPLOYED model on a fixed eval subset with gen_cfg_eval (n=64,
# T=0.7), every eval_ablate_every steps. route/route2: deploy = quarantine
# T=0.7), every eval_ablate_every steps. route/routeV: deploy = quarantine
# knob zeroed (ablate_quarantine), and the claim is this hacks far less than
# the training-time model (per-step hack_s, knob still on). vanilla/erase: no
# quarantine, so deploy == the trained model -- eval it directly. Running the
@@ -1458,7 +1468,7 @@ def main(cfg: Config) -> int:
if cfg.eval_ablate_every > 0 and (step % cfg.eval_ablate_every == 0 or step == steps - 1):
_was_training = model.training
model.eval()
is_route = cfg.intervention in ("route", "route2")
is_route = cfg.intervention in ("route", "routeV")
with (ablate_quarantine(wrappers) if is_route else nullcontext()):
ev = eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg_eval, device, max_new)
hack_deploy, solve_deploy = ev["hack"], ev["solve"]
@@ -1645,7 +1655,7 @@ def main(cfg: Config) -> int:
logger.error(
f"DIVERGED at step {step}: lp_t={lp_t_mean:.1f} (ppl_t={ppl_t:.0e}), {lp_t_best - lp_t_mean:.1f} "
f"nats below best {lp_t_best:.1f}, for {diverged_steps} steps -- policy collapsed "
f"(gn={gn:.1f}). Aborting to save GPU. Likely lr too high (route2: lower --route2-quar-lr-scale).")
f"(gn={gn:.1f}). Aborting to save GPU. Likely lr too high (routeV: lower --routeV-quar-lr-scale).")
if last_gen_sample:
_s, _r = last_gen_sample
logger.error(f"--- last student gen (step {_s}, reward={_r['reward']:+.2f}) ---\n"
@@ -1722,13 +1732,13 @@ def main(cfg: Config) -> int:
dsh_norm = float(sum(info["delta_S_hack"].data.float().pow(2).sum().item()
for info in wrappers.values()) ** 0.5)
logger.info(f"||delta_S_hack|| = {dsh_norm:.4f} "
f"(SHOULD: >0 for route/route2, ==0 for none/erase; ELSE routing broke)")
if cfg.intervention in ("route", "route2") and cfg.route2_random_v_seed is None:
f"(SHOULD: >0 for route/routeV, ==0 for none/erase; ELSE routing broke)")
if cfg.intervention in ("route", "routeV") and cfg.routeV_random_v_seed is None:
assert dsh_norm > 0.0, f"{cfg.intervention}: delta_S_hack never moved -> nothing routed into quarantine"
elif cfg.route2_random_v_seed is not None and dsh_norm == 0.0:
elif cfg.routeV_random_v_seed is not None and dsh_norm == 0.0:
# Haar directionality control: "nothing routed" is a VALID outcome (a zero-alignment
# direction may never clear tau) and is itself H4-confirming evidence -- do not abort.
logger.warning("route2 Haar control: ||delta_S_hack||==0 -> the random direction routed "
logger.warning("routeV Haar control: ||delta_S_hack||==0 -> the random direction routed "
"NOTHING. This is a real result (favours H4: alignment needed), not a failure.")
# Last training generation -- a fast eyeball for coherence before the eval
@@ -1744,17 +1754,28 @@ def main(cfg: Config) -> int:
# ── final eval + BLUF ──
# Final per-mode train-vs-deploy eval -- run for EVERY arm on the SAME fixed
# eval subset so the all-arms overlay reads identical numbers. For route/route2
# eval subset so the all-arms overlay reads identical numbers. For route/routeV
# this is the absorption test: TRAIN keeps the quarantine knob on (still hacks),
# DEPLOY deletes it (the shipped model). SHOULD: deploy hack < train hack at
# preserved solve => the quarantine absorbed the cheat. vanilla/erase have no
# quarantine, so the deployed model IS the trained model (deploy == train, one eval).
model.eval()
ev_train = eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg_eval, device, max_new)
has_quarantine = cfg.intervention in ("route", "route2")
# Paper-grade final eval: eval_n_prompts_final distinct prompts (>> the periodic
# eval_n_prompts curve), and a FIXED gen seed before each pass so every arm/seed sees
# common random numbers -> cross-arm deltas reflect the intervention, not eval sampling
# noise (gen is do_sample T=0.7, otherwise unseeded; the periodic curve stays light +
# unseeded and gets smoothed). Capped at the available pool size.
EVAL_GEN_SEED = 12345
eval_idxs_final = list(range(min(cfg.eval_n_prompts_final, len(problems))))
logger.info(f"FINAL EVAL: {len(eval_idxs_final)} distinct prompts x G={group} = "
f"{len(eval_idxs_final) * group} completions (periodic curve used {len(eval_idxs)})")
torch.manual_seed(EVAL_GEN_SEED)
ev_train = eval_hack_solve(model, tok, problems, eval_idxs_final, gen_cfg_eval, device, max_new)
has_quarantine = cfg.intervention in ("route", "routeV")
if has_quarantine:
with ablate_quarantine(wrappers):
ev_deploy = eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg_eval, device, max_new)
torch.manual_seed(EVAL_GEN_SEED)
ev_deploy = eval_hack_solve(model, tok, problems, eval_idxs_final, gen_cfg_eval, device, max_new)
else:
ev_deploy = ev_train
logger.info(