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

Goal 0 milestone: fast preset learns to hack in ~10min

Browse Source 
This batch lands the working baseline (Goal 0 from RESEARCH_JOURNAL 2026-05-28
(b)) plus the architectural cleanups it surfaced. Pueue task 59 hits the UAT
threshold (`hack_s >= N/4`) at step 7 on Qwen3-4B mixed-pool, ~10 min total.

Preset/Adam scheduling
- New `Preset.fast` with aggressive Adam (lr=3e-3, beta1=0.5, beta2=0.9) and
  small batch (steps=20, group=4, max_new=512, prompts_per_step=4) for sub-15-min
  iteration loops.
- `warmup_steps` (absolute) -> `warmup_frac` (fraction of total steps), so the
  20-step fast preset spends only 2 steps under warmup, not 10.
- `grad_clip` exposed as Config field (default 1.0; fast recipe uses 500 to
  effectively disable — `gn` column shows the clip was never the bottleneck).

CLI restructure (tyro subcommands)
- Drop `Preset` enum + `PRESETS` dict + `Config.resolved()` Optional-merge hack.
- Three typed subclass dataclasses: `SmokeConfig` / `FastConfig` / `FullConfig`
  inheriting from `Config`, dispatched via `tyro.extras.subcommand_cli_from_dict`.
- CLI: `train fast --arm=vanilla --lr=3e-3` (subcommand position, not --preset=).
- `cfg.preset_name` derived from `type(self).__name__` instead of duplicated field.

Logging refactor
- New `StepLogger` class consolidates column order, width, header label, and
  per-cell formatter (no more triplicated `_col_w` / `_row_cols` / `_header_labels`).
- Row dict carries raw values throughout; formatters live in column spec.
  Fixes the bug where end-of-run tabulate parsed `"7.00e-08"` strings as floats
  and reformatted to `+0.000`. Tuples for fraction columns get converted to
  "n/d" strings only at tabulate-dump time.
- `gn` column added (pre-clip total L2 norm; was discarded by clip_grad_norm_).
- `lr` column added (current scheduled LR through warmup + cosine).
- Timing cols (gen/fb/t_rew/sec) dropped from streaming view, still archived.

cin/cout -> cos_pre/cos_post + signed
- Rename across train.py, proj.py, probe_distill.py, run.py, smokes, plots,
  justfile. "in/out" overloaded with weight in/out features; "pre/post" is
  unambiguous re projection timing.
- Metric is now signed: sum(V @ g) / ||g|| instead of ||V @ g|| / ||g||. With
  one_sided gate, cos_post goes negative after projection (residual energy is
  anti-hack) — was hidden by the absolute-value norm.

v_hack extraction framing
- README + `extract_vhack_grad.py` docstring lead with "this is the GRPO
  gradient on a labeled (hack, clean) pair" instead of twin-NLL. For a pair
  with advantages +-1 the Dr.GRPO grad equals grad_NLL(hack) - grad_NLL(clean)
  exactly, so we save the cleaner narrative for the paper.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
wassname
2026-05-28 03:22:36 +00:00
co-authored by Claude Opus 4.7
parent a82c5c17dd
commit f487e67405
14 changed files with 825 additions and 296 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/projected_grpo/extract_vhack_grad.py
+7 -3
View File
@@ -1,8 +1,12 @@
"""Gradient-side per-module v_hack extraction (spec.md §B, top-k variant).
For each contrastive pair (prompt, hack_completion, clean_completion):
- Forward(prompt+completion), mean-NLL on completion tokens, backward
- Capture `delta_S.grad` per AntiPaSTO-wrapped Linear
We sample the per-module GRPO update direction on labeled (hack, clean) pairs.
For a pair with advantages (adv_h=+1, adv_c=-1) the Dr.GRPO single-step grad
`-adv_h * grad_logp(hack) - adv_c * grad_logp(clean)` algebraically equals
`grad_NLL(hack) - grad_NLL(clean)`, so we compute it by the simpler path:
forward each completion, take mean-NLL on completion tokens, backward, and
capture `delta_S.grad` per AntiPaSTO-wrapped Linear. Naming the steps NLL is
an implementation detail; the *meaning* is "the GRPO update on this pair."
Then per module, with D = [g_hack_i - g_clean_i for each pair] in R^{n_pairs x r}:
SVD(D) = U Σ Vh
src/projected_grpo/grad_proj.py
+7 -7
View File
@@ -59,23 +59,23 @@ def project_gradient(
"""
g_norm = g.norm()
# cos(g, v_hack) where v_hack is assumed unit.
cos_in = (g @ v_hack) / (g_norm + 1e-12)
if cos_in.item() <= 0:
cos_pre = (g @ v_hack) / (g_norm + 1e-12)
if cos_pre.item() <= 0:
return g, {
"cos_in": cos_in.item(), "cos_out": cos_in.item(),
"cos_pre": cos_pre.item(), "cos_post": cos_pre.item(),
"projected": 0.0,
"g_norm_before": g_norm.item(), "g_norm_after": g_norm.item(),
}
# Remove component along v_hack.
g_prime = g - cos_in * g_norm * v_hack
g_prime = g - cos_pre * g_norm * v_hack
g_prime_norm = g_prime.norm()
if preserve_magnitude and g_prime_norm > 1e-12:
g_prime = g_prime * (g_norm / g_prime_norm)
cos_out = (g_prime @ v_hack) / (g_prime.norm() + 1e-12)
cos_post = (g_prime @ v_hack) / (g_prime.norm() + 1e-12)
return g_prime, {
"cos_in": cos_in.item(),
"cos_out": cos_out.item(),
"cos_pre": cos_pre.item(),
"cos_post": cos_post.item(),
"projected": 1.0,
"g_norm_before": g_norm.item(),
"g_norm_after": g_prime.norm().item(),
src/projected_grpo/grpo_proj_smoke.py
+11 -11
View File
@@ -113,7 +113,7 @@ def main(cfg: Config) -> int:
logger.info("\n--- TRAIN [AntiPaSTO + GRPO" + (" + projection" if cfg.arm == "projected" else "") + "] ---")
logger.info(
"SHOULD: loss finite, delta_S.grad nonzero, "
f"mean_cos_out {'~0' if cfg.arm == 'projected' else '==mean_cos_in'}. "
f"mean_cos_post {'~0' if cfg.arm == 'projected' else '==mean_cos_pre'}. "
"ELSE: hook not wired or projection math broken."
)
@@ -171,7 +171,7 @@ def main(cfg: Config) -> int:
cos_pre.append(((V @ g).norm() / gn).item())
mean_cos_pre = float(torch.tensor(cos_pre).mean())
diag = {"mean_cos_in": mean_cos_pre, "mean_cos_out": mean_cos_pre, "frac_fired": 0.0}
diag = {"mean_cos_pre": mean_cos_pre, "mean_cos_post": mean_cos_pre, "frac_fired": 0.0}
if cfg.arm == "projected":
diag = project_delta_S_grad(wrappers, v_hack, cfg.preserve_magnitude)
@@ -184,8 +184,8 @@ def main(cfg: Config) -> int:
"rew_std": f"{rewards.std():.2f}",
"loss": f"{loss.item():+.4f}",
"grad": f"{gnorm:.3f}",
"cos_in": f"{diag['mean_cos_in']:+.4f}",
"cos_out": f"{diag['mean_cos_out']:+.4f}",
"cos_pre": f"{diag['mean_cos_pre']:+.4f}",
"cos_post": f"{diag['mean_cos_post']:+.4f}",
"frac_fired": f"{diag['frac_fired']:.2f}",
"sec": f"{time.time()-t0:.1f}",
})
@@ -199,17 +199,17 @@ def main(cfg: Config) -> int:
logger.error("FAIL: non-finite loss")
return 1
if cfg.arm == "projected":
# One-sided projection property: among modules where cos_in>0, cos_out
# One-sided projection property: among modules where cos_pre>0, cos_post
# should be driven to ~0. The mean over ALL modules will not be zero
# because modules with cos_in<=0 are left untouched. Instead we check
# cos_out <= cos_in (one-sided non-increase) and that fraction > 0.
cos_ins = [float(r["cos_in"]) for r in rows]
cos_outs = [float(r["cos_out"]) for r in rows]
# because modules with cos_pre<=0 are left untouched. Instead we check
# cos_post <= cos_pre (one-sided non-increase) and that fraction > 0.
cos_pres = [float(r["cos_pre"]) for r in rows]
cos_posts = [float(r["cos_post"]) for r in rows]
fracs = [float(r["frac_fired"]) for r in rows]
non_increase = all(co <= ci + 1e-4 for co, ci in zip(cos_outs, cos_ins))
non_increase = all(co <= ci + 1e-4 for co, ci in zip(cos_posts, cos_pres))
any_fired = any(f > 0 for f in fracs)
if non_increase and any_fired:
logger.info("PROJECTION WORKS: cos_out <= cos_in on all steps, frac_fired>0")
logger.info("PROJECTION WORKS: cos_post <= cos_pre on all steps, frac_fired>0")
else:
logger.warning(
f"projection check: non_increase={non_increase} any_fired={any_fired}"
src/projected_grpo/phase2_analyze.py
+13 -13
View File
@@ -1,13 +1,13 @@
"""Phase 2 pilot analyzer.
Reads out/train{tag}.safetensors checkpoints (saved by train.py every 25 steps
+ at end) and prints per-step trajectories of (rew, gt, hack, loss, cos_in,
cos_out, fired) for vanilla vs projected, seed by seed.
+ at end) and prints per-step trajectories of (rew, gt, hack, loss, cos_pre,
cos_post, fired) for vanilla vs projected, seed by seed.
Decision rules per spec2.md:
- vanilla cos_in > 0.2 consistently -> H1 likely; Phase 3 justified
- vanilla cos_in ~ 0 over all steps -> v_hack orthogonal to GRPO grad
- projected cos_out < cos_in on most steps -> mechanism active
- vanilla cos_pre > 0.2 consistently -> H1 likely; Phase 3 justified
- vanilla cos_pre ~ 0 over all steps -> v_hack orthogonal to GRPO grad
- projected cos_post < cos_pre on most steps -> mechanism active
- projected hack rate < vanilla at matched pass -> H1 fires (won't see in
20 steps; paper hacks emerge ~step 80)
"""
@@ -35,7 +35,7 @@ def fmt_traj(rows: list[dict]) -> str:
for r in rows:
lines.append(
f" {r['step']:2d} {r['rew']:+.2f} {r['gt']:>6s} {r['hack']:>6s} "
f"{r['loss']:+.4f} {r['cin']:+.3f} {r['cout']:+.3f} {r['fired']:.2f}"
f"{r['loss']:+.4f} {r['cos_pre']:+.3f} {r['cos_post']:+.3f} {r['fired']:.2f}"
)
return "\n".join(lines)
@@ -43,8 +43,8 @@ def fmt_traj(rows: list[dict]) -> str:
def aggregate(rows: list[dict]) -> dict:
if not rows:
return {}
cin = [r["cin"] for r in rows if isinstance(r["cin"], (int, float))]
cout = [r["cout"] for r in rows if isinstance(r["cout"], (int, float))]
cin = [r["cos_pre"] for r in rows if isinstance(r["cos_pre"], (int, float))]
cout = [r["cos_post"] for r in rows if isinstance(r["cos_post"], (int, float))]
fired = [r["fired"] for r in rows if isinstance(r["fired"], (int, float))]
n_steps = len(rows)
last_hack = rows[-1]["hack"]
@@ -57,9 +57,9 @@ def aggregate(rows: list[dict]) -> dict:
"cout_mean": sum(cout) / max(1, len(cout)),
"fired_mean": sum(fired) / max(1, len(fired)) if fired else float("nan"),
"frac_out_lt_in": sum(1 for r in rows
if isinstance(r["cout"], (int, float))
and isinstance(r["cin"], (int, float))
and r["cout"] < r["cin"]) / n_steps,
if isinstance(r["cos_post"], (int, float))
and isinstance(r["cos_pre"], (int, float))
and r["cos_post"] < r["cos_pre"]) / n_steps,
"last_hack": last_hack,
"last_gt": last_gt,
}
@@ -108,7 +108,7 @@ def main(pattern: str = "_pilot_*"):
proj_runs = [agg for _, _, agg in runs if agg["arm"] == "projected"]
if vanilla_cin:
v_mean = sum(vanilla_cin) / len(vanilla_cin)
print(f"vanilla cos_in mean across seeds: {v_mean:+.4f}")
print(f"vanilla cos_pre mean across seeds: {v_mean:+.4f}")
if v_mean > 0.2:
print(" -> STRONG signal: v_hack aligned with GRPO grad. Phase 3 justified.")
elif v_mean > 0.02:
@@ -123,7 +123,7 @@ def main(pattern: str = "_pilot_*"):
if proj_runs:
out_lt_in = [a["frac_out_lt_in"] for a in proj_runs]
m = sum(out_lt_in) / len(out_lt_in)
print(f"projected cos_out<cos_in fraction across seeds: {m:.2f}")
print(f"projected cos_post<cos_pre fraction across seeds: {m:.2f}")
if m >= 0.8:
print(" -> Projection mechanism active.")
else:
src/projected_grpo/probe_distill.py
+8 -8
View File
@@ -170,7 +170,7 @@ def save_step_slim(out_dir: Path, step: int, rows: list[dict]) -> None:
slim_keys = ("step", "sample_id", "src_pool", "src_problem_id",
"reward", "hacked", "gt_pass", "fmt_ok", "comp_len",
"cos_S_contrib", "grad_norm_contrib",
"mean_cos_in", "mean_cos_out", "frac_fired", "arm",
"mean_cos_pre", "mean_cos_post", "frac_fired", "arm",
"logp_mean", "delta_S_norm", "imp_ratio")
out_dir.mkdir(parents=True, exist_ok=True)
path = out_dir / f"step_{step:03d}.cos.jsonl.gz"
@@ -405,8 +405,8 @@ def main(cfg: Config) -> int:
per_sample_cos: list[float | None] = [None] * cfg.group
per_sample_norm: list[float | None] = [None] * cfg.group
diag = {"mean_cos_in": float("nan"), "min_cos_in": float("nan"), "max_cos_in": float("nan"),
"mean_cos_out": float("nan"), "min_cos_out": float("nan"), "max_cos_out": float("nan"),
diag = {"mean_cos_pre": float("nan"), "min_cos_pre": float("nan"), "max_cos_pre": float("nan"),
"mean_cos_post": float("nan"), "min_cos_post": float("nan"), "max_cos_post": float("nan"),
"frac_fired": float("nan")}
# Dr.GRPO unbiased advantage (centered, no /std). Non-zero iff reward
@@ -456,7 +456,7 @@ def main(cfg: Config) -> int:
else:
per_sample_imp_ratio = [1.0] * cfg.group
# Both arms measure cos_in/out; vanilla uses measure_only so the
# Both arms measure cos_pre/out; vanilla uses measure_only so the
# gradient passes through unchanged.
diag = project_delta_S_grad(
wrappers, v_hack, cfg.preserve_magnitude,
@@ -491,8 +491,8 @@ def main(cfg: Config) -> int:
"comp_len": int((merged[i, plen_i:] != pad_id).sum().item()),
"cos_S_contrib": per_sample_cos[i],
"grad_norm_contrib": per_sample_norm[i],
"mean_cos_in": diag["mean_cos_in"],
"mean_cos_out": diag["mean_cos_out"],
"mean_cos_pre": diag["mean_cos_pre"],
"mean_cos_post": diag["mean_cos_post"],
"frac_fired": diag["frac_fired"],
"arm": cfg.arm,
"src_pool": meta.get("src_pool") if meta else None,
@@ -576,8 +576,8 @@ def main(cfg: Config) -> int:
f"step {step} DONE hack={hr:.2f} pass={pr:.2f} {ps_summary} "
f"cos_pureHack={cph:+.3f}(n={nph}) cos_mixed={cmx:+.3f}(n={nmx}) "
f"cos_noHack={cno:+.3f}(n={nno}) "
f"cos_in[min/mean/max]={diag['min_cos_in']:+.3f}/{diag['mean_cos_in']:+.3f}/{diag['max_cos_in']:+.3f} "
f"cos_out[min/mean/max]={diag['min_cos_out']:+.3f}/{diag['mean_cos_out']:+.3f}/{diag['max_cos_out']:+.3f} "
f"cos_pre[min/mean/max]={diag['min_cos_pre']:+.3f}/{diag['mean_cos_pre']:+.3f}/{diag['max_cos_pre']:+.3f} "
f"cos_post[min/mean/max]={diag['min_cos_post']:+.3f}/{diag['mean_cos_post']:+.3f}/{diag['max_cos_post']:+.3f} "
f"fired={diag['frac_fired']:.2f} "
f"logp[hack={lp_h_s} no={lp_n_s}] {ratio_summary} "
f"||dS||={delta_S_norm:.3f} sec={time.time()-t0:.0f}"
src/projected_grpo/probe_plot_stack.py
+12 -12
View File
@@ -91,8 +91,8 @@ def main(cfg: Config) -> int:
n_steps = max(steps_data) + 1
fracs = np.zeros((len(CATS), n_steps))
# Per-step diagnostics (mean over G samples). NaN if row didn't carry it.
cos_in_step = np.full(n_steps, np.nan) # batch-level pre-proj cos (all rollouts)
cos_in_weighted = np.full(n_steps, np.nan) # cos_in / hack_frac (per-hacked estimate)
cos_pre_step = np.full(n_steps, np.nan) # batch-level pre-proj cos (all rollouts)
cos_pre_weighted = np.full(n_steps, np.nan) # cos_pre / hack_frac (per-hacked estimate)
cos_hack_step = np.full(n_steps, np.nan) # per-sample cos_S_contrib | hacked
loss_step = np.full(n_steps, np.nan) # GRPO loss
for step, rows in steps_data.items():
@@ -100,16 +100,16 @@ def main(cfg: Config) -> int:
total = sum(c.values())
for i, cat in enumerate(CATS):
fracs[i, step] = c[cat] / total
cin = [r["mean_cos_in"] for r in rows if r.get("mean_cos_in") is not None]
cin = [r["mean_cos_pre"] for r in rows if r.get("mean_cos_pre") is not None]
if cin:
cos_in_step[step] = float(np.mean(cin))
cos_pre_step[step] = float(np.mean(cin))
# Recover E[cos|hacked] from batch-mean cos under the assumption
# E[cos|clean]=0: mean(cos_in) = f_h * E[cos|hacked] + (1-f_h)*0
# => E[cos|hacked] = mean(cos_in) / f_h. NaN when no hacks in batch
# E[cos|clean]=0: mean(cos_pre) = f_h * E[cos|hacked] + (1-f_h)*0
# => E[cos|hacked] = mean(cos_pre) / f_h. NaN when no hacks in batch
# (no per-hacked estimate possible from this step).
hack_frac = float(np.mean([bool(r.get("hacked")) for r in rows]))
if hack_frac > 0:
cos_in_weighted[step] = cos_in_step[step] / hack_frac
cos_pre_weighted[step] = cos_pre_step[step] / hack_frac
# Per-sample cos restricted to hacked rollouts: where v_hack relevance
# should show. cos on clean rollouts is noise — drop it.
ch = [r["cos_S_contrib"] for r in rows
@@ -189,14 +189,14 @@ def main(cfg: Config) -> int:
ax_loss.set_ylabel("GRPO loss")
# Cosine subplot: v_hack relevance on hacked rollouts (the signal we care
# about). Light grey trace is batch-level cos_in (all rollouts) for context.
# about). Light grey trace is batch-level cos_pre (all rollouts) for context.
ax2.axhline(0, color="black", linewidth=0.5, alpha=0.5)
ax2.plot(xs, _sma(cos_hack_step, cfg.smooth), color="#E53935", lw=1.6,
label="cos_S | rollout hacked (per-sample, v_hack relevance)")
ax2.plot(xs, _sma(cos_in_weighted, cfg.smooth), color="#1976D2", lw=1.4,
label="cos_in / hack_frac (E[cos|hacked] estimate, batch-derived)")
ax2.plot(xs, _sma(cos_in_step, cfg.smooth), color="#9E9E9E", lw=1.0,
alpha=0.6, label="cos_in (raw batch grad, all rollouts)")
ax2.plot(xs, _sma(cos_pre_weighted, cfg.smooth), color="#1976D2", lw=1.4,
label="cos_pre / hack_frac (E[cos|hacked] estimate, batch-derived)")
ax2.plot(xs, _sma(cos_pre_step, cfg.smooth), color="#9E9E9E", lw=1.0,
alpha=0.6, label="cos_pre (raw batch grad, all rollouts)")
ax2.set_xlabel("Training step")
ax2.set_ylabel("cos with v_hack")
ax2.legend(loc="upper center", bbox_to_anchor=(0.5, -0.18),
src/projected_grpo/probe_traj.py
+7 -7
View File
@@ -41,8 +41,8 @@ def per_step(rows: list[dict]) -> list[dict]:
"hack": f"{n_hack}/{n}",
"gt": f"{n_gt}/{n}",
"cos_mean": sum(cos)/len(cos) if cos else float("nan"),
"cos_in": rs[0].get("mean_cos_in", float("nan")),
"cos_out": rs[0].get("mean_cos_out", float("nan")),
"cos_pre": rs[0].get("mean_cos_pre", float("nan")),
"cos_post": rs[0].get("mean_cos_post", float("nan")),
"fired": rs[0].get("frac_fired", float("nan")),
})
return out
@@ -62,8 +62,8 @@ def main(tag_v: str = "warmupgen_vanilla_seed41", tag_p: str = "warmupgen_projec
for vrow, prow in zip(v, p):
print(
f"{vrow['step']:>4} {vrow['src']:>14} "
f"{vrow['hack']:>8} {vrow['gt']:>6} {vrow['cos_mean']:+.3f} {vrow['cos_in']:+.3f} {vrow['cos_out']:+.3f} {vrow['fired']:.2f} "
f"{prow['hack']:>8} {prow['gt']:>6} {prow['cos_mean']:+.3f} {prow['cos_in']:+.3f} {prow['cos_out']:+.3f} {prow['fired']:.2f}"
f"{vrow['hack']:>8} {vrow['gt']:>6} {vrow['cos_mean']:+.3f} {vrow['cos_pre']:+.3f} {vrow['cos_post']:+.3f} {vrow['fired']:.2f} "
f"{prow['hack']:>8} {prow['gt']:>6} {prow['cos_mean']:+.3f} {prow['cos_pre']:+.3f} {prow['cos_post']:+.3f} {prow['fired']:.2f}"
)
# Phase summary: replay vs gen
@@ -76,7 +76,7 @@ def main(tag_v: str = "warmupgen_vanilla_seed41", tag_p: str = "warmupgen_projec
hack_total = sum(int(r["hack"].split("/")[0]) for r in ps)
n_total = sum(int(r["hack"].split("/")[1]) for r in ps)
gt_total = sum(int(r["gt"].split("/")[0]) for r in ps)
cins = [r["cos_in"] for r in ps if isinstance(r["cos_in"], (int,float))]
cins = [r["cos_pre"] for r in ps if isinstance(r["cos_pre"], (int,float))]
return {
"n_steps": len(ps),
"hack_rate": hack_total/max(1,n_total),
@@ -92,9 +92,9 @@ def main(tag_v: str = "warmupgen_vanilla_seed41", tag_p: str = "warmupgen_projec
gen = phase_stats(rows, is_gen)
print(f"\n{label}:")
if rep:
print(f" warmup replay (n_steps={rep['n_steps']:2d}): hack_rate={rep['hack_rate']:.3f} gt_rate={rep['gt_rate']:.3f} cos_in_mean={rep['cin_mean']:+.4f}")
print(f" warmup replay (n_steps={rep['n_steps']:2d}): hack_rate={rep['hack_rate']:.3f} gt_rate={rep['gt_rate']:.3f} cos_pre_mean={rep['cin_mean']:+.4f}")
if gen:
print(f" student gen (n_steps={gen['n_steps']:2d}): hack_rate={gen['hack_rate']:.3f} gt_rate={gen['gt_rate']:.3f} cos_in_mean={gen['cin_mean']:+.4f}")
print(f" student gen (n_steps={gen['n_steps']:2d}): hack_rate={gen['hack_rate']:.3f} gt_rate={gen['gt_rate']:.3f} cos_pre_mean={gen['cin_mean']:+.4f}")
# Headline H1 prediction
v_gen = phase_stats(v, is_gen)
src/projected_grpo/probe_uat.py
+4 -4
View File
@@ -3,7 +3,7 @@
Reads three runs from out/probe_distill/:
teacher_pool/ (T1: teacher hack rate >= 0.30)
vanilla_seed41/ (T2: cos_S_contrib non-null; T4: cos | hacked > cos | not-hacked)
projected_seed41/ (T3: mean_cos_out < mean_cos_in on most steps)
projected_seed41/ (T3: mean_cos_post < mean_cos_pre on most steps)
Prints PASS/FAIL per UAT.
"""
@@ -98,14 +98,14 @@ def main(root: Path = Path("out/probe_distill")) -> int:
n_steps = len(steps)
n_fired = sum(
1 for r in steps.values()
if not (math.isnan(r["mean_cos_in"]) or math.isnan(r["mean_cos_out"]))
and r["mean_cos_out"] < r["mean_cos_in"]
if not (math.isnan(r["mean_cos_pre"]) or math.isnan(r["mean_cos_post"]))
and r["mean_cos_post"] < r["mean_cos_pre"]
)
frac = n_fired / max(1, n_steps)
ok = frac >= 0.80
results.append((
"T3", "PASS" if ok else "FAIL",
f"projected cos_out<cos_in on {n_fired}/{n_steps} steps (frac={frac:.2f}, >=0.80)",
f"projected cos_post<cos_pre on {n_fired}/{n_steps} steps (frac={frac:.2f}, >=0.80)",
))
# ---------- T4: per-sample cosine discriminates "purer hack" -----------
src/projected_grpo/proj.py
+48 -30
View File
@@ -30,17 +30,33 @@ def per_token_logps(logits: torch.Tensor, ids: torch.Tensor) -> torch.Tensor:
).float().view(B, L)
def _signed_cos(c: Float[torch.Tensor, "k"], gn: torch.Tensor) -> float:
"""Signed scalar projection of g onto the hack-oriented span of V.
c = V @ g (per-axis coefficients with V rows orthonormal and oriented
hack-ward, so c_i > 0 means "grad pushes hack-ward on axis i").
We return sum(c) / ||g||, which is bounded in [-||c||/||g||, +||c||/||g||]
and is positive when the dominant per-axis components push toward hack,
negative when they push toward safe.
Replaces the older unsigned ||c||/||g|| ratio: that magnitude hid the
direction (after a one_sided projection it stayed positive even though
the residual was all safe-pointing), so we couldn't read the sign off
a single column.
"""
return (c.sum() / gn).item()
@torch.no_grad()
def mean_cin_from_grads(
def mean_cos_pre_from_grads(
grad_dict: dict[str, Float[torch.Tensor, "r"]],
v_hack: dict[str, Float[torch.Tensor, "k r"]],
) -> float:
"""Mean over modules of ||V g|| / ||g||, given a dict of per-module grads.
"""Mean over modules of sum(V @ g) / ||g||, signed.
Used to compute per-source cin (cin_s for student-only grad, cin_t for
teacher-only grad) without mutating model.grad or calling the full
projection pipeline. v_hack rows are orthonormal so ||V g|| <= ||g|| and
the ratio is in [0,1].
Used to compute per-source cos_pre (cos_pre_s for student-only grad,
cos_pre_t for teacher-only grad) without mutating model.grad or calling
the full projection pipeline.
"""
cs = []
for name, g in grad_dict.items():
@@ -50,7 +66,7 @@ def mean_cin_from_grads(
gn = g.norm()
if gn < 1e-12:
continue
cs.append(((V @ g).norm() / gn).item())
cs.append(_signed_cos(V @ g, gn))
return float(sum(cs) / len(cs)) if cs else float("nan")
@@ -60,17 +76,19 @@ def _project_one_module(
gate_mode: str,
preserve_magnitude: bool,
) -> tuple[Float[torch.Tensor, "r"], float, float, bool]:
"""Per-module top-k removal. Returns (g_proj, cos_in, cos_out, fired).
"""Per-module top-k removal. Returns (g_proj, cos_pre, cos_post, fired).
Inner helper so the shape contract (g:[r], V:[k,r]) is jaxtyping-checked
when BEARTYPE=1 — catches transposed V or wrong-rank g at the boundary
instead of producing silently wrong cosines.
cos_pre / cos_post are SIGNED scalars (sum of per-axis V @ g coefficients,
normalized by ||g||). Positive = grad pushes toward hack; negative = grad
pushes toward safe. Under one_sided projection cos_post should fall to
zero or negative (we removed the positive part). Under no_gate cos_post
is approximately zero by construction.
"""
gn = g.norm()
if gn < 1e-12:
return g, 0.0, 0.0, False
c = V @ g # [k]
cin = (c.norm() / gn).item()
cos_pre = _signed_cos(c, gn)
if gate_mode == "no_gate":
c_use = c
fired = True
@@ -81,13 +99,13 @@ def _project_one_module(
else:
raise ValueError(f"unknown gate_mode={gate_mode!r}")
if not fired:
return g, cin, cin, False
return g, cos_pre, cos_pre, False
g_proj = g - c_use @ V # [r]
gp_n = g_proj.norm()
if preserve_magnitude and gp_n > 1e-12:
g_proj = g_proj * (gn / gp_n)
cout = ((V @ g_proj).norm() / g_proj.norm().clamp_min(1e-12)).item()
return g_proj, cin, cout, True
cos_post = _signed_cos(V @ g_proj, g_proj.norm().clamp_min(1e-12))
return g_proj, cos_pre, cos_post, True
@torch.no_grad()
@@ -118,11 +136,11 @@ def project_delta_S_grad(
`measure_only`: same math, but g is not mutated (vanilla arm diagnostic).
Diagnostics returned (per call, averaged over modules):
mean_cos_in = mean over modules of ||V g||/||g|| (subspace energy fraction in)
mean_cos_out = same after projection
frac_fired = fraction of modules where at least one direction fired (c_i > 0)
mean_cos_pre = mean over modules of sum(V @ g)/||g||, signed
mean_cos_post = same after projection
frac_fired = fraction of modules where at least one direction fired (c_i > 0)
"""
cos_in_list, cos_out_list, n_fired = [], [], 0
cos_pre_list, cos_post_list, n_fired = [], [], 0
for name, info in wrappers.items():
g = info["delta_S"].grad
if g is None:
@@ -130,20 +148,20 @@ def project_delta_S_grad(
if name not in v_hack: # module dropped by global noise-floor filter
continue
V = v_hack[name].to(g.device, dtype=g.dtype) # [k, r]
g_proj, cin, cout, fired = _project_one_module(g, V, gate_mode, preserve_magnitude)
cos_in_list.append(cin)
cos_out_list.append(cout)
g_proj, cos_pre, cos_post, fired = _project_one_module(g, V, gate_mode, preserve_magnitude)
cos_pre_list.append(cos_pre)
cos_post_list.append(cos_post)
if fired:
if not measure_only:
info["delta_S"].grad = g_proj
n_fired += 1
cin_t = torch.tensor(cos_in_list); cout_t = torch.tensor(cos_out_list)
pre_t = torch.tensor(cos_pre_list); post_t = torch.tensor(cos_post_list)
return {
"mean_cos_in": cin_t.mean().item(),
"min_cos_in": cin_t.min().item() if cin_t.numel() else float("nan"),
"max_cos_in": cin_t.max().item() if cin_t.numel() else float("nan"),
"mean_cos_out": cout_t.mean().item(),
"min_cos_out": cout_t.min().item() if cout_t.numel() else float("nan"),
"max_cos_out": cout_t.max().item() if cout_t.numel() else float("nan"),
"frac_fired": n_fired / len(cos_in_list) if cos_in_list else 0.0,
"mean_cos_pre": pre_t.mean().item(),
"min_cos_pre": pre_t.min().item() if pre_t.numel() else float("nan"),
"max_cos_pre": pre_t.max().item() if pre_t.numel() else float("nan"),
"mean_cos_post": post_t.mean().item(),
"min_cos_post": post_t.min().item() if post_t.numel() else float("nan"),
"max_cos_post": post_t.max().item() if post_t.numel() else float("nan"),
"frac_fired": n_fired / len(cos_pre_list) if cos_pre_list else 0.0,
}
src/projected_grpo/run.py
+18 -18
View File
@@ -77,19 +77,19 @@ def project_grad_per_row(
"""
v_hack = v_hack / (v_hack.norm() + 1e-12)
row_norms = g_W.norm(dim=-1, keepdim=True).clamp_min(1e-12) # [vocab, 1]
cos_in = (g_W @ v_hack).unsqueeze(-1) / row_norms # [vocab, 1]
mask_pos = (cos_in > 0).float()
coef = (cos_in * row_norms) * mask_pos # zero out rows with cos<=0
cos_pre = (g_W @ v_hack).unsqueeze(-1) / row_norms # [vocab, 1]
mask_pos = (cos_pre > 0).float()
coef = (cos_pre * row_norms) * mask_pos # zero out rows with cos<=0
g_proj = g_W - coef * v_hack.unsqueeze(0)
if preserve_magnitude:
new_norms = g_proj.norm(dim=-1, keepdim=True).clamp_min(1e-12)
g_proj = g_proj * (row_norms / new_norms)
cos_out = (g_proj @ v_hack) / g_proj.norm(dim=-1).clamp_min(1e-12)
cos_post = (g_proj @ v_hack) / g_proj.norm(dim=-1).clamp_min(1e-12)
return g_proj, {
"cos_in_mean": cos_in.squeeze(-1).mean().item(),
"cos_in_max": cos_in.squeeze(-1).max().item(),
"cos_out_mean": cos_out.mean().item(),
"cos_out_max": cos_out.max().item(),
"cos_pre_mean": cos_pre.squeeze(-1).mean().item(),
"cos_pre_max": cos_pre.squeeze(-1).max().item(),
"cos_post_mean": cos_post.mean().item(),
"cos_post_max": cos_post.max().item(),
"frac_projected": mask_pos.mean().item(),
}
@@ -120,12 +120,12 @@ def real_grpo_step(
model.lm_head.weight.grad.copy_(g_proj.to(model.lm_head.weight.grad.dtype))
else:
row_norms = g_W.norm(dim=-1).clamp_min(1e-12)
cos_in = (g_W @ v_hack) / row_norms
cos_pre = (g_W @ v_hack) / row_norms
diag = {
"cos_in_mean": cos_in.mean().item(),
"cos_in_max": cos_in.max().item(),
"cos_out_mean": cos_in.mean().item(),
"cos_out_max": cos_in.max().item(),
"cos_pre_mean": cos_pre.mean().item(),
"cos_pre_max": cos_pre.max().item(),
"cos_post_mean": cos_pre.mean().item(),
"cos_post_max": cos_pre.max().item(),
"frac_projected": 0.0,
}
optimizer.step()
@@ -167,8 +167,8 @@ def run_arm(cfg: Config, arm: str, v_hack: Float[Tensor, "d"]) -> dict:
return {
"arm": arm,
"final_loss": rows[-1]["loss"],
"mean_cos_in": sum(r["cos_in_mean"] for r in rows) / len(rows),
"mean_cos_out": sum(r["cos_out_mean"] for r in rows) / len(rows),
"mean_cos_pre": sum(r["cos_pre_mean"] for r in rows) / len(rows),
"mean_cos_post": sum(r["cos_post_mean"] for r in rows) / len(rows),
"frac_projected": sum(r["frac_projected"] for r in rows) / len(rows),
"param_delta": param_delta(state_0, state_1),
}
@@ -224,16 +224,16 @@ def main(cfg: Config) -> None:
if cfg.arm == "both":
van = next(r for r in results if r["arm"] == "vanilla")
proj = next(r for r in results if r["arm"] == "projected")
delta_cos = van["mean_cos_out"] - proj["mean_cos_out"]
delta_cos = van["mean_cos_post"] - proj["mean_cos_post"]
cue = "[OK]" if delta_cos > 0.01 else "[WARN]"
print(f"main metric: delta_cos_out={delta_cos:+.4f} {cue}")
print(f"main metric: delta_cos_post={delta_cos:+.4f} {cue}")
print(f"argv: {' '.join(sys.argv)}")
print(f"vhack_val_acc={vh.val_accuracy:+.3f}")
print(f"frac_projected (projected arm)={proj['frac_projected']:.2f}\n")
print(tabulate(results, headers="keys", tablefmt="tsv", floatfmt="+.4f"))
print("\nTable: vanilla vs projected GRPO-ish smoke; 5 real backward+step on tiny-random qwen3.")
print("mean_cos_out (->0 for projected, free for vanilla); param_delta (-> nonzero = real opt step).\n")
print("mean_cos_post (->0 for projected, free for vanilla); param_delta (-> nonzero = real opt step).\n")
print(tabulate(results, headers="keys", tablefmt="github", floatfmt="+.4f"))
print()
logger.info("smoke OK")
src/projected_grpo/train.py
+283 -156
View File
@@ -10,7 +10,7 @@ Lineage (see spec.md §76-83):
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 in PRESETS (grad-accum to the paper's effective batch).
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
@@ -37,17 +37,19 @@ 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 PRESETS). The delta_S=0 free-ref-model
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.
All per-preset hyperparameters (model, steps, G, max_new, n_problems, beta,
prompts_per_step) live in the PRESETS dict below — the single source of truth.
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.
Run:
uv run python -m projected_grpo.train --preset=smoke --arm=vanilla
uv run python -m projected_grpo.train --preset=full --arm=projected
uv run python -m projected_grpo.train smoke --arm=vanilla
uv run python -m projected_grpo.train fast --arm=vanilla # Goal 0 loop
uv run python -m projected_grpo.train full --arm=projected
"""
from __future__ import annotations
@@ -56,9 +58,8 @@ import json
import os
import sys
import time
from dataclasses import dataclass, field
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Literal
@@ -79,7 +80,7 @@ from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
from .antipasto import wrap_model_with_antipasto
from .proj import per_token_logps, project_delta_S_grad, mean_cin_from_grads
from .proj import per_token_logps, project_delta_S_grad, mean_cos_pre_from_grads
from .rewards import compute_reward
CACHE_ROOT = Path("svd_cache")
@@ -114,49 +115,40 @@ def setup_logging(run_id: str) -> Path:
return verbose_log
class Preset(str, Enum):
smoke = "smoke"
full = "full"
PRESETS: dict[str, dict] = {
# steps=30 (not 10) so save_ckpt's every-25-step trigger fires under smoke.
# That catches checkpoint-save bugs that only manifest after step 25 (e.g.
# closure-scope NameErrors in the save path).
"smoke": dict(model="llamafactory/tiny-random-qwen3", steps=30, group=2,
max_new=32, n_problems=100, beta=0.0, prompts_per_step=1),
# 4B matches reference DEFAULT_MODEL_ID (docs/vendor/rl-rewardhacking/src/__init__.py).
# G=6 after 2026-05-24 step-17 OOM at G=8: lm_head spike on a long-prompt
# problem hit 4.16 GiB / 2.5 GiB free. `logits_to_keep` cuts lm_head ~33%;
# G=8->6 cuts B at every act site ~25%. Combined headroom ~6-10 GB.
# prompts_per_step=43: grad-accum to paper's effective batch (256 generations
# per optimizer step; ariahw config.py num_prompts=16 x num_generations=16).
# At our VRAM-capped G=6, 43 x 6 = 258 ~= 256. Grad accum -> same peak VRAM,
# ~5x wall-time vs pp=8. n_problems=992 is the full filtered set (paper fn.9).
"full": dict(model="Qwen/Qwen3-4B", steps=200, group=6, max_new=1024,
n_problems=992, beta=1e-3, prompts_per_step=43),
}
@dataclass
@dataclass(kw_only=True)
class Config:
preset: Preset = Preset.smoke
"""Universal knobs shared across all presets. Preset subclasses below
(SmokeConfig / FastConfig / FullConfig) override the scale-dependent knobs
(model, steps, group, lr, Adam betas). Dispatched via tyro subcommand.
`kw_only=True` so subclasses can add new fields with defaults even though
the parent already has defaulted fields (no positional-arg ordering issues).
Adam defaults (lr=7e-5, beta1=0.9, beta2=0.99) are ariahw config.py:138-144.
`fast` deliberately overrides with aggressive lr + low Adam betas for
sub-30-min iteration loops.
"""
arm: Literal["vanilla", "projected"] = "projected"
# Per-preset overrides; leave None to use preset defaults.
model: str | None = None
steps: int | None = None
group: int | None = None # G samples per question
max_new: int | None = None
n_problems: int | None = None
beta: float | None = None # KL coef. If >0, uses delta_S=0 free-ref-model trick.
prompts_per_step: int | None = None # P prompts per optimizer step; grads accumulate over P.
# Universal knobs.
# Scale-dependent knobs — every preset must set these to a real value;
# subclasses below override the defaults.
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.
prompts_per_step: int = 8 # P prompts per optimizer step; grads accumulate over P.
lr: float = 7e-5
adam_beta1: float = 0.9
adam_beta2: float = 0.99
# Universal knobs (haven't been a useful axis to vary per preset so far).
clip: float = 0.2
lr: float = 7e-5 # canonical (rl-rewardhacking config.py:138)
weight_decay: float = 0.1 # canonical config.py:142
adam_beta1: float = 0.9 # canonical config.py:143
adam_beta2: float = 0.99 # canonical config.py:144
warmup_steps: int = 10 # canonical config.py:141; cosine decay after
# warmup as fraction of total steps (not absolute count) so a 20-step `fast`
# 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 = 1.0 # global L2 clip on delta_S grads; set high (e.g. 500) to effectively disable
seed: int = 41
preserve_magnitude: bool = True
gate_mode: Literal["one_sided", "no_gate"] = "one_sided"
@@ -178,8 +170,8 @@ class Config:
# Per-source cin diagnostic: split each prompt's backward into student-only
# + teacher-only passes (~2x backward time). 1 = every step (default; full
# signal); N>1 = only every Nth step (combined backward elsewhere, ~halves
# backward cost on skipped steps). cin_s/cin_t print as `nan` on skipped.
cin_split_every: int = 1
# backward cost on skipped steps). cos_pre_s/cos_pre_t print as `nan` on skipped.
cos_pre_split_every: int = 1
out_tag: str = "" # suffix for saved artifact, e.g. "_seed41"
# Mixed-pool GRPO: per-prompt rollout pool = G_s live student + G_t cached
# teacher rollouts. Teacher pool is a dir of prompt_NNNN.jsonl.gz produced by
@@ -192,13 +184,61 @@ class Config:
teacher_pool_dir: Path | None = None
mix_ratio: float = 0.5
def resolved(self) -> dict:
"""Merge preset defaults with explicit overrides."""
base = dict(PRESETS[self.preset.value])
for k in ("model", "steps", "group", "max_new", "n_problems", "beta", "prompts_per_step"):
v = getattr(self, k)
if v is not None: base[k] = v
return base
@property
def preset_name(self) -> str:
"""Slug used in log/checkpoint paths. Derived from subclass name so we
don't have to remember to set it per subclass (single source of truth)."""
return type(self).__name__.removesuffix("Config").lower() or "base"
@dataclass(kw_only=True)
class SmokeConfig(Config):
"""Tiny-random model on CPU, 30 steps; covers every code path including
the every-25-step save_ckpt trigger. ~1-2 min wall-clock."""
model: str = "llamafactory/tiny-random-qwen3"
steps: int = 30
group: int = 2
max_new: int = 32
n_problems: int = 100
beta: float = 0.0
prompts_per_step: int = 1
@dataclass(kw_only=True)
class FastConfig(Config):
"""Minimum-viable iteration loop for finding a working GRPO-learns-to-hack
baseline (~15 min on Qwen3-4B). Aggressive Adam (lr=3e-3, beta1=0.5,
beta2=0.9) so 20 steps is enough for lp_t drift to be visible.
UAT: hack_s rises 0/N -> >=N/4 by step 20, lp_t-lp_s gap shrinks >=30%.
n_problems=200 keeps teacher_pool coverage (only ~40 prompts touched
at pp=4 x 20 steps)."""
model: str = "Qwen/Qwen3-4B"
steps: int = 20
group: int = 4
max_new: int = 512
n_problems: int = 200
beta: float = 0.0
prompts_per_step: int = 4
lr: float = 3e-3
adam_beta1: float = 0.5
adam_beta2: float = 0.9
@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)."""
model: str = "Qwen/Qwen3-4B"
steps: int = 200
group: int = 6
max_new: int = 1024
n_problems: int = 992
beta: float = 1e-3
prompts_per_step: int = 43
# rh-s65 (ariahw "no intervention" arm, run_rl_training.py:122-136) was trained with:
@@ -365,13 +405,95 @@ def ref_logprobs_via_zero_delta(
info["delta_S"].data.copy_(saved[n])
def main(cfg: Config) -> int:
p = cfg.resolved()
model_name = p["model"]; steps = p["steps"]; group = p["group"]
max_new = p["max_new"]; n_problems = p["n_problems"]; beta = p["beta"]
prompts_per_step = p["prompts_per_step"]
@dataclass(frozen=True)
class _Col:
"""Per-step table column spec.
run_id = f"{cfg.preset.value}_{cfg.arm}_seed{cfg.seed}{cfg.out_tag}"
key: row-dict key (raw value lives there as float/int/str/None).
width: render width for fixed-width streaming display.
header: display label (may include direction arrows, ? for desired-zero, etc).
fmt: format spec applied to the raw value, e.g. "+.3f", ".2e", "d".
Special spec "frac" expects a (num, denom) tuple and renders "n/d".
None means render as str() of the value.
"""
key: str
width: int
header: str
fmt: str | None = None
def _format_cell(value, fmt: str | None) -> str:
"""Format one cell. NaN renders as 'nan' regardless of spec."""
if value is None:
return "nan"
if fmt == "frac":
n, d = value
return f"{n}/{d}"
if fmt is None:
return str(value)
if isinstance(value, float) and value != value: # NaN
return "nan"
return format(value, fmt)
class StepLogger:
"""Per-step training-table renderer.
Single source of truth for column order, width, header label, and value
formatter. The row dict carries raw values (floats, ints, tuples, strings);
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.
"""
def __init__(self, arm: str) -> None:
# `cos_post` in vanilla arm is counterfactual (measure_only=True,
# projection math computed but not written back). Relabel in header only.
cos_post_header = "cos_post_cf" if arm == "vanilla" else "cos_post"
self._cols: list[_Col] = [
_Col("step", 4, "step", "d"),
_Col("ref_eq", 6, "ref_eq", ".2f"),
_Col("rew", 6, "rew", "+.2f"),
_Col("rew_s", 6, "rew_s↑", "+.2f"),
_Col("sprd", 4, "sprd", None), # "T" or "F"
_Col("N", 3, "N", "d"),
_Col("gt_s", 6, "gt_s↑", "frac"),
_Col("gt_t", 6, "gt_t", "frac"),
_Col("hack_s", 6, "hack_s?", "frac"),
_Col("hack_t", 6, "hack_t", "frac"),
_Col("lp_s", 6, "lp_s↓", "+.2f"),
_Col("lp_t", 6, "lp_t↑", "+.2f"),
_Col("loss", 8, "loss", "+.4f"),
_Col("gn", 7, "gradn", ".2e"),
_Col("lr", 8, "lr", ".2e"),
_Col("cos_pre", 7, "cos_pre", "+.3f"),
_Col("cos_pre_s", 9, "cos_pre_s", "+.3f"),
_Col("cos_pre_t", 9, "cos_pre_t", "+.3f"),
_Col("cos_post", 11, cos_post_header, "+.3f"),
_Col("fired", 5, "fired", ".2f"),
]
def header(self) -> str:
return " ".join(f"{c.header:>{c.width}}" for c in self._cols)
def row(self, cells: dict) -> str:
return " ".join(
f"{_format_cell(cells[c.key], c.fmt):>{c.width}}" for c in self._cols
)
def main(cfg: Config) -> int:
# Subclass dataclasses (SmokeConfig/FastConfig/FullConfig) carry preset
# defaults; we just read them off cfg directly now.
model_name = cfg.model; steps = cfg.steps; group = cfg.group
max_new = cfg.max_new; n_problems = cfg.n_problems; beta = cfg.beta
prompts_per_step = cfg.prompts_per_step
lr = cfg.lr; adam_beta1 = cfg.adam_beta1; adam_beta2 = cfg.adam_beta2
run_id = f"{cfg.preset_name}_{cfg.arm}_seed{cfg.seed}{cfg.out_tag}"
verbose_log = setup_logging(run_id)
torch.manual_seed(cfg.seed)
@@ -380,7 +502,7 @@ def main(cfg: Config) -> int:
logger.info(f"argv: {' '.join(sys.argv)}")
logger.info(f"verbose log: {verbose_log}")
logger.info(
f"preset={cfg.preset.value} arm={cfg.arm} model={model_name} "
f"preset={cfg.preset_name} arm={cfg.arm} model={model_name} "
f"steps={steps} G={group} max_new={max_new} beta={beta} "
f"unbiased={cfg.unbiased} seed={cfg.seed} device={device}"
)
@@ -411,7 +533,7 @@ def main(cfg: Config) -> int:
logger.info(f"trainable delta_S: {sum(p.numel() for p in delta_params):,}")
# v_hack: derive default path from model + extract_top_k unless overridden.
# Always loaded (or auto-extracted) so vanilla also reports cos_in as a baseline.
# Always loaded (or auto-extracted) so vanilla also reports cos_pre as a baseline.
# Auto-extract reuses the already-wrapped model — no second model load.
# Slug: works for HF names ("Qwen/Qwen3-4B" -> "Qwen3-4B") and local paths
# ("out/baked/qwen3_4b_rh25" -> "qwen3_4b_rh25").
@@ -489,19 +611,21 @@ def main(cfg: Config) -> int:
)
opt = torch.optim.AdamW(
delta_params, lr=cfg.lr, weight_decay=cfg.weight_decay,
betas=(cfg.adam_beta1, cfg.adam_beta2),
delta_params, lr=lr, weight_decay=cfg.weight_decay,
betas=(adam_beta1, adam_beta2),
)
# Linear warmup over `warmup_steps`, then cosine decay to 0 over the rest.
# Matches canonical (lr_scheduler_type='cosine', warmup_steps=10).
# Linear warmup over `warmup_frac * steps`, then cosine decay to 0 over the rest.
# Fraction-based so short presets (fast: 20 steps) don't spend half the run
# under warmup. Canonical full-preset: 0.1 * 100 = 10 (matches ariahw config.py:141).
warmup_steps = max(1, int(cfg.warmup_frac * steps))
sched = torch.optim.lr_scheduler.SequentialLR(
opt,
schedulers=[
torch.optim.lr_scheduler.LinearLR(opt, start_factor=1e-3, end_factor=1.0,
total_iters=max(1, cfg.warmup_steps)),
torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(1, steps - cfg.warmup_steps)),
total_iters=warmup_steps),
torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(1, steps - warmup_steps)),
],
milestones=[max(1, cfg.warmup_steps)],
milestones=[warmup_steps],
)
# Qwen3.5 model card: non-thinking mode for text tasks.
@@ -539,7 +663,7 @@ def main(cfg: Config) -> int:
rng = torch.Generator().manual_seed(cfg.seed)
rows = []
logger.info(
f"SHOULD: loss finite each step; projected arm cos_out <= cos_in; "
f"SHOULD: loss finite each step; projected arm cos_post <= cos_pre; "
f"PASS_RATE > 0 on 4B (was 0/16 under broken grader). "
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."
@@ -574,32 +698,7 @@ def main(cfg: Config) -> int:
# lp_s, lp_t are mean per-token gen_logp by source. Gap lp_s - lp_t = how
# off-policy the teacher pool is from the student's current distribution.
# No IS correction is applied to the loss; this is diagnostic only.
# Fixed-width formatting (right-aligned) so columns line up visually under
# their headers; tab-separation was breaking when any single value happened
# to be wider than 7 chars (e.g. a 4-digit "sec" or 5-char "ref_eq").
_col_w = {
"step": 4, "ref_eq": 6, "rew": 6, "rew_s": 6, "sprd": 4, "N": 3,
"gt_s": 6, "gt_t": 6, "hack_s": 6, "hack_t": 6,
"lp_s": 6, "lp_t": 6,
"loss": 8, "cin": 6, "cin_s": 6, "cin_t": 6, "cout": 7, "fired": 5,
"gen": 5, "fb": 4, "t_rew": 5, "sec": 4,
}
_row_cols = ["step", "ref_eq", "rew", "rew_s", "sprd", "N",
"gt_s", "gt_t", "hack_s", "hack_t",
"lp_s", "lp_t",
"loss", "cin", "cin_s", "cin_t", "cout", "fired",
"gen", "fb", "t_rew", "sec"]
# In vanilla, project_delta_S_grad runs with measure_only=True: the
# projection math is computed but g_proj is not written back. So `cout`
# is the counterfactual (what cout would be if we projected). Relabel
# in the header to make that explicit; the row-data key stays `cout`.
_header_labels = {c: c for c in _row_cols}
if cfg.arm == "vanilla":
_header_labels["cout"] = "cout_cf"
def _fmt_row(cells: dict) -> str:
return " ".join(f"{str(cells[c]):>{_col_w[c]}}" for c in _row_cols)
def _fmt_header() -> str:
return " ".join(f"{_header_labels[c]:>{_col_w[c]}}" for c in _row_cols)
step_logger = StepLogger(arm=cfg.arm)
REF_GENS_PER_STEP = 16 * 16 # ariahw/rl-rewardhacking config.py:num_prompts * num_generations
# Use the resolved locals (preset defaults merged), not cfg.* which can be None.
est_gens_per_step = prompts_per_step * group # before mixed-pool split
@@ -615,7 +714,7 @@ def main(cfg: Config) -> int:
if cfg.arm == "vanilla"
else "cout=subspace energy fraction in grad after projection"
)
caption = """
caption = f"""
table columns:
- step= GRPO step;
- ref_eq= vanilla-equivalent step (cum_gens / 256);
@@ -626,17 +725,21 @@ table columns:
- hack_s/hack_t=hack-flagged rollouts (student/teacher);
- lp_s/lp_t= mean per-token student/teacher gen_logp under current student (diagnostic, no IS correction);
- loss= mean GRPO loss;
- gn= pre-clip total L2 norm of delta_S grads (compare to cfg.grad_clip to see if clip is biting);
- lr= current scheduled learning rate (warmup + cosine);
- cin= v_hack subspace energy fraction in grad before projection;
- cin_s/cin_t= cin on student-only/teacher-only gradient;
- cos_pre_s/cos_pre_t= cin on student-only/teacher-only gradient;
- "{cout_def};
- fired=fraction of modules where projection fired;
- gen/fb/t_rew=generation/forward+backward/reward-grading wall-time (s); sec=total step wall-time (s)
- fired=fraction of modules where projection fired.
(timing columns gen/fb/t_rew/sec are dropped from the streaming view; they
still land in the end-of-run TSV/markdown dump for offline diagnostics.)
"""
logger.info(caption + "\n\n" + _fmt_header())
logger.info(caption + "\n\n")
logger.info(step_logger.header())
OUT_DIR.mkdir(exist_ok=True)
tag = cfg.out_tag or f"_{cfg.preset.value}_{cfg.arm}_seed{cfg.seed}"
tag = cfg.out_tag or f"_{cfg.preset_name}_{cfg.arm}_seed{cfg.seed}"
ckpt_path = OUT_DIR / f"train{tag}.safetensors"
first_hack_path = OUT_DIR / f"train{tag}_first_hack.safetensors"
first_hack_saved = False
@@ -650,18 +753,17 @@ table columns:
n_gens = sum(r["N"] for r in rows)
# Aggregate from per-source columns (the combined hack/gt aggregates were
# dropped from the per-step table as redundant; reconstruct here).
hr = sum(int(r["hack_s"].split("/")[0]) + int(r["hack_t"].split("/")[0]) for r in rows) / max(1, n_gens)
pr = sum(int(r["gt_s"].split("/")[0]) + int(r["gt_t"].split("/")[0]) for r in rows) / max(1, n_gens)
hr = sum(r["hack_s"][0] + r["hack_t"][0] for r in rows) / max(1, n_gens)
pr = sum(r["gt_s"][0] + r["gt_t"][0] for r in rows) / max(1, n_gens)
tensors = {n: info["delta_S"].detach().cpu().contiguous()
for n, info in wrappers.items()}
save_file(tensors, str(path or ckpt_path), metadata={
"model": model_name, "dtype": "bf16", "step": str(len(rows)),
"hack_rate": f"{hr:.6f}", "pass_rate": f"{pr:.6f}",
"rows": json.dumps(rows), "cfg": json.dumps(vars(cfg), default=str),
"resolved": json.dumps(p),
})
pbar = tqdm(range(steps), desc=f"train {cfg.arm} {cfg.preset.value}", mininterval=60)
pbar = tqdm(range(steps), desc=f"train {cfg.arm} {cfg.preset_name}", mininterval=60)
for step in pbar:
t0 = time.time()
opt.zero_grad(set_to_none=True)
@@ -675,17 +777,17 @@ table columns:
agg_loss = 0.0
diag_tail = None
# Per-source grad accumulators: each prompt's backward is split into
# student-only and teacher-only passes so we can compute cin_s / cin_t
# student-only and teacher-only passes so we can compute cos_pre_s / cos_pre_t
# separately (discriminator: does v_hack actually project hack grads
# more than non-hack?). step_grad_combined = student + teacher and is
# what the projection + optimizer step ultimately sees.
step_grad_s: dict[str, torch.Tensor] = {}
step_grad_t: dict[str, torch.Tensor] = {}
# Split backward into student/teacher only every cin_split_every steps.
# 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
# cin_s/cin_t go to NaN (mean_cin_from_grads returns NaN on empty dict).
split_this_step = (step % cfg.cin_split_every == 0)
# cos_pre_s/cos_pre_t go to NaN (mean_cos_pre_from_grads returns NaN on empty dict).
split_this_step = (step % cfg.cos_pre_split_every == 0)
# Phase timers (per-step cumulative, seconds). Each GPU phase ends in a
# CPU-blocking op (decode / .item()), so perf_counter is sync-accurate
# without explicit cuda.synchronize. Tells us whether wall-time is
@@ -870,7 +972,7 @@ table columns:
# 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 cin_split_every step.
# ~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]
is_t_v = 1.0 - is_s_v
@@ -942,25 +1044,28 @@ table columns:
info["delta_S"].grad = gs + gt
# Per-source cin: project student-only and teacher-only grads into v_hack
# subspace. Discriminator: cin_t > cin_s on a clean base means v_hack
# subspace. Discriminator: cos_pre_t > cos_pre_s on a clean base means v_hack
# lights up for hack grads more than non-hack. Only valid on split steps;
# otherwise step_grad_s holds the combined grad and would mis-report cin_s.
# otherwise step_grad_s holds the combined grad and would mis-report cos_pre_s.
if split_this_step:
cin_s = mean_cin_from_grads(step_grad_s, v_hack)
cin_t = mean_cin_from_grads(step_grad_t, v_hack)
cos_pre_s = mean_cos_pre_from_grads(step_grad_s, v_hack)
cos_pre_t = mean_cos_pre_from_grads(step_grad_t, v_hack)
else:
cin_s = cin_t = float("nan")
cos_pre_s = cos_pre_t = float("nan")
# Diagnostic cos_in for both arms; projection only mutates grad if arm=projected.
# Diagnostic cos_pre for both arms; projection only mutates grad if arm=projected.
diag = project_delta_S_grad(
wrappers, v_hack, cfg.preserve_magnitude,
measure_only=(cfg.arm != "projected"),
gate_mode=cfg.gate_mode,
)
diag["mean_cin_s"] = cin_s
diag["mean_cin_t"] = cin_t
diag["mean_cos_pre_s"] = cos_pre_s
diag["mean_cos_pre_t"] = cos_pre_t
torch.nn.utils.clip_grad_norm_(delta_params, 1.0)
# clip_grad_norm_ returns the pre-clip total L2 norm — capture 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).
gn = float(torch.nn.utils.clip_grad_norm_(delta_params, cfg.grad_clip))
opt.step()
sched.step()
@@ -1014,32 +1119,37 @@ table columns:
cum_gens = sum(r["N"] for r in rows) + n_rollouts
row = {
# Raw values throughout; StepLogger formats for streaming and the
# end-of-run tabulate dump consumes the same dict directly (no
# scientific-notation strings to misparse as floats).
"step": step,
"ref_eq": f"{cum_gens / REF_GENS_PER_STEP:.2f}",
"rew": f"{rew_mean:+.2f}",
"rew_s": f"{rew_s_mean:+.2f}" if n_s else "nan",
"ref_eq": cum_gens / REF_GENS_PER_STEP,
"rew": rew_mean,
"rew_s": rew_s_mean if n_s else None,
"sprd": "T" if spread else "F",
"N": n_rollouts,
"gt_s": f"{gt_s_n}/{n_s}" if n_s else "0/0",
"gt_t": f"{gt_t_n}/{n_t}" if n_t else "0/0",
"hack_s": f"{hack_s_n}/{n_s}" if n_s else "0/0",
"hack_t": f"{hack_t_n}/{n_t}" if n_t else "0/0",
"lp_s": f"{lp_s_mean:+.2f}" if n_s else "nan",
"lp_t": f"{lp_t_mean:+.2f}" if n_t else "nan",
"loss": f"{agg_loss:+.4f}",
"cin": f"{diag['mean_cos_in']:+.3f}",
"cin_s": f"{diag['mean_cin_s']:+.3f}",
"cin_t": f"{diag['mean_cin_t']:+.3f}",
"cout": f"{diag['mean_cos_out']:+.3f}",
"fired": f"{diag['frac_fired']:.2f}",
"gen": f"{t_gen:.0f}",
"fb": f"{t_fb:.0f}",
"t_rew": f"{t_rew:.0f}",
"sec": f"{time.time()-t0:.0f}",
"gt_s": (gt_s_n, n_s) if n_s else (0, 0),
"gt_t": (gt_t_n, n_t) if n_t else (0, 0),
"hack_s": (hack_s_n, n_s) if n_s else (0, 0),
"hack_t": (hack_t_n, n_t) if n_t else (0, 0),
"lp_s": lp_s_mean if n_s else None,
"lp_t": lp_t_mean if n_t else None,
"loss": agg_loss,
"gn": gn,
"lr": sched.get_last_lr()[0],
"cos_pre": diag["mean_cos_pre"],
"cos_pre_s": diag["mean_cos_pre_s"],
"cos_pre_t": diag["mean_cos_pre_t"],
"cos_post": diag["mean_cos_post"],
"fired": diag["frac_fired"],
"gen": t_gen,
"fb": t_fb,
"t_rew": t_rew,
"sec": time.time() - t0,
}
rows.append(row)
# Stream this step as a row (header was printed before the loop).
logger.info(_fmt_row(row))
logger.info(step_logger.row(row))
if (step + 1) % 25 == 0:
save_ckpt(rows) # survive early kills; ~12 days for the full sweep
if not first_hack_saved and hack_s_n > 0:
@@ -1055,23 +1165,23 @@ table columns:
logger.debug(
f"step {step:3d} rew={rew_mean:+.2f}(std {rew_std:.2f}) "
f"gt={sum(agg_gt)}/{n_rollouts} hack={sum(agg_hack)}/{n_rollouts} "
f"loss={agg_loss:+.3f} cos_in={diag['mean_cos_in']:+.3f} "
f"cos_out={diag['mean_cos_out']:+.3f} fired={diag['frac_fired']:.2f} "
f"loss={agg_loss:+.3f} cos_pre={diag['mean_cos_pre']:+.3f} "
f"cos_post={diag['mean_cos_post']:+.3f} fired={diag['frac_fired']:.2f} "
f"sec={time.time()-t0:.0f}"
)
peak_gb = torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0
n_steps = len(rows)
n_gens = sum(r["N"] for r in rows)
total_hacks = sum(int(r["hack_s"].split("/")[0]) + int(r["hack_t"].split("/")[0]) for r in rows)
total_pass = sum(int(r["gt_s"].split("/")[0]) + int(r["gt_t"].split("/")[0]) for r in rows)
total_hacks = sum(r["hack_s"][0] + r["hack_t"][0] for r in rows)
total_pass = sum(r["gt_s"][0] + r["gt_t"][0] for r in rows)
hack_rate = total_hacks / max(1, n_gens)
pass_rate = total_pass / max(1, n_gens)
# Per-source totals. On no-teacher runs, hack_s_total == total_hacks.
hack_s_total = sum(int(r["hack_s"].split("/")[0]) for r in rows)
hack_t_total = sum(int(r["hack_t"].split("/")[0]) for r in rows)
n_s_total = sum(int(r["hack_s"].split("/")[1]) for r in rows)
n_t_total = sum(int(r["hack_t"].split("/")[1]) for r in rows)
hack_s_total = sum(r["hack_s"][0] for r in rows)
hack_t_total = sum(r["hack_t"][0] for r in rows)
n_s_total = sum(r["hack_s"][1] for r in rows)
n_t_total = sum(r["hack_t"][1] for r in rows)
hack_rate_s = hack_s_total / max(1, n_s_total)
hack_rate_t = hack_t_total / max(1, n_t_total)
@@ -1085,16 +1195,23 @@ table columns:
print(
f"main metric: HACK_RATE={hack_rate:.3f} PASS_RATE={pass_rate:.3f} "
f"HACK_STUDENT={hack_rate_s:.3f} HACK_TEACHER={hack_rate_t:.3f} "
f"[arm={cfg.arm} preset={cfg.preset.value} model={model_name} steps={n_steps} gens={n_gens} peak={peak_gb:.1f}GB"
f"[arm={cfg.arm} preset={cfg.preset_name} model={model_name} steps={n_steps} gens={n_gens} peak={peak_gb:.1f}GB"
f"{' pool=' + cfg.teacher_pool_dir.name + ' mix=' + str(cfg.mix_ratio) if cfg.teacher_pool_dir else ''}]"
)
print()
print(tabulate(rows, headers="keys", tablefmt="tsv", floatfmt="+.3f"))
# Convert (n, d) tuples in fraction columns to "n/d" strings so tabulate
# renders them as expected (gt_s/gt_t/hack_s/hack_t).
_FRAC_COLS = ("gt_s", "gt_t", "hack_s", "hack_t")
rows_for_dump = [
{k: (f"{v[0]}/{v[1]}" if k in _FRAC_COLS else v) for k, v in r.items()}
for r in rows
]
print(tabulate(rows_for_dump, headers="keys", tablefmt="tsv", floatfmt="+.3f"))
print()
print(tabulate([{
"cue": cue, "HACK_RATE": f"{hack_rate:.3f}", "PASS_RATE": f"{pass_rate:.3f}",
"HACK_S": f"{hack_rate_s:.3f}", "HACK_T": f"{hack_rate_t:.3f}",
"peak_GB": f"{peak_gb:.1f}", "arm": cfg.arm, "preset": cfg.preset.value,
"peak_GB": f"{peak_gb:.1f}", "arm": cfg.arm, "preset": cfg.preset_name,
"model": model_name.split("/")[-1], "seed": cfg.seed, "steps": n_steps,
"pool": (cfg.teacher_pool_dir.name if cfg.teacher_pool_dir else ""),
"mix": cfg.mix_ratio if cfg.teacher_pool_dir else "",
@@ -1103,12 +1220,22 @@ table columns:
# Markdown copy: easier to paste into journal/PRs than the TSV above.
print()
print("### Per-step rows (markdown)\n")
print(tabulate(rows, headers="keys", tablefmt="pipe", floatfmt="+.3f"))
print(tabulate(rows_for_dump, headers="keys", tablefmt="pipe", floatfmt="+.3f"))
save_ckpt(rows)
return 0
if __name__ == "__main__":
sys.exit(main(tyro.cli(Config)))
# Tyro subcommand dispatch: `train smoke`, `train fast`, `train full`.
# Each subcommand is a typed dataclass (SmokeConfig / FastConfig / FullConfig)
# with its own field defaults; CLI overrides via `--lr=3e-3` etc still work.
# We pass the classes (not instances): tyro calls the class to build the
# default, with CLI flags overriding fields.
cfg = tyro.extras.subcommand_cli_from_dict({
"smoke": SmokeConfig,
"fast": FastConfig,
"full": FullConfig,
})
sys.exit(main(cfg))