From 4359dc53a8083ef512d22b65ac9413afc9c034e6 Mon Sep 17 00:00:00 2001 From: wassname Date: Sun, 31 May 2026 10:16:13 +0000 Subject: [PATCH] feat: route2 distinct-basis quarantine + per-sample act-mask detach-route Adds intervention=route2: a LoRA quarantine (A_q,B_q) with its own basis, always summed into the forward, plus a per-sample activation-cosine mask that detaches the kept adapter for flagged samples. Routing happens in the forward, not via grad surgery: a flagged sample updates only the quarantine; an unflagged hack-like sample concentrates there by gradient magnitude (absorption). Deploy zeroes A_q,B_q. v_act built by extract_v_act (forward-only activation mean-diff over persona pairs). Fixes the per-prompt zero_grad wiping quarantine grads before opt.step. scripts/make_random_vhack.py = the random-V route control. vhack_refresh_every default 0->5 (0 is ablation-only). Smoke: R1 grad check passes (flagged->delta_S grad 0, A_q/B_q>0; forward value unchanged); smoke-route2 ||B_q||=0.109, deploy eval + asserts pass. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com> --- justfile | 11 +++ scripts/make_random_vhack.py | 52 ++++++++++ src/projected_grpo/antipasto.py | 84 +++++++++++++---- src/projected_grpo/extract_vhack_grad.py | 70 ++++++++++++++ src/projected_grpo/train.py | 115 ++++++++++++++++++----- 5 files changed, 294 insertions(+), 38 deletions(-) create mode 100644 scripts/make_random_vhack.py diff --git a/justfile b/justfile index db11286..247363d 100644 --- a/justfile +++ b/justfile @@ -44,6 +44,17 @@ 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): distinct-basis quarantine (A_q,B_q) + per-sample +# act-mask detach-route in the FORWARD. Fires extract_v_act, the quarantine +# optimizer params, the act-mask cosine, the B_q-moved assert, and the deploy +# ablation (B_q zeroed). tau=-1 forces all-flagged so the seed path is exercised +# on tiny inputs (real runs calibrate tau against R3 fire-ratio). +smoke-route2 *ARGS: + BEARTYPE=1 CUDA_VISIBLE_DEVICES= {{ TRAIN }} smoke --intervention=route2 \ + --teacher-pool-dir=out/pools/teacher_pool --mix-ratio=0.5 \ + --route2-quarantine-rank=8 --route2-tau=-1.0 \ + --eval-ablate-every=10 --eval-n-prompts=2 {{ ARGS }} + # Run smoke twice: first warms the v_hack cache (cache-miss path), second hits # the cache (cache-hit path). Catches scope/save bugs that only manifest in one. smoke-both: diff --git a/scripts/make_random_vhack.py b/scripts/make_random_vhack.py new file mode 100644 index 0000000..c4fc636 --- /dev/null +++ b/scripts/make_random_vhack.py @@ -0,0 +1,52 @@ +"""Haar-random orthonormal V control for the route arm (#157). + +Decisive discriminator (gpt-5.5 review Q5, design brainstorm): is route's +deploy-hack-drop + solve-jump *directional* (the extracted v_hack subspace +matters) or just *adapter regularization* (deleting any rank-k quarantine of +matched norm reverts toward base)? If a random orthonormal V of identical +per-module shape/rank/singular-values reproduces run 31's effect, the additive +result was an ablation artifact, not directional specificity. + +We keep `_sv/*` (singular values) and metadata byte-identical so the noise-floor +filter survives the same modules and per-module scaling is matched. Only the +[k, r] direction rows are replaced with Haar-random orthonormal rows. +""" +import sys +from pathlib import Path + +import torch +from loguru import logger +from safetensors import safe_open +from safetensors.torch import save_file + +SRC = Path("out/vhack/v_hack_pairset_prog_wide.safetensors") # run 31's v_hack +DST = Path("out/vhack/v_hack_pairset_prog_wide_randomV.safetensors") +SEED = 157 # the task number; fixed so the control is reproducible + + +def haar_orthonormal_rows(k: int, r: int, generator: torch.Generator) -> torch.Tensor: + """k orthonormal rows in r-dim space, uniform over the Stiefel manifold (QR of Gaussian).""" + g = torch.randn(r, k, generator=generator, dtype=torch.float32) + q, _ = torch.linalg.qr(g) # q: [r, k], orthonormal columns + return q.mT.contiguous() # [k, r], orthonormal rows + + +def main(): + gen = torch.Generator().manual_seed(SEED) + out = {} + with safe_open(str(SRC), "pt") as f: + metadata = f.metadata() + for name in f.keys(): + t = f.get_tensor(name) + if name.startswith("_sv/"): + out[name] = t # keep singular values identical -> matched norm + noise floor + else: + k, r = t.shape # [k directions, r SVD coords] + out[name] = haar_orthonormal_rows(k, r, gen).to(t.dtype) + save_file(out, str(DST), metadata=metadata) + n_dir = sum(1 for k in out if not k.startswith("_sv/")) + logger.info(f"wrote {DST} | {n_dir} random-V modules | seed={SEED} | metadata={metadata}") + + +if __name__ == "__main__": + main() diff --git a/src/projected_grpo/antipasto.py b/src/projected_grpo/antipasto.py index f9f2540..ced4771 100644 --- a/src/projected_grpo/antipasto.py +++ b/src/projected_grpo/antipasto.py @@ -65,24 +65,48 @@ def is_target(name: str) -> bool: def _delta_hook(layer: nn.Linear, args: tuple, y: Tensor) -> Tensor: - """Add U @ (delta_S * (Vh @ x)) to y. Cast delta to y.dtype to match. + """Add the AntiPaSTO delta to y. Two regimes, branched on whether a route2 + quarantine (`_antipasto_A_q`) was attached: - Note: hook input tuple is (x,) for nn.Linear forward. + SHARED-BASIS (none/erase/route -- A_q is None): + y += U @ ((delta_S + delta_S_hack) * (Vh @ x)) + delta_S = main knob; delta_S_hack = shared-basis routing quarantine (proj.py + parks the removed hack-ward grad here). Both 0 at init -> identity. + + DISTINCT-BASIS route2 (A_q set): + a = Vh @ x # [..., r] activation in SVD coords + kept= U @ (delta_S * a) # the deployed adapter (delta_S_hack stays 0) + m = cos(a, v_act) > tau # [...] per-sample/token persona flag + kept= where(m, kept.detach(), kept)# flagged -> delta_S gets no grad + quar= B_q @ (A_q @ x) # distinct-basis quarantine, always summed + y += kept + quar + Routing is done in the forward, not the gradient: a flagged sample detaches + the kept path so only the quarantine learns it; an unflagged sample updates + both, so a hack-like one concentrates in the quarantine where its gradient is + larger (absorption, Cloud 2410.04332 / SGTM 2512.05648). Deploy zeroes A_q,B_q. """ (x,) = args Vh = layer._antipasto_Vh # [r, d_in] U = layer._antipasto_U # [d_out, r] - # delta_S = main trainable knob; delta_S_hack = routing quarantine knob. - # Forward uses their sum, so a routed (hack-ward) update via delta_S_hack - # still moves the training-time model; zeroing delta_S_hack at eval ablates - # the routed capability. Both are 0 at init -> identity preserved. delta_S = layer._antipasto_delta_S # [r] delta_S_hack = layer._antipasto_delta_S_hack # [r] - # match input dtype for matmul - h = torch.nn.functional.linear(x, Vh) # [..., r] - h = h * (delta_S + delta_S_hack).to(h.dtype) # [..., r] - delta = torch.nn.functional.linear(h, U) # [..., d_out] - return y + delta.to(y.dtype) + A_q = layer._antipasto_A_q # None or [k_q, d_in] + + a = torch.nn.functional.linear(x, Vh) # [..., r] + if A_q is None: + h = a * (delta_S + delta_S_hack).to(a.dtype) + return y + torch.nn.functional.linear(h, U).to(y.dtype) + + # --- route2: distinct-basis quarantine + per-sample act-mask detach-route --- + kept = torch.nn.functional.linear(a * delta_S.to(a.dtype), U) # [..., d_out] + v_act = layer._antipasto_v_act # [r] unit, hack-ward, in Vh coords + # per-position cosine of the SVD-coord activation with the persona direction + cos = (a @ v_act) / (a.norm(dim=-1).clamp_min(1e-6) * v_act.norm().clamp_min(1e-6)) + m = cos > layer._antipasto_tau # [...] bool + kept = torch.where(m.unsqueeze(-1), kept.detach(), kept) + B_q = layer._antipasto_B_q # [d_out, k_q] + quar = torch.nn.functional.linear(torch.nn.functional.linear(x, A_q), B_q) # [..., d_out] + return y + (kept + quar).to(y.dtype) def wrap_model_with_antipasto( @@ -90,12 +114,19 @@ def wrap_model_with_antipasto( model_name: str, cache_root: Path = Path("svd_cache"), svd_device: torch.device | str = "cuda", + quarantine_rank: int | None = None, + route2_tau: float = 0.0, ) -> dict[str, dict]: """Attach AntiPaSTO hooks to every target nn.Linear in `model` (in place). Returns dict[qualified_name -> dict(layer, delta_S, handle, r)]. Frozen U/Vh stored on the layer as buffers `_antipasto_{U,Vh}` in the layer's native dtype. delta_S kept in fp32 (tiny, ~r per module). + + `quarantine_rank` (route2 only): if set, also attach a DISTINCT-basis LoRA + quarantine per module -- A_q [k_q, d_in] (kaiming), B_q [d_out, k_q] (zeros so + quar=0 at init), and a unit `v_act` buffer [r] (filled by the act-mask + extraction in train.py) + scalar `route2_tau`. None -> shared-basis path only. """ svd_device_t = torch.device(svd_device) if isinstance(svd_device, str) else svd_device safe = model_name.replace("/", "__") @@ -123,12 +154,28 @@ def wrap_model_with_antipasto( delta_S_hack = nn.Parameter(torch.zeros(r, device=dev, dtype=torch.float32)) linear.register_parameter("_antipasto_delta_S", delta_S) linear.register_parameter("_antipasto_delta_S_hack", delta_S_hack) - handle = linear.register_forward_hook(_delta_hook) - out[name] = {"layer": linear, "delta_S": delta_S, - "delta_S_hack": delta_S_hack, "handle": handle, "r": r} + info = {"layer": linear, "delta_S": delta_S, + "delta_S_hack": delta_S_hack, "handle": None, "r": r} + if quarantine_rank is None: + linear._antipasto_A_q = None # plain attr -> shared-basis hook branch + else: + k_q = min(quarantine_rank, r) + A_q = nn.Parameter(torch.empty(k_q, d_in, device=dev, dtype=torch.float32)) + nn.init.kaiming_uniform_(A_q, a=5 ** 0.5) # LoRA-A init + B_q = nn.Parameter(torch.zeros(d_out, k_q, device=dev, dtype=torch.float32)) # quar=0 at init + linear.register_parameter("_antipasto_A_q", A_q) + linear.register_parameter("_antipasto_B_q", B_q) + linear.register_buffer("_antipasto_v_act", + torch.zeros(r, device=dev, dtype=torch.float32), persistent=True) + linear._antipasto_tau = route2_tau + info["A_q"], info["B_q"], info["k_q"] = A_q, B_q, k_q + info["handle"] = linear.register_forward_hook(_delta_hook) + out[name] = info - # freeze everything except the two AntiPaSTO knobs (delta_S, delta_S_hack) - trainable = ("_antipasto_delta_S", "_antipasto_delta_S_hack") + # freeze everything except the AntiPaSTO knobs. A_q/B_q (route2) are trainable + # too; v_act is a buffer (not a param) so it stays frozen by construction. + trainable = ("_antipasto_delta_S", "_antipasto_delta_S_hack", + "_antipasto_A_q", "_antipasto_B_q") for n, p in model.named_parameters(): if not n.endswith(trainable): p.requires_grad_(False) @@ -142,6 +189,9 @@ def detach_antipasto(model: nn.Module, attached: dict) -> None: for attr in ("_antipasto_U", "_antipasto_Vh"): if attr in layer._buffers: del layer._buffers[attr] - for attr in ("_antipasto_delta_S", "_antipasto_delta_S_hack"): + for attr in ("_antipasto_delta_S", "_antipasto_delta_S_hack", + "_antipasto_A_q", "_antipasto_B_q"): if attr in layer._parameters: del layer._parameters[attr] + if "_antipasto_v_act" in layer._buffers: + del layer._buffers["_antipasto_v_act"] diff --git a/src/projected_grpo/extract_vhack_grad.py b/src/projected_grpo/extract_vhack_grad.py index 118c092..39c585e 100644 --- a/src/projected_grpo/extract_vhack_grad.py +++ b/src/projected_grpo/extract_vhack_grad.py @@ -218,6 +218,76 @@ def extract_v_hack( return v_hack, v_sv, raw_grads, rows +@torch.no_grad() +def extract_v_act( + model, + tokenizer, + wrappers: dict, + pairs: list, + n_heldout: int, + device, +) -> dict[str, Float[torch.Tensor, "r"]]: + """Activation-space persona direction per module for route2 Arm B. + + Forward-only analogue of extract_v_hack: capture the SVD-coord activation + a = Vh @ x at each wrapped Linear, averaged over completion tokens, and return + the unit hack-minus-clean direction per module (in R^r, the same space the + route2 forward computes cos(a, v_act) in). No backward -> cheaper than v_hack. + + Oriented hack-ward by construction (mean_hack - mean_clean). The k=1 mean-diff; + we deliberately do NOT SVD here -- the mask only needs one direction, and the + sign/orientation is unambiguous from the paired mean difference. + """ + train_pairs = pairs[:-n_heldout] if n_heldout > 0 else pairs + sums: dict[str, dict] = { + name: {"hack": torch.zeros(info["r"]), "clean": torch.zeros(info["r"]), + "n_hack": 0, "n_clean": 0} + for name, info in wrappers.items() + } + captured: dict[str, torch.Tensor] = {} + handles = [] + for name, info in wrappers.items(): + def mk(nm): + def pre(_mod, args): + captured[nm] = args[0].detach() # [1, L, d_in] input to the Linear + return pre + handles.append(info["layer"].register_forward_pre_hook(mk(name))) + + try: + for pi, pair in enumerate(train_pairs): + for label, completion in (("hack", pair.hack), ("clean", pair.clean)): + prompt_ids = tokenizer(pair.prompt, return_tensors="pt").input_ids.to(device) + full_ids = tokenizer(pair.prompt + completion, return_tensors="pt").input_ids.to(device) + n_prompt = prompt_ids.shape[1] + captured.clear() + model(full_ids) + for name, info in wrappers.items(): + x = captured[name].float() # [1, L, d_in] + Vh = info["layer"]._antipasto_Vh.float() # [r, d_in] + a = torch.nn.functional.linear(x, Vh) # [1, L, r] + a_comp = a[:, n_prompt:, :].mean(dim=(0, 1)).cpu() # [r] mean over completion toks + sums[name][label] += a_comp + sums[name]["n_" + label] += 1 + if (pi + 1) % 5 == 0: + logger.info(f" v_act pair {pi+1}/{len(train_pairs)}") + finally: + for h in handles: + h.remove() + + v_act: dict[str, torch.Tensor] = {} + n_zero = 0 + for name, s in sums.items(): + diff = s["hack"] / max(1, s["n_hack"]) - s["clean"] / max(1, s["n_clean"]) + nrm = diff.norm() + if nrm < 1e-12: + n_zero += 1 + v_act[name] = torch.zeros_like(diff) + else: + v_act[name] = (diff / nrm).contiguous() + logger.info(f"v_act: modules={len(v_act)} zero-||diff||={n_zero} (activation mean-diff, unit, hack-ward)") + return v_act + + def main(cfg: Config) -> int: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = resolve_dtype(cfg.dtype) diff --git a/src/projected_grpo/train.py b/src/projected_grpo/train.py index 8cd0f63..124883e 100644 --- a/src/projected_grpo/train.py +++ b/src/projected_grpo/train.py @@ -140,9 +140,17 @@ class Config: # erase = today's projection: subtract the hack-ward component from delta_S # route = park the hack-ward component in the delta_S_hack quarantine knob # (Gradient Routing, Cloud 2410.04332); ablate it at eval. + # route2 = DISTINCT-basis quarantine (A_q,B_q LoRA) + per-sample act-mask + # detach-route in the FORWARD (no proj.py grad surgery). Absorption + # design (SGTM 2512.05648); see docs/spec/20260531_routing_v2_*.md. # Replaces the old `arm` flag (vanilla/projected); `arm` survives as a derived # display name (see property below) so log/run-id formatting is unchanged. - intervention: Literal["none", "erase", "route"] = "erase" + intervention: Literal["none", "erase", "route", "route2"] = "erase" + # route2-only: quarantine LoRA rank (per module, capped at r) and the act-mask + # cosine threshold. tau is the routing-fraction knob (higher -> route less -> + # less starvation, weaker seed). Calibrate against R3 (fire >2x on hack vs clean). + route2_quarantine_rank: int = 16 + route2_tau: float = 0.0 # Scale-dependent knobs — every preset must set these to a real value; # subclasses below override the defaults. model: str = "Qwen/Qwen3-4B" @@ -187,12 +195,11 @@ class Config: # global threshold get filtered out entirely (projection skips them — they # didn't carry hack signal anyway). 0 = no filter. v_hack_drop_bottom_frac: float = 0.25 - # Online refresh: every N optimizer steps, re-run extract_v_hack against the - # current (delta_S-modified) model using the same hand-crafted PAIRS. Motivated - # by Goal-1 negative result (2026-05-28 c) where v_hack appeared to lose - # discriminative power as the student drifted. 0 = off (load once at start - # and freeze). Refresh cost ~14*2 backwards on Qwen3-4B ~ 1-2 min wall. - vhack_refresh_every: int = 0 + # Online refresh: every N optimizer steps, re-extract v_hack against the + # current (delta_S-modified) model so it tracks the student's drifting hack + # subspace rather than the step-0 one. 0 = freeze at load (ablation only). + # Refresh cost ~14*2 backwards on Qwen3-4B ~ 1-2 min wall. + vhack_refresh_every: int = 5 # Route eval-time ablation: every N steps (and at the end), zero delta_S_hack # and eval hack/solve on a fixed prompt subset -> the `hack_deploy`/`solve_deploy` # columns. This is the series the dynamics plot uses for route, because the @@ -245,7 +252,8 @@ class Config: def arm(self) -> str: """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"}[self.intervention] + return {"none": "vanilla", "erase": "projected", + "route": "routing", "route2": "routing2"}[self.intervention] @dataclass(kw_only=True) @@ -521,17 +529,23 @@ def ref_logprobs_via_zero_delta( @contextmanager def ablate_quarantine(wrappers: dict): - """Zero delta_S_hack for the duration -- the eval-time ablation of the - routed hack capability. Save -> zero -> (eval) -> restore. The route arm's - deployment model IS this ablated state.""" + """Zero the routing quarantine for the duration -- the eval-time ablation of + the routed hack capability. Save -> zero -> (eval) -> restore. The route arm's + deployment model IS this ablated state. Covers BOTH quarantines: delta_S_hack + (shared-basis route) and B_q (distinct-basis route2; zeroing B_q -> quar=0).""" saved = {n: info["delta_S_hack"].data.clone() for n, info in wrappers.items()} + saved_Bq = {n: info["B_q"].data.clone() for n, info in wrappers.items() if "B_q" in info} for info in wrappers.values(): info["delta_S_hack"].data.zero_() + if "B_q" in info: + info["B_q"].data.zero_() try: yield finally: for n, info in wrappers.items(): info["delta_S_hack"].data.copy_(saved[n]) + if "B_q" in info: + info["B_q"].data.copy_(saved_Bq[n]) @torch.no_grad() @@ -729,26 +743,55 @@ def main(cfg: Config) -> int: # call below: True for autoregressive decode, False for the single loss forwards. model.config.use_cache = False - wrappers = wrap_model_with_antipasto(model, model_name, CACHE_ROOT, device) + is_route2 = cfg.intervention == "route2" + wrappers = wrap_model_with_antipasto( + model, model_name, CACHE_ROOT, device, + quarantine_rank=cfg.route2_quarantine_rank if is_route2 else None, + route2_tau=cfg.route2_tau, + ) # Both knobs are trainable params. delta_S_hack only ever gets a grad under # intervention=route (the routing split in proj.py); under none/erase its # grad stays None so AdamW skips it and it stays exactly 0 (forward adds 0, # so none/erase reproduce the pre-quarantine behaviour bit-for-bit). delta_params = [info["delta_S"] for info in wrappers.values()] delta_hack_params = [info["delta_S_hack"] for info in wrappers.values()] + # route2: the distinct-basis quarantine LoRA (A_q,B_q). Trained for all arms + # it exists in (route2 only); routing happens in the forward, not the grad. + quar_params = ([info["A_q"] for info in wrappers.values()] + + [info["B_q"] for info in wrappers.values()]) if is_route2 else [] 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, route-only)") + f"(+{sum(p.numel() for p in delta_hack_params):,} delta_S_hack, route-only" + f"{f'; +{sum(p.numel() for p in quar_params):,} A_q/B_q route2' if is_route2 else ''})") # v_hack: the hack-direction subspace the erase/route arms project against. # VANILLA (intervention=none) is a pure GRPO baseline and ignores v_hack # entirely -- loading it there only to print a cos_pre diagnostic was misleading # (and could trigger a needless ~5-min extraction). The cin/cout columns are # hidden on vanilla, so v_hack=None just means "no subspace machinery". - if cfg.intervention == "none": - if cfg.v_hack_path is not None: + if cfg.intervention in ("none", "route2"): + 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 + v_hack = None # route2 routes in the FORWARD via v_act, not via grad surgery + if is_route2: + # Build the activation-space persona direction (Arm B mask source) from + # the SAME persona pairs the grad extract uses, then load it into each + # module's _antipasto_v_act buffer. No oracle: pairs are the weak detector. + from .extract_vhack_grad import extract_v_act + if cfg.vhack_pairs_path is not None: + from .pairs_from_pool import load_pairs_json + VACT_PAIRS = load_pairs_json(cfg.vhack_pairs_path) + logger.info(f"v_act pairs: pool-derived ({cfg.vhack_pairs_path}) -> {len(VACT_PAIRS)} pairs") + else: + from .pairs import PAIRS as VACT_PAIRS + logger.info(f"v_act pairs: hand-crafted PAIRS -> {len(VACT_PAIRS)} pairs") + model.eval() + v_act = extract_v_act(model, tok, wrappers, VACT_PAIRS, n_heldout=2, device=device) + model.train() + for name, info in wrappers.items(): + info["layer"]._antipasto_v_act.data.copy_(v_act[name].to(device)) + logger.info(f"route2: loaded v_act into {len(v_act)} modules; " + f"quarantine_rank={cfg.route2_quarantine_rank} tau={cfg.route2_tau}") else: # derive default path from model + extract_top_k unless overridden. # Auto-extract reuses the already-wrapped model — no second model load. @@ -853,7 +896,7 @@ def main(cfg: Config) -> int: ) opt = torch.optim.AdamW( - delta_params + delta_hack_params, lr=lr, weight_decay=cfg.weight_decay, + delta_params + delta_hack_params + quar_params, lr=lr, weight_decay=cfg.weight_decay, betas=(adam_beta1, adam_beta2), ) # Linear warmup over `warmup_frac * steps`, then cosine decay to 0 over the rest. @@ -1040,6 +1083,24 @@ 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 quarantine grads must survive the per-pass model.zero_grad (which + # exists to isolate delta_S's per-source grad). A_q/B_q need neither source + # split nor projection, just plain accumulation, so we stash and re-inject + # them exactly as delta_S is. Keyed ".". + step_grad_quar: dict[str, torch.Tensor] = {} + + def _stash_quar_grads(): + if not is_route2: + return + for name, info in wrappers.items(): + for sub in ("A_q", "B_q"): + p = info[sub] + if p.grad is None: + continue + key = f"{name}.{sub}" + step_grad_quar[key] = (step_grad_quar[key] + p.grad.detach().clone() + if key in step_grad_quar else p.grad.detach().clone()) + # Split backward into student/teacher only every cos_pre_split_every steps. # 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 @@ -1306,6 +1367,7 @@ def main(cfg: Config) -> int: step_grad_s[name] = (step_grad_s[name] + gs.detach().clone() if name in step_grad_s else gs.detach().clone()) + _stash_quar_grads() model.zero_grad(set_to_none=True) # Pass 2: teacher. loss_t.backward() @@ -1316,6 +1378,7 @@ def main(cfg: Config) -> int: step_grad_t[name] = (step_grad_t[name] + gt.detach().clone() if name in step_grad_t else gt.detach().clone()) + _stash_quar_grads() model.zero_grad(set_to_none=True) agg_loss += (loss_s + loss_t).item() else: @@ -1336,6 +1399,7 @@ def main(cfg: Config) -> int: step_grad_s[name] = (step_grad_s[name] + g.detach().clone() if name in step_grad_s else g.detach().clone()) + _stash_quar_grads() model.zero_grad(set_to_none=True) agg_loss += loss.item() t_fb += time.perf_counter() - _tfb @@ -1354,6 +1418,11 @@ def main(cfg: Config) -> int: info["delta_S"].grad = gs else: info["delta_S"].grad = gs + gt + # route2: re-inject the stashed quarantine grads (student+teacher summed + # across passes) so clip + opt.step move A_q/B_q. + for key, g in step_grad_quar.items(): + name, sub = key.rsplit(".", 1) + wrappers[name][sub].grad = g # Per-source cin: project student-only and teacher-only grads into v_hack # subspace. Discriminator: cos_pre_t > cos_pre_s on a clean base means v_hack @@ -1407,7 +1476,7 @@ def main(cfg: Config) -> int: # Clip over both knobs. For none/erase, delta_S_hack.grad is None so it's # ignored -> identical norm to before (R4). For route it bounds the # combined update (main + quarantine). - gn = float(torch.nn.utils.clip_grad_norm_(delta_params + delta_hack_params, cfg.grad_clip)) + gn = float(torch.nn.utils.clip_grad_norm_(delta_params + delta_hack_params + quar_params, cfg.grad_clip)) opt.step() sched.step() @@ -1416,7 +1485,7 @@ def main(cfg: Config) -> int: # than at step 0). Same PAIRS, same extract code; we just discard the # 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 - if cfg.vhack_refresh_every > 0 and (step + 1) % cfg.vhack_refresh_every == 0: + if v_hack is not None and cfg.vhack_refresh_every > 0 and (step + 1) % cfg.vhack_refresh_every == 0: from .extract_vhack_grad import extract_v_hack if cfg.vhack_pairs_path is not None: from .pairs_from_pool import load_pairs_json @@ -1488,7 +1557,7 @@ def main(cfg: Config) -> int: # which still hacks because training keeps the knob on). This is the curve the # plot uses for route. NaN on non-eval steps / non-route arms. hack_deploy = solve_deploy = float("nan") - if (cfg.intervention == "route" and cfg.eval_ablate_every > 0 + if (cfg.intervention in ("route", "route2") and cfg.eval_ablate_every > 0 and (step % cfg.eval_ablate_every == 0 or step == steps - 1)): _was_training = model.training model.eval() @@ -1697,12 +1766,16 @@ def main(cfg: Config) -> int: f"(SHOULD: >0 for route, ==0 for none/erase; ELSE routing broke)") if cfg.intervention == "route": assert dsh_norm > 0.0, "route: delta_S_hack never moved -> degenerated to erasure" + if is_route2: + bq_norm = sum(info["B_q"].data.norm().item() for info in wrappers.values()) + logger.info(f"||B_q|| sum = {bq_norm:.4f} (SHOULD: >0; ELSE quarantine never seeded)") + assert bq_norm > 0.0, "route2: B_q never moved -> quarantine never seeded (mask never fired?)" # Route: final training-vs-deployed eval -- the absorption test. TRAIN keeps # the quarantine knob on (training-time model, still hacks); DEPLOY deletes it # (the deployed model). SHOULD: deploy hack < train hack at preserved solve # => the quarantine knob absorbed the cheat. ELSE routing didn't localize it. - if cfg.intervention == "route": + if cfg.intervention in ("route", "route2"): model.eval() ev_train = eval_hack_solve(model, tok, problems, eval_idxs, gen_cfg_eval, device, max_new) with ablate_quarantine(wrappers):