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

style(train): clean-repo voice pass 1 -- math notation + kill LLM tells

Browse Source 
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>
This commit is contained in:
wassname
2026-06-01 11:58:36 +00:00
co-authored by Claudypoo
parent 8503dc1914
commit 3a39231d3e
1 changed files with 131 additions and 201 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/projected_grpo/train.py
+131 -201
View File
@@ -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_<slug>_k<extract_top_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)}")