From 3a39231d3e0421f900e7ce95122852b2950dea10 Mon Sep 17 00:00:00 2001 From: wassname Date: Mon, 1 Jun 2026 11:58:36 +0000 Subject: [PATCH] style(train): clean-repo voice pass 1 -- math notation + kill LLM tells MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Behavior-preserving (smoke + smoke-route2 exit 0, headline metrics identical: HACK_RATE=0.492 PASS=0.117 HACK_T=0.983). route2/delta_S_hack/hk_abl untouched. - GRPO loss reads like the equations: greek vars in the code itself (pol_logp->logπ, gen_logp->logπ_old, ref_logp->logπ_ref, ratio->ρ, adv->A, per_tok_loss->Lp, inline K3 KL). Scoped rename, no collisions. - Docstrings use unicode mirroring the math (δS, π_ref, ‖·‖, σ_R, Vᵀ, Sᵢ) not ASCII transliteration (delta_W, pi_ref, ||g||). - Em-dashes -> ASCII (grep -P '—' = 0). - Dropped LLM tells: past-reader war-stories (job 46, "was 0/16", dated journal refs, step-17 OOM anecdotes), jargon used before defined. - Module docstring + Config rewritten terse; one # ── section ── banner. Voice pass is not complete: section banners through main() and the cross-file decomposition (helpers -> antipasto.py/extract_vhack_grad.py/problems.py) follow in the next commits. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com> --- src/projected_grpo/train.py | 332 ++++++++++++++---------------------- 1 file changed, 131 insertions(+), 201 deletions(-) diff --git a/src/projected_grpo/train.py b/src/projected_grpo/train.py index 15f5084..2c6b86d 100644 --- a/src/projected_grpo/train.py +++ b/src/projected_grpo/train.py @@ -1,55 +1,28 @@ -"""Canonical training entry point: AntiPaSTO + GRPO (Dr.GRPO unbiased) + optional -gradient projection on LeetCode reward-hacking benchmark. +"""GRPO / Dr.GRPO loop with SVD-basis gradient projection on the LeetCode +reward-hacking benchmark. -Lineage (see spec.md §76-83): - - The inner GRPO_step (per_token_logps, ratio + clip + min, K3 KL, per-token - loss, completion mask) is a direct port of lsdefine/simple_GRPO's - `GRPO_step` in `grpo_vllm_one.py` (lines 64-95). - - The OUTER loop adopts simple_GRPO's `Q_batch_size` pattern (multiple - prompts per optimizer step, per-prompt GRPO advantage groups, grad - accumulation across prompts). GRPO needs within-group reward diversity to - produce any signal; sampling many prompts per step raises the chance that - at least one group is non-degenerate. simple_GRPO uses Q_batch_size=5; our - prompts_per_step is set per preset (grad-accum to the paper's effective batch). - - Deviations from simple_GRPO are deliberate, listed in spec.md: - 1. Loss normalization: Dr.GRPO unbiased (Liu et al. 2025, arXiv - 2503.20783) replaces simple_GRPO's `(R-mean)/std` + per-response-len - denominator. Drops two biases: - - length norm `1/|o_i|` (favors short correct, long incorrect) - - group-std norm `/std(R)` (overweights easy/hard questions) - Toggle via `--unbiased` (default on); flipping to False recovers - simple_GRPO's classic GRPO advantage normalization. - 2. Reference model: simple_GRPO runs a separate base model via an HTTP - `ref_server`. We use the AntiPaSTO `delta_S=0` zero-adapter trick - (W' = W + U diag(0) Vh = W exactly) — no second model loaded. - 3. Rollout: simple_GRPO uses vLLM in a separate process. We use HF - `model.generate` in-process. - 4. Adapter: simple_GRPO is full FT (with DeepSpeed ZeRO). Canonical - (ariahw/rl-rewardhacking) is LoRA r=32. We use AntiPaSTO full-rank - SVD adapter (the research artifact). + generate -> grade -> backward -> project -> step -Hyperparameters (lr, weight_decay, betas, warmup, cosine, beta=KL) are taken -from the closest-in-param-count reference: ariahw/rl-rewardhacking config.py -(LoRA r=32 on 4B ≈ 30M params) rather than simple_GRPO (full FT on 7B). See -docs/grpo_hyperparams.md. +Inner GRPO step ported from lsdefine/simple_GRPO grpo_vllm_one.py:64-95; the +outer loop accumulates grads over prompts_per_step prompts (simple_GRPO's +Q_batch_size), so at least one per-prompt group has reward variance. +Unbiased normalization: Dr.GRPO, Liu et al. 2025, arXiv:2503.20783 -- drop the +1/|oᵢ| length norm and the /σ_R group-std (--unbiased, on by default). -Reference-model term (`--beta`): Dr.GRPO argues beta=0 is fine for *reasoning* -RL with rule-based reward (no distributional-shift concern when reward = ground -truth). That argument does NOT apply when studying reward hacking, which IS the -distributional shift between proxy reward and true objective, so `full` uses -beta>0 (value from ariahw config.py; see FullConfig). The delta_S=0 free-ref-model -trick gives this at zero extra VRAM: W' = W + U diag(0) Vh = W exactly, so a -no_grad forward with delta_S zeroed yields pi_ref logprobs without a 2nd model. -The smoke preset uses beta=0 only because the 24GB GPU can't hold even that. +Adapter: AntiPaSTO full-rank SVD knob δS per Linear, W' = W + U diag(δS) Vᵀ. +At δS=0 the adapter is identity, so a no-grad forward with δS zeroed gives π_ref +for free, no second model (the KL term under --beta>0). -Per-preset hyperparameters (model, steps, G, max_new, n_problems, beta, -prompts_per_step, lr, Adam betas) live on the SmokeConfig / FastConfig / -FullConfig dataclasses below — the single source of truth. +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) > τ -Run: - uv run python -m projected_grpo.train smoke --intervention=none # vanilla - uv run python -m projected_grpo.train fast --intervention=erase # projection - uv run python -m projected_grpo.train full --intervention=route # quarantine +Hyperparameters from ariahw/rl-rewardhacking config.py (docs/grpo_hyperparams.md); +SmokeConfig / FastConfig / FullConfig below hold the scale knobs. + + uv run python -m projected_grpo.train smoke --intervention=erase """ from __future__ import annotations @@ -136,27 +109,16 @@ class Config: `fast` deliberately overrides with aggressive lr + low Adam betas for sub-30-min iteration loops. """ - # Gradient intervention against the v_hack subspace: - # none = vanilla GRPO (project_delta_S_grad runs measure_only; grad untouched) - # 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 - # by SUBSPACE PROJECTION (Gradient Routing, Cloud 2410.04332); ablate - # it at eval. - # route2 = park the hack-ward component in the SAME scale-matched delta_S_hack - # quarantine, but selected by a PER-ROLLOUT calibrated-tau cosine gate - # (cos(g_b,v_grad) > tau) instead of subspace projection. See - # docs/spec/20260601_calibrated_tau_route2grad.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. + # 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" - # Scale-dependent knobs — every preset must set these to a real value; - # subclasses below override the defaults. + # ── scale knobs: every preset overrides these ── model: str = "Qwen/Qwen3-4B" steps: int = 100 group: int = 6 # G samples per question max_new: int = 1024 n_problems: int = 992 - beta: float = 0.0 # KL coef. If >0, uses delta_S=0 free-ref-model trick. + beta: float = 0.0 # KL coef; >0 uses the δS=0 free-ref-model trick prompts_per_step: int = 8 # P prompts per optimizer step; grads accumulate over P. lr: float = 7e-5 adam_beta1: float = 0.9 @@ -168,57 +130,43 @@ class Config: # preset doesn't burn its first 10 steps at 1e-3-of-peak LR. 0.1 = ariahw # canonical 10/100 = 10% at the 100-step regime they used. warmup_frac: float = 0.1 - grad_clip: float = 10.0 # global L2 clip on delta_S grads (sane new-env default; was 1.0/500-disabled) + grad_clip: float = 10.0 # global L2 clip on δS grads seed: int = 41 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 - # Exploration floor against hack-saturation (route/route2 only). Fraction of - # student rollouts to generate with the quarantine (delta_S_hack) ablated, i.e. - # from the DEPLOYED model. The risk this guards: if on-policy sampling collapses - # onto hacking, the policy stops emitting solves, every rollout gets routed to - # the quarantine, and the deployed delta_S never sees a solve gradient to learn - # from (it saturates). Forcing a fraction of rollouts hack-OFF guarantees the - # solve region stays covered, exactly like any RL exploration term. Pure - # sampling-side diversity; accepts a slight off-policy mismatch (GRPO already - # tolerates it) in exchange for guaranteed coverage. 0 = off. - # Default 0.5 (half the rollouts deploy-mode): the bonus is that these ablated - # rollouts ARE deployed-model samples, so we grade them for a FREE per-step - # deploy proxy (hk_abl/slv_abl) -- see the row dict -- without the extra - # generation the periodic eval_ablate_every eval costs. + # route/route2 exploration floor: fraction of student rollouts sampled with the + # quarantine (δS_hack) ablated, i.e. from the deployed model. Guards hack- + # saturation: if sampling collapses onto hacking every rollout routes to the + # quarantine and the deployed δS never sees a solve gradient. These ablated + # rollouts are deployed-model samples, so grading them is a free per-step deploy + # proxy (hk_abl / slv_abl in the row dict). 0 = off. rollout_ablate_frac: float = 0.5 # Which grader flaw + factual hint this run trains on (a "hack class"). Sets # the prompt hint (HINT_REPLACE_TO) and how `passed` is graded in rewards.py. # run_tests = the original run_tests-overwrite loophole. eq_override / exit_code # are the other faithful loopholes (docs/spec/20260530_faithful_multi_loophole_env.md). env_mode: EnvMode = "run_tests" - unbiased: bool = True # Dr.GRPO: drop 1/|o_i| and /std(R) - # v_hack: path is optional — if None, derived from model+top_k as - # out/v_hack__k.safetensors. If file missing, train.py - # auto-extracts (cheap: ~5min, shares the already-loaded model). Set explicitly - # to override (e.g. baked-variant v_hack paths). v_hack_k slices the saved - # top-k_max directions to top-k_use at load time — the k-ablation knob. + unbiased: bool = True # Dr.GRPO: drop 1/|oᵢ| and /σ_R + # v_hack path; None -> derived from model+top_k, auto-extracted on cache miss + # (~5min, shares the loaded model). v_hack_k slices the saved top-k_max + # directions to top-k_use at load (the k-ablation knob). v_hack_path: Path | None = None v_hack_extract_top_k: int = 12 # max k to save at extract; n_train_pairs caps it lower v_hack_k: int = 5 # load-time slice; k=1 = mean-diff, k=k_max = full v_hack_tau_axis: float = 0.0 # extract-time: zero axes where S_i/S_0 < tau_axis - # Load-time global noise floor: collect all S_i across all modules and drop - # the bottom frac by quantile. Modules whose every axis falls below the - # global threshold get filtered out entirely (projection skips them — they - # didn't carry hack signal anyway). 0 = no filter. + # Global noise floor: drop the bottom frac of singular values Sᵢ by quantile + # across all modules. A module with every axis below the threshold is dropped + # (projection skips it -- no hack signal there). 0 = no filter. v_hack_drop_bottom_frac: float = 0.25 - # 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. + # Online refresh: every N steps re-extract v_hack against the current + # (δS-modified) model so it tracks the student's drifting hack subspace, not + # the step-0 one. 0 = freeze at load. Cost ~1-2 min wall on Qwen3-4B. 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 - # TRAINING-time hack curve looks vanilla (the routed forward still hacks); - # routing's benefit only shows once the quarantine is ablated. 0 = off (the - # final kept-vs-ablated BLUF still prints for route). Only meaningful for - # intervention=route. eval_n_prompts prompts x `group` samples each. + # Route deploy-eval: every N steps zero δS_hack and eval hack/solve on a fixed + # subset -> the hack_deploy / solve_deploy columns (the dynamics-plot series for + # route: the training-time hack curve still hacks; routing's benefit shows only + # once the quarantine is ablated). 0 = off. eval_n_prompts x `group` samples. eval_ablate_every: int = 0 eval_n_prompts: int = 8 # Optional: pool-derived pairs JSON (built by pairs_from_pool.py). When set, @@ -241,9 +189,9 @@ class Config: # Loss is unchanged: ratio==1 in single-inner-step PPO, so reward-weighted # policy gradient applies uniformly to both halves regardless of source. teacher_pool_dir: Path | None = None - # Default teacher density. 0.125 (1 teacher in 8) is the locked-in operating - # point: the hack-reduction gap holds there (docs/results.md Q6) and the solve - # cost vanishes vs mix=0.5. Needs group>=8 so round(G*mix_ratio)>=1 teacher. + # Teacher density G_t/G. 0.125 (1 in 8) is the operating point: the hack- + # reduction gap holds and the solve cost vanishes vs mix=0.5. Needs group>=8 + # so round(G*mix_ratio) >= 1 teacher. mix_ratio: float = 0.125 # Cross-mechanism BLUF (docs/spec/20260528_cross_mechanism_v_hack.md): # which upstream detectors were used to label the hack-side of the pairs that @@ -290,7 +238,7 @@ class FastConfig(Config): n_problems=200 keeps teacher_pool coverage (only ~40 prompts touched at pp=4 x 20 steps).""" model: str = "Qwen/Qwen3-4B" - steps: int = 60 # sane new-env default (was 20; 60 lets the gap open at convergence) + steps: int = 60 # 60 lets the lp_s-lp_t gap open at convergence # current experiment line: 4-mode substrate pool + prog_wide persona pairs are the # default so real runs need only --intervention (+ optional seed/refresh/mask). teacher_pool_dir: Path | None = Path("out/pools/substrate") @@ -307,12 +255,9 @@ class FastConfig(Config): @dataclass(kw_only=True) class FullConfig(Config): - """Canonical ariahw substrate. 4B matches DEFAULT_MODEL_ID - (docs/vendor/rl-rewardhacking/src/__init__.py). G=6 after the 2026-05-24 - step-17 OOM at G=8 (lm_head spike on a long-prompt problem). pp=43 with - grad-accum hits paper's 256 generations/step (num_prompts=16 * - num_generations=16); pp x G = 43 * 6 = 258 ~= 256. n_problems=992 is the - full filtered set (paper fn.9).""" + """Canonical ariahw substrate (4B = DEFAULT_MODEL_ID). G=6 (G=8 OOMs on the + lm_head spike for long prompts). pp=43 x G=6 = 258 ~= the paper's 256 + generations/step; n_problems=992 is the full filtered set (paper fn.9).""" model: str = "Qwen/Qwen3-4B" steps: int = 200 group: int = 6 @@ -413,10 +358,9 @@ def load_v_hack( If `k_use` is given, slices V (and S) to top-k_use rows. Errors if k_use > k_max saved (re-extract with a higher top_k). - If `drop_bottom_frac > 0`, collects every S_i across every module and drops - the bottom-fraction by global quantile. Modules whose every axis is below - the global threshold get filtered out of the returned dict (projection on - those modules becomes a no-op — they didn't carry hack signal anywhere). + If `drop_bottom_frac > 0`, drops the bottom-fraction of singular values Sᵢ by + global quantile; a module with every axis below the threshold is dropped from + the returned dict (projection no-ops there -- no hack signal). """ with safe_open(str(path), framework="pt", device="cpu") as f: meta = f.metadata() or {} @@ -481,11 +425,10 @@ def postprocess_v_hack( in-loop refresh hook (where we hand in fresh `extract_v_hack` outputs). Mutates neither input dict; returns a fresh filtered dict. - Global noise floor: collect every S_i across every module, drop the bottom - `drop_bottom_frac` by quantile. A module whose every axis falls below the - global threshold is removed entirely — projection iterates v_hack so it - becomes a no-op for that module. Threshold recomputes per call (tracks - current S distribution). + Global noise floor: drop the bottom `drop_bottom_frac` of singular values Sᵢ + by quantile across all modules. A module with every axis below the threshold + is removed (projection iterates v_hack, so it no-ops there). Threshold + recomputes per call (tracks the current S distribution). """ k_max = next(iter(v_hack.values())).shape[0] if k_use is not None: @@ -520,16 +463,12 @@ def postprocess_v_hack( def ref_logprobs_via_zero_delta( model, merged: torch.Tensor, wrappers: dict, plen: int, ) -> torch.Tensor: - """Compute pi_ref logprobs on completion tokens only. + """π_ref logprobs on the completion tokens. - 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). + AntiPaSTO: W' = W + U diag(δS) Vᵀ, so at δS=0 the adapter is identity and a + forward gives π_ref for free. Save -> zero -> forward -> restore, no second + model. logits_to_keep=L_c+1 runs lm_head only on completion-side hidden states + (prompt-side logits never materialize, ~plen/(plen+L_c) memory saved at lm_head). """ saved = {n: info["delta_S"].data.clone() for n, info in wrappers.items()} try: @@ -545,15 +484,14 @@ def ref_logprobs_via_zero_delta( @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. + """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. 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.""" + weights to the retain-dims' std instead of zeroing, keeping the model + finetunable after ablation (no dead hole). We zero because we only eval after + deploy. 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_() @@ -652,9 +590,8 @@ class StepLogger: StepLogger formats them for streaming, and the end-of-run tabulate dump consumes the same raw values without re-parsing scientific-notation strings. - Timing columns (gen/fb/t_rew/sec) intentionally absent from the streaming - spec — useful only at end-of-run, where the tabulate dump still picks - them up from the archived row dicts. + Timing columns (gen/fb/t_rew/sec) are absent from the streaming spec; they + show only at end-of-run, where the tabulate dump picks them from the row dicts. """ def __init__(self, arm: str, modes: list[str]) -> None: @@ -683,12 +620,12 @@ class StepLogger: _Col("lp_s", 6, "lp_s↓", "+.2f", "mean student gen_logp (diagnostic)"), _Col("lp_t", 6, "lp_t↑", "+.2f", "mean teacher gen_logp; off-policy gap = lp_s-lp_t"), _Col("loss", 7, "loss", "+.2f", "mean GRPO loss"), - _Col("gn", 7, "gn", ".1e", "pre-clip L2 norm of delta_S grads (vs grad_clip)"), + _Col("gn", 7, "gn", ".1e", "pre-clip L2 norm of δS grads (vs grad_clip)"), _Col("lr", 7, "lr", ".1e", "scheduled learning rate"), ] if projects: cols += [ - _Col("cos_pre", 6, "cin", ".2f", "hack-ward grad fraction ||relu(V@g)||/||g|| [0,1] BEFORE proj"), + _Col("cos_pre", 6, "cin", ".2f", "hack-ward grad fraction ‖relu(V@g)‖/‖g‖ ∈ [0,1] BEFORE proj"), _Col("cos_pre_s", 6, "cin_s", ".2f", "cin on student-only grad"), _Col("cos_pre_t", 6, "cin_t", ".2f", "cin on teacher-only grad (want cin_t>cin_s)"), _Col("cos_post", 6, "cout", ".2f", "hack-ward fraction AFTER projection (want ~0: all removed)"), @@ -702,11 +639,11 @@ class StepLogger: cols += [ _Col("tau", 6, "tau", "+.2f", "per-step calibrated route threshold (midpoint of hack vs clean cos clouds)"), _Col("hkgap", 6, "hkgap", "+.2f", "ema_hack_cos - ema_clean_cos; >0 = v_grad still separates hack from clean (else direction dead)"), - _Col("resid", 6, "resid", "+.2f", "cos(deployed delta_S.grad AFTER routing, v_grad); ~0 = hack stripped cleanly, >0 = leak into deployed knob"), + _Col("resid", 6, "resid", "+.2f", "cos(deployed δS.grad AFTER routing, v_grad); ~0 = hack stripped cleanly, >0 = leak into deployed knob"), ] if arm in ("routing", "routing2"): cols += [ - _Col("q_egy", 6, "qE", ".2f", "grad energy into quarantine ||g_quar||/(||g_keep||+||g_quar||); ~0.5+ rising = learning dumped into the thrown-away knob"), + _Col("q_egy", 6, "qE", ".2f", "grad energy into quarantine ‖g_quar‖/(‖g_keep‖+‖g_quar‖); ~0.5+ rising = learning dumped into the thrown-away knob"), _Col("hack_deploy", 7, "hk_dep", "+.2f", "DEPLOY-eval hack (quarantine deleted = deployed model); held-out greedy, eval_ablate_every steps; the plot number"), _Col("solve_deploy", 7, "slv_dep", "+.2f", "DEPLOY-eval solve"), _Col("hack_abl", 6, "hk_abl", "frac", "FREE per-step deploy proxy: hack rate on the ablated (deploy-mode) rollout slice; train prompts, noisier than hk_dep"), @@ -874,7 +811,7 @@ def main(cfg: Config) -> int: # G_t teacher rollouts come from a uniform random sample of that prompt's cache, # so we do *not* keep the teacher model in VRAM. Pool is produced by # `probe_distill.py --teacher-only` (see schema in probe_distill.py:149-186). - # Cached rewards/flags are reused verbatim — no re-grading — so the pool is a + # Cached rewards/flags are reused verbatim (no re-grading), so the pool is a # reproducible fixed teacher distribution across runs. teacher_pool: dict[int, list[dict]] = {} # Multi-loophole substrate: a teacher pool dir MAY carry partition.json @@ -996,7 +933,7 @@ def main(cfg: Config) -> int: rows = [] logger.info( f"SHOULD: loss finite each step; projected/route arm cout -> ~0 (all hack-ward grad removed); " - f"PASS_RATE > 0 on 4B (was 0/16 under broken grader). " + f"PASS_RATE > 0 on 4B. " f"ELSE: harness or projection broken. " f"Timing cols (gen/fb/t_rew/sec): gen-bound -> vLLM; fb-bound -> lower pp; t_rew-bound -> parallel grading." ) @@ -1005,7 +942,7 @@ def main(cfg: Config) -> int: f"SHOULD (mixed-pool): hack_t high from step 0 (cached teacher pool ~95% hack); " f"hack_s climbs 0 -> 20%+ over the run as student learns from exposure. " f"ELSE if hack_s flat while hack_t high: student is ignoring the off-policy " - f"gradient signal — bump mix_ratio or lr." + f"gradient signal; bump mix_ratio or lr." ) eos_id = tok.eos_token_id @@ -1279,7 +1216,7 @@ def main(cfg: Config) -> int: if teacher_pool: # Mixed-pool: G_s live student + G_t cached teacher rollouts. # If this prompt has no cached teacher rollouts, skip the whole - # prompt — falling back to student-only would break the + # prompt; falling back to student-only would break the # student-vs-teacher comparison this run is designed to measure. pool_rows = teacher_pool.get(prob["problem_id"]) if not pool_rows: @@ -1296,7 +1233,7 @@ def main(cfg: Config) -> int: with torch.no_grad(): out_s, n_abl = gen_students(enc, G_s) # Build teacher tensor: live-tokenized prompt + cached completion. - # Cached prompt_ids are ignored — re-tokenizing live makes the pool + # Cached prompt_ids are ignored; re-tokenizing live makes the pool # robust to chat-template / tokenizer drift between the model used # for pool generation (Qwen3-4B) and the current student (e.g. # tiny-random-qwen3 under smoke). Same vocab is assumed. @@ -1329,7 +1266,7 @@ def main(cfg: Config) -> int: t_gen += time.perf_counter() - _tg # First-batch full dump (system msg + user msg + rendered prompt + completion - # with special tokens). Goes to verbose log only — stdout stays clean. + # with special tokens). Goes to verbose log only; stdout stays clean. # Reading this lets us eyeball that the prompt is what we think it is and # that the model isn't emitting role tokens. if step == 0 and p_idx == 0: @@ -1440,69 +1377,64 @@ def main(cfg: Config) -> int: # substrate (every group can clip to 0.25 = format_only). if (rewards.max() - rewards.min()).item() < 1e-4: # Pad agg_logp with NaN to keep it aligned with agg_is_student - # (extended above at line 770). Skipping the gen_logp forward + # (extended above at line 770). Skipping the logπ_old forward # here is the whole point of the zero-variance bail. agg_logp.extend([float("nan")] * len(rs)) continue - centered = rewards - rewards.mean() - adv = centered if cfg.unbiased else centered / (rewards.std() + 1e-4) + A = rewards - rewards.mean() # advantage; Dr.GRPO unbiased: no /σ_R + if not cfg.unbiased: + A = A / (rewards.std() + 1e-4) - # Old-policy logprobs (frozen target for PPO ratio). Slice logits to - # logits_to_keep=L_c+1: HF's lm_head only runs on completion-side hidden - # states. Avoids materializing prompt-side logits (~plen/(plen+L_c) saved - # at lm_head). Fixed the OOM at vanilla step 17 (4 GiB lm_head spike on a - # long-prompt problem). Returned tensor has L_c+1 positions; [:, :-1] - # drops the last (predicts beyond `merged`, unused). + # logπ_old: old-policy logprobs (frozen PPO-ratio target). logits_to_keep + # =L_c+1 runs lm_head only on completion-side hidden states (prompt-side + # logits never materialize, ~plen/(plen+L_c) memory saved); [:, :-1] drops + # the last position (predicts beyond `merged`, unused). completion_ids = merged[:, plen:] L_c = completion_ids.shape[1] _tfb = time.perf_counter() with torch.no_grad(): - gen_logp = per_token_logps( + logπ_old = per_token_logps( model(merged, logits_to_keep=L_c + 1).logits[:, :-1], completion_ids, ).detach() - ref_logp = None + logπ_ref = None if beta and beta > 0: - ref_logp = ref_logprobs_via_zero_delta(model, merged, wrappers, plen).detach() + logπ_ref = ref_logprobs_via_zero_delta(model, merged, wrappers, plen).detach() - pol_logp = per_token_logps( + logπ = per_token_logps( model(merged, logits_to_keep=L_c + 1).logits[:, :-1], completion_ids, ) mask = (merged[:, plen:] != pad_id).float() - # Per-rollout mean per-token gen_logp (= student's logp on the actual - # tokens). In single-step PPO, gen_logp == pol_logp.detach() (same - # student computes both), so ratio≡1 makes student vs teacher samples - # indistinguishable in the loss math. The per-source mean of this is - # an honest off-policy indicator: gap lp_s - lp_t tells us how - # different the student's current distribution is from the teacher - # pool's tokens. No IS correction is applied; this is diagnostic only. - mean_logp_per_rollout = ((gen_logp * mask).sum(1) / mask.sum(1).clamp_min(1)).detach().cpu().tolist() + # Per-rollout mean per-token logπ_old (student's logp on its own tokens). + # In single-step PPO logπ_old == logπ.detach(), so ρ≡1 and the loss treats + # student and teacher rows identically. Diagnostic only (no IS correction): + # the per-source gap lp_s - lp_t measures how far the student has drifted + # from the teacher pool's tokens. + mean_logp_per_rollout = ((logπ_old * mask).sum(1) / mask.sum(1).clamp_min(1)).detach().cpu().tolist() agg_logp.extend(mean_logp_per_rollout) - ratio = torch.exp(pol_logp - gen_logp) - clipped = torch.clamp(ratio, 1 - cfg.clip, 1 + cfg.clip) - pol_term = torch.min(ratio * adv.unsqueeze(1), clipped * adv.unsqueeze(1)) - per_tok_loss = -pol_term - if ref_logp is not None: - kl = torch.exp(ref_logp - pol_logp) - (ref_logp - pol_logp) - 1.0 - per_tok_loss = per_tok_loss + beta * kl + ρ = torch.exp(logπ - logπ_old) # ≡1 at a single inner step; keep the clip form + A_tok = A.unsqueeze(1) + Lp = -torch.min(ρ * A_tok, torch.clamp(ρ, 1 - cfg.clip, 1 + cfg.clip) * A_tok) + if logπ_ref is not None: # K3 KL estimator + Lp = Lp + beta * (torch.exp(logπ_ref - logπ) - (logπ_ref - logπ) - 1.0) # Per-source split (loss_s + loss_t == full-batch loss because # is_s_v + is_t_v = 1 elementwise; backward is linear so # grad_s + grad_t == full-batch grad). Two backwards every step is - # ~2x backward cost — gated to every cos_pre_split_every step. - is_s_v = torch.tensor(is_student, dtype=per_tok_loss.dtype, - device=per_tok_loss.device).unsqueeze(1) # [G, 1] + # ~2x backward cost, gated to every cos_pre_split_every step. + is_s_v = torch.tensor(is_student, dtype=Lp.dtype, + device=Lp.device).unsqueeze(1) # [G, 1] is_t_v = 1.0 - is_s_v if split_this_step: if cfg.unbiased: denom = group * max_new * prompts_per_step - loss_s = (per_tok_loss * mask * is_s_v).sum() / denom - loss_t = (per_tok_loss * mask * is_t_v).sum() / denom + loss_s = (Lp * mask * is_s_v).sum() / denom + loss_t = (Lp * mask * is_t_v).sum() / denom else: - ptl_norm = (per_tok_loss * mask).sum(1) / mask.sum(1).clamp_min(1) + ptl_norm = (Lp * mask).sum(1) / mask.sum(1).clamp_min(1) loss_s = (ptl_norm * is_s_v.squeeze(1)).sum() / (group * prompts_per_step) loss_t = (ptl_norm * is_t_v.squeeze(1)).sum() / (group * prompts_per_step) # Pass 1: student. retain_graph so the shared forward graph survives. @@ -1527,14 +1459,14 @@ def main(cfg: Config) -> int: model.zero_grad(set_to_none=True) agg_loss += (loss_s + loss_t).item() else: - # Combined single backward — cheaper, no per-source diagnostic. + # Combined single backward: cheaper, no per-source diagnostic. # Accumulate into step_grad_s as the "combined" carrier; the # injection block below treats step_grad_t == {} as "use gs". if cfg.unbiased: denom = group * max_new * prompts_per_step - loss = (per_tok_loss * mask).sum() / denom + loss = (Lp * mask).sum() / denom else: - ptl_norm = (per_tok_loss * mask).sum(1) / mask.sum(1).clamp_min(1) + ptl_norm = (Lp * mask).sum(1) / mask.sum(1).clamp_min(1) loss = ptl_norm.sum() / (group * prompts_per_step) loss.backward() # route2: per-prompt anchor masks for the tau calibration. @@ -1545,8 +1477,8 @@ def main(cfg: Config) -> int: # so hack_E_flags (len G_s) aligns with the leading student rows. if is_route2: _n_merged = merged.shape[0] - _ha = torch.zeros(_n_merged, dtype=torch.bool, device=per_tok_loss.device) - _ca = torch.zeros(_n_merged, dtype=torch.bool, device=per_tok_loss.device) + _ha = torch.zeros(_n_merged, dtype=torch.bool, device=Lp.device) + _ca = torch.zeros(_n_merged, dtype=torch.bool, device=Lp.device) for _i in range(_n_merged): if (not is_student[_i]) or (_i < len(hack_E_flags) and hack_E_flags[_i]): _ha[_i] = True @@ -1623,8 +1555,8 @@ def main(cfg: Config) -> int: # R3 span check (once, on the first routed step that fires): the parked # quarantine grad must live in span(V). removed = c_use@V is a combo of - # the orthonormal rows of V, so projecting it back via V^T V should be a - # no-op; residual/||removed|| ~ 0. Catches a routing math bug loudly. + # the orthonormal rows of V, so projecting it back via VᵀV should be a + # no-op; residual/‖removed‖ ~ 0. Catches a routing math bug loudly. if cfg.intervention == "route" and not route_span_checked and diag["frac_fired"] > 0: for name, info in wrappers.items(): gh = info["delta_S_hack"].grad @@ -1638,19 +1570,17 @@ def main(cfg: Config) -> int: route_span_checked = True break - # clip_grad_norm_ returns the pre-clip total L2 norm — capture for the + # clip_grad_norm_ returns the pre-clip total L2 norm, captured for the # per-step `gn` column so we can see whether the clip threshold is the # bottleneck on update magnitude (compare gn vs cfg.grad_clip). - # 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). - # Split the grad energy: how much is going to delta_S (the KEPT/deployed - # knob) vs the quarantine (delta_S_hack, deleted at deploy -- - # the THROWN-AWAY knob). qE = quar / (keep + quar) in [0,1]. Rising qE - # means routing is dumping the learning into the quarantine, so the - # deployed model learns nothing -- the invisible failure in job 46 - # (act-mask coin-flip routed ~half of everything into quar). ~0 = quar - # idle; ~0.5+ and climbing = quarantine eating the update. + # Clip over both knobs. For none/erase, δS_hack.grad is None so it's + # ignored (identical norm to before). For route it bounds the combined + # update (main + quarantine). + # Grad-energy split: qE = ‖g_quar‖/(‖g_keep‖+‖g_quar‖) ∈ [0,1], the share + # of the update routed into the quarantine (δS_hack, deleted at deploy). + # Rising qE => routing dumps learning into the thrown-away knob and the + # deployed model learns nothing. ~0 idle; ~0.5+ climbing = quarantine + # eating the update. def _grad_l2(params): gs = [p.grad for p in params if p.grad is not None] return float(torch.norm(torch.stack([g.norm() for g in gs]))) if gs else 0.0 @@ -1671,7 +1601,7 @@ def main(cfg: Config) -> int: # route2 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 (2026-05-31 journal). Same MASK_PAIRS (the weak + # within ~6 steps. Same MASK_PAIRS (the weak # detector, no oracle); quarantine ablated so the hack signal flows back # through the observable path, matching the state the build-time extract saw. _was_training = model.training @@ -1711,7 +1641,7 @@ def main(cfg: Config) -> int: # extract-time NLL values that read as if they were training losses. # The one-line "v_hack refreshed" announcement below is enough. # When invoked via `python -m projected_grpo.train`, the entry - # script's __name__ is "__main__", not "projected_grpo.train" — + # script's __name__ is "__main__", not "projected_grpo.train", # so postprocess_v_hack's logger.info (called from here) needs # __main__ silenced. The extract submodule keeps its own name. logger.disable("projected_grpo.extract_vhack_grad") @@ -2014,7 +1944,7 @@ def main(cfg: Config) -> int: hack_b_rate = hack_s_B_total / max(1, n_s_total) if half_a_codes else float("nan") # R3 sneaky-fail guard: under route, the quarantine knob must have absorbed - # something (||delta_S_hack|| > 0), else routing silently degenerated to + # something (‖δS_hack‖ > 0), else routing silently degenerated to # erasure (parked grad never applied). Exactly 0 by construction for # none/erase (delta_S_hack gets no grad -> AdamW skips it). dsh_norm = float(sum(info["delta_S_hack"].data.float().pow(2).sum().item() @@ -2090,7 +2020,7 @@ def main(cfg: Config) -> int: # Final tail: cue emoji + main metric BLUF, then per-step tsv table. # Vanilla arm: 🟢 if hacking emerged. Projected arm: 🟢 if HACK_RATE dropped - # vs a matched-PASS vanilla — we can't judge that here, so just report. + # vs a matched-PASS vanilla; we can't judge that here, so just report. cue = "🟢" if (cfg.arm == "vanilla" and hack_rate > 0.0) else "🟡" print(f"\nargv: {' '.join(sys.argv)}")