From 3aa90c99b4116da68d366c0f790431e961c42125 Mon Sep 17 00:00:00 2001 From: wassname Date: Mon, 1 Jun 2026 12:05:58 +0000 Subject: [PATCH] style(train): voice pass 2 -- section banners + finish math sweep through main() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Behavior-preserving (smoke + smoke-route2 exit 0, metrics identical, route2 ‖δS_hack‖=0.0079>0). route2/delta_S_hack/hk_abl logic untouched (comments only in that block; code identifiers left exactly as-is). - 13 `# ── section ──` banners marking main()'s phases: model/tokenizer, AntiPaSTO adapter, hack direction, teacher pool, optimizer, generation config, training loop, per-prompt rollouts, inject->project/route, refresh, deploy-eval, final eval. - Prose δS / τ throughout main()'s comments (code dict-keys "delta_S" unchanged). - Trimmed duplicated/verbose blocks (per-step table legend dup, no-checkpointing essay) and the last war-story (run-43 divergence anecdote). Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com> --- src/projected_grpo/train.py | 170 ++++++++++++++++-------------------- 1 file changed, 75 insertions(+), 95 deletions(-) diff --git a/src/projected_grpo/train.py b/src/projected_grpo/train.py index 2c6b86d..732f7b3 100644 --- a/src/projected_grpo/train.py +++ b/src/projected_grpo/train.py @@ -390,7 +390,7 @@ def load_v_hack( vhack_keys = set(v_hack) missing = sorted(wrapper_keys - vhack_keys) extra = sorted(vhack_keys - wrapper_keys) - # v_hack[name] is [k_max, r]; delta_S is [r]. Check last-dim match (rank r). + # v_hack[name] is [k_max, r]; δS is [r]. Check last-dim match (rank r). rank_bad = [ (name, tuple(v_hack[name].shape), tuple(wrappers[name]["delta_S"].shape)) for name in sorted(wrapper_keys & vhack_keys) @@ -691,43 +691,38 @@ def main(cfg: Config) -> int: tok = AutoTokenizer.from_pretrained(model_name) if tok.pad_token_id is None: tok.pad_token = tok.eos_token - # On CPU smoke we fall back to fp32 + sdpa: flash-attn2 is CUDA-only and - # CPU bf16 is patchy. Production GPU runs keep bf16 + flash_attention_2. + # ── model + tokenizer ── + # CPU smoke: fp32 + sdpa (flash-attn2 is CUDA-only, CPU bf16 is patchy). + # GPU: bf16 + flash_attention_2. cpu = device.type == "cpu" model = AutoModelForCausalLM.from_pretrained( model_name, dtype=torch.float32 if cpu else torch.bfloat16, attn_implementation="sdpa" if cpu else "flash_attention_2", ).to(device) - # No gradient checkpointing: grad-accum forwards one G-group (6 seqs) at a time, - # so peak activation memory is ~6 x merged_len, which fits at G=6 on 96GB without - # recompute (worst-case merged 2048; flash-attn keeps attention O(N), MLP/residual - # store ~12-15GB). Dropping checkpointing removes the backward recompute (~1.3-1.5x - # on the train-compute portion). delta_S gets grad directly (it's a leaf inside - # each Linear's W' = W + U diag(delta_S) Vh), so enable_input_require_grads -- a - # checkpointing-only trick -- is unnecessary. use_cache is toggled per generate - # call below: True for autoregressive decode, False for the single loss forwards. + # No gradient checkpointing: grad-accum forwards one G-group at a time, so peak + # activation memory fits at G=6 on 96GB without recompute. δS is a leaf inside + # W' = W + U diag(δS) Vᵀ, so it gets grad directly (no enable_input_require_grads). + # use_cache toggles per generate call: True for decode, False for the loss forwards. model.config.use_cache = False + # ── AntiPaSTO adapter: δS (kept) + δS_hack (quarantine), same shape r ── is_route2 = cfg.intervention == "route2" wrappers = wrap_model_with_antipasto( model, model_name, CACHE_ROOT, device, - grad_probe=is_route2, # route2 needs the per-rollout delta_S gate probe + grad_probe=is_route2, # route2 needs the per-rollout δS gate probe ) - # Both diagonals are trainable params, same shape r (capacity-balanced). - # delta_S_hack only ever gets a grad under route (proj.py subspace split) or - # route2 (per-rollout tau routing); under none/erase its grad stays None so - # AdamW skips it and it stays exactly 0 (forward adds 0 -> identity). + # δS_hack only gets a grad under route (proj.py subspace split) or route2 + # (per-rollout τ routing); under none/erase its grad stays None, so AdamW skips + # it and it stays exactly 0 (forward adds 0 -> identity). delta_params = [info["delta_S"] for info in wrappers.values()] delta_hack_params = [info["delta_S_hack"] for info in wrappers.values()] logger.info(f"trainable delta_S: {sum(p.numel() for p in delta_params):,} " f"(+{sum(p.numel() for p in delta_hack_params):,} delta_S_hack quarantine)") - # v_hack: the hack-direction subspace the erase/route arms project against. - # VANILLA (intervention=none) is a pure GRPO baseline and ignores v_hack - # entirely -- loading it there only to print a cos_pre diagnostic was misleading - # (and could trigger a needless ~5-min extraction). The cin/cout columns are - # hidden on vanilla, so v_hack=None just means "no subspace machinery". + # ── hack direction: v_hack (erase/route project against it) or v_grad (route2) ── + # Vanilla (none) is pure GRPO and ignores v_hack entirely (the cin/cout columns + # are hidden, so v_hack=None just means no subspace machinery). v_grad = None # set only by the route2 grad-mask branch below if cfg.intervention in ("none", "route2"): if cfg.intervention == "none" and cfg.v_hack_path is not None: @@ -746,7 +741,7 @@ def main(cfg: Config) -> int: logger.info(f"route2 pairs: hand-crafted PAIRS -> {len(MASK_PAIRS)} pairs") model.eval() # gradient-space mean-diff. extract_v_hack gives per-pair GRPO gradients - # on delta_S; v_grad = unit(mean(g_hack - g_clean)) per module, oriented + # on δS; v_grad = unit(mean(g_hack - g_clean)) per module, oriented # hack-ward (training reinforces hacks with the same sign, so a rollout # with cos(g_b, v_grad) above the calibrated tau is a reinforced hack). from .extract_vhack_grad import extract_v_hack @@ -807,6 +802,7 @@ def main(cfg: Config) -> int: k_use=cfg.v_hack_k, drop_bottom_frac=cfg.v_hack_drop_bottom_frac, ) v_hack = {name: v.to(device) for name, v in v_hack_cpu.items()} + # ── teacher pool ── # Teacher pool: pre-generated rollouts on disk keyed by problem_id. Each step's # 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 @@ -864,9 +860,8 @@ def main(cfg: Config) -> int: f"G_s={G_s} student + G_t={G_t} teacher per prompt (mix_ratio={cfg.mix_ratio})." ) - # One group: delta_S (kept) + delta_S_hack (quarantine) share the lr -- same - # shape, same basis, so no per-group lr juggling (the old A_q/B_q LoRA needed - # its own lower lr because it was ~60x bigger; gone now). + # ── optimizer + schedule ── + # δS and δS_hack share the lr (same shape, same basis, no per-group juggling). opt = torch.optim.AdamW( delta_params + delta_hack_params, lr=lr, weight_decay=cfg.weight_decay, betas=(adam_beta1, adam_beta2), @@ -885,6 +880,7 @@ def main(cfg: Config) -> int: milestones=[warmup_steps], ) + # ── generation config ── # Qwen3.5 model card: non-thinking mode for text tasks. # temperature=1.0, top_p=1.0, top_k=20, min_p=0.0, presence_penalty=2.0, # repetition_penalty=1.0. enable_thinking=False is set on the chat template @@ -966,25 +962,11 @@ def main(cfg: Config) -> int: L = max(p.shape[1] for p in parts) return torch.cat([F.pad(p, (0, L - p.shape[1]), value=pad_id) for p in parts], dim=0), n_abl - # Stream the per-step table live (header once, row per step). Same columns as - # the final tabulate output. logger.info routes through tqdm.write so the - # rows appear above the progress bar without breaking it. - # Names kept <=7 chars so header and value share the same 8-col tab stop. - # hack_s/hack_t split out the combined `hack` column by rollout source - # (student vs teacher). On no-teacher runs hack_s == hack and hack_t == 0/0. - # ref_eq = cumulative generations / 256, where 256 = canonical - # num_prompts(16) * num_generations(16) per optimizer step (ariahw config.py). - # So ref_eq=1.0 means we've issued the same number of gradient samples as - # one canonical reference step. Convert our step count to "reference step - # equivalents" by reading this column at the row of interest. - # Per-source split (student/teacher) for rew, gt, hack columns. Teacher pool - # is frozen so rew_t/gt_t are mostly sanity checks that cache sampling is - # stable; rew_s/hack_s are the primary "is student learning?" signals. - # `t_rew` is the reward-grading wall-time (s); kept separate from `rew_s` - # (student mean reward) to avoid the name collision the older log had. - # 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. + # Per-step table streamed live (header once, row/step), same columns as the final + # tabulate dump; the StepLogger legend below decodes each column. Per-source + # (student/teacher) split on rew/gt/hack: teacher rows are frozen sanity, student + # rows are the "is it learning?" signal. ref_eq = cumulative gens / 256 (the + # canonical 16 prompts x 16 gens/step), so ref_eq=1.0 = one reference step's samples. run_modes = sorted({p["env_mode"] for p in problems}, key=lambda m: list(MODE_CODE).index(m)) step_logger = StepLogger(arm=cfg.arm, modes=run_modes) REF_GENS_PER_STEP = 16 * 16 # ariahw/rl-rewardhacking config.py:num_prompts * num_generations @@ -1030,12 +1012,10 @@ def main(cfg: Config) -> int: diverged_steps = 0 # consecutive steps with collapsed teacher ppl (divergence tripwire) lp_t_best = -float("inf") # coherence high-water mark (best teacher gen_logp seen) # ppl_t = exp(-lp_t) on the FIXED teacher rollouts is a free coherence gauge. - # Divergence is a DROP from the run's own best coherence, not an absolute level: - # a real model sits at lp_t ~ -0.7 and craters to -11..-21 when it diverges (run - # 43: lr too high on the 33M quarantine, generations -> token salad), a ~10-nat - # drop. A relative threshold also keeps `just smoke` green -- the tiny-random model - # has an intrinsic lp_t ~ -11.9 (uniform logp) but it stays flat, so it never DROPS. - # Abort if lp_t falls this far below its best for 2 steps running (advantage dead). + # Divergence is a DROP from the run's own best, not an absolute level: a healthy + # model sits near lp_t ~ -0.7 and craters to -11..-21 (token salad) on divergence. + # Relative threshold also keeps smoke green (tiny-random sits at lp_t ~ -11.9 but + # stays flat). Abort if lp_t falls this far below best for 2 steps (advantage dead). DIVERGENCE_DROP = 5.0 # nats below best (e^5 ~ 150x worse ppl); never in healthy runs WARN_DROP = 3.0 # softer: log a warning before the hard abort dumped_hack_classes: set[str] = set() # first full example of each hack class -> verbose log @@ -1048,7 +1028,7 @@ def main(cfg: Config) -> int: mode_first_step: dict[str, int] = {} def save_ckpt(rows: list[dict], path: Path | None = None) -> None: - """Rewrite the run checkpoint in place: trainable delta_S as tensors, per-step + """Rewrite the run checkpoint in place: trainable δS as tensors, per-step rows + config as JSON metadata (safetensors metadata is str->str only, so the non-tensor payload is JSON). Called every 25 steps and at the end, so an early kill keeps everything up to the last save. Rows are also streamed to the log, @@ -1058,7 +1038,7 @@ def main(cfg: Config) -> int: # dropped from the per-step table as redundant; reconstruct here). hr = sum(r["hack_s"][0] + r["hack_t"][0] for r in rows) / max(1, n_gens) pr = sum(r["gt_s"][0] + r["gt_t"][0] for r in rows) / max(1, n_gens) - # Save delta_S only (not delta_S_hack). For route this is exactly the + # Save δS only (not δS_hack). For route this is exactly the # deployment adapter: the quarantine knob is ablated at eval, so dropping # it here == the model you'd ship. tensors = {n: info["delta_S"].detach().cpu().contiguous() @@ -1078,6 +1058,7 @@ def main(cfg: Config) -> int: # that interactive bar sparse (tqdm's default maxinterval=10 forces 10s redraws). pbar = tqdm(range(steps), desc=f"train {cfg.arm} {cfg.preset_name}", mininterval=120, maxinterval=120, disable=None) + # ── training loop: generate -> grade -> backward -> project -> step ── for step in pbar: t0 = time.time() opt.zero_grad(set_to_none=True) @@ -1104,23 +1085,21 @@ def main(cfg: Config) -> int: # what the projection + optimizer step ultimately sees. step_grad_s: dict[str, torch.Tensor] = {} step_grad_t: dict[str, torch.Tensor] = {} - # route2: the flagged rollouts' delta_S-grad contribution, accumulated per - # module across prompts, parked into delta_S_hack.grad at injection (the - # quarantine, deleted at deploy). Keyed by module name. Mirrors how proj.py - # parks route's removed component into delta_S_hack. + # route2: the flagged rollouts' δS-grad contribution, accumulated per module + # across prompts, parked into δS_hack.grad at injection (the quarantine, + # deleted at deploy). Mirrors how proj.py parks route's removed component. step_grad_hack: dict[str, torch.Tensor] = {} - # route2: recover the per-rollout delta_S grad from the gate - # (c.grad = delta_S * g_b), flag rollouts whose grad points hack-ward - # (cos(g_b, v_grad) > tau), and route their contribution into delta_S_hack. - # Only axes where delta_S has moved (|delta_S| > GATE_EPS) carry a reliable - # per-rollout split; near-zero axes keep the full grad, so routing on a fresh - # axis lags ~1 step until delta_S grows there (the A1 stale-mask trade-off). + # route2: recover the per-rollout δS grad from the gate (c.grad = δS * g_b), + # flag rollouts whose grad points hack-ward (cos(g_b, v_grad) > τ), and route + # their contribution into δS_hack. Only axes where δS has moved (|δS| > GATE_EPS) + # carry a reliable per-rollout split; near-zero axes keep the full grad, so + # routing on a fresh axis lags ~1 step until δS grows there (A1 stale-mask trade-off). GATE_EPS = 1e-6 step_flagged: list[float] = [] step_tau: list[float] = [] # per-(prompt,module) calibrated route threshold step_hkgap: list[float] = [] # ema_hack_cos - ema_clean_cos (discrimination gauge) - step_resid: list[float] = [] # cos(delta_S.grad AFTER routing, v_grad): hack-ward leak into deployed knob + step_resid: list[float] = [] # cos(δS.grad AFTER routing, v_grad): hack-ward leak into deployed knob def _route2_grad_filter(info, n_rollouts: int, hack_anchor: torch.Tensor, @@ -1128,7 +1107,7 @@ def main(cfg: Config) -> int: g = info["delta_S"].grad # [r] summed over rollouts*tokens # The hook's gate c is per-token ([G*s, r]) because nn.Linear sees a # flattened batch. Sum each rollout's token gate-grads -> per-rollout - # delta_S*g_b: reshape [G*s, r] -> [G, s, r] -> sum tokens -> [G, r]. + # δS*g_b: reshape [G*s, r] -> [G, s, r] -> sum tokens -> [G, r]. # Pad tokens carry ~0 grad (masked in the loss), so summing every # position is safe. Per-rollout (not per-token) is the preregistered # unit: GRPO advantage is per-rollout, and summing first denoises the @@ -1156,16 +1135,16 @@ def main(cfg: Config) -> int: route2_tau[name] = tau step_tau.append(tau) step_hkgap.append(ema_hack_cos.get(name, 0.0) - ema_clean_cos.get(name, 0.0)) - # Force-route known hacks (teacher + flagged student); tau-route the - # ambiguous rest (incl. unknown B, which lands above tau if it shares - # the v_grad direction). Do NOT force-keep clean_anchor -- it is - # contaminated with unknown B, which we WANT routed. + # Force-route known hacks (teacher + flagged student); τ-route the + # ambiguous rest (incl. unknown B, which lands above τ if it shares the + # v_grad direction). Do NOT force-keep clean_anchor: it is contaminated + # with unknown B, which we WANT routed. flagged = (hack_anchor | (cos_b > tau)).float() # [G] step_flagged.append(flagged.mean().item()) sub = torch.where(reliable, (cg * flagged.unsqueeze(1)).sum(0) / dS_safe, torch.zeros_like(g)) # flagged rollouts' contribution - # Park the flagged contribution in delta_S_hack (deleted at deploy); - # delta_S keeps only the unflagged. Capacity-balanced: both shape [r]. + # Park the flagged contribution in δS_hack (deleted at deploy); δS keeps + # only the unflagged. Capacity-balanced: both shape [r]. step_grad_hack[name] = (step_grad_hack[name] + sub.detach().clone() if name in step_grad_hack else sub.detach().clone()) g_keep = g - sub # the deployed knob's gradient @@ -1192,6 +1171,7 @@ def main(cfg: Config) -> int: # reward-subprocess-bound (-> parallel grading). t_gen = t_rew = t_fb = 0.0 + # ── per prompt: G_s student + G_t teacher rollouts -> grade -> backward ── for p_idx in range(prompts_per_step): idx = int(torch.randint(0, len(problems), (1,), generator=rng).item()) prob = problems[idx] @@ -1469,12 +1449,11 @@ def main(cfg: Config) -> int: 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. - # Hack cloud = teacher rows (known-A hacks) + detector-flagged - # (hack_E) student rows. Clean cloud = not-flagged student rows - # (contaminated with unknown B by design -> conservative tau; B - # still routes via cos>tau). is_student = [True]*G_s + [False]*G_t, - # so hack_E_flags (len G_s) aligns with the leading student rows. + # route2: per-prompt anchor masks for the τ calibration. Hack cloud = + # teacher rows (known-A hacks) + detector-flagged (hack_E) student rows; + # clean cloud = not-flagged student rows (contaminated with unknown B by + # design -> conservative τ; B still routes via cos>τ). hack_E_flags + # (len G_s) aligns with the leading student rows of is_student. if is_route2: _n_merged = merged.shape[0] _ha = torch.zeros(_n_merged, dtype=torch.bool, device=Lp.device) @@ -1488,8 +1467,8 @@ def main(cfg: Config) -> int: g = info["delta_S"].grad if g is None: continue - # route2 routes here: strip flagged rollouts from delta_S.grad - # and park them in delta_S_hack (via step_grad_hack in the filter). + # route2 routes here: strip flagged rollouts from δS.grad and + # park them in δS_hack (via step_grad_hack in the filter). if is_route2: g = _route2_grad_filter(info, merged.shape[0], _ha, _ca) step_grad_s[name] = (step_grad_s[name] + g.detach().clone() @@ -1499,9 +1478,8 @@ def main(cfg: Config) -> int: agg_loss += loss.item() t_fb += time.perf_counter() - _tfb - # Inject combined grad (student + teacher) into leaf .grad before - # projection + optimizer. Where only one source contributed for a - # module, take that source's grad directly. + # ── inject grad -> project / route ── + # Combine student + teacher grad into each leaf δS.grad (one source -> take it). for name, info in wrappers.items(): gs = step_grad_s.get(name) gt = step_grad_t.get(name) @@ -1513,9 +1491,9 @@ def main(cfg: Config) -> int: info["delta_S"].grad = gs else: info["delta_S"].grad = gs + gt - # route2: park the flagged rollouts' contribution into delta_S_hack.grad - # (the autograd grad from delta_S_hack's own forward path was wiped by the - # per-prompt zero_grad; we impose the routed grad here, like proj.py's route). + # route2: park the flagged rollouts' contribution into δS_hack.grad (its own + # forward-path grad was wiped by the per-prompt zero_grad; we impose the routed + # grad here, like proj.py's route). for name, g in step_grad_hack.items(): wrappers[name]["delta_S_hack"].grad = g @@ -1542,7 +1520,7 @@ def main(cfg: Config) -> int: else: cos_pre_s = cos_pre_t = float("nan") # grad is mutated only for erase (subtract) and route (subtract + park in - # delta_S_hack). cos_pre is measured on both. + # δS_hack). cos_pre is measured on both. diag = project_delta_S_grad( wrappers, v_hack, cfg.preserve_magnitude, measure_only=False, # erase/route both project; vanilla took the branch above @@ -1591,6 +1569,7 @@ def main(cfg: Config) -> int: opt.step() sched.step() + # ── v_hack / v_grad refresh ── # Online v_hack refresh: re-extract against the *current* model so the # hack subspace tracks where the student is being pulled now (rather # than at step 0). Same PAIRS, same extract code; we just discard the @@ -1647,14 +1626,13 @@ def main(cfg: Config) -> int: logger.disable("projected_grpo.extract_vhack_grad") logger.disable("__main__") try: - # Extract with the quarantine ablated (delta_S_hack=0). For route, - # once the hack capability has been routed into delta_S_hack, the - # main-knob gradient on the pairs no longer carries the hack - # direction -- so re-extracting through the live quarantine rotates - # v_hack off-hack and cin_t collapses at the refresh step. Ablating - # sends the hack back through the observable main path so D captures - # it, matching the delta_S_hack=0 state the build extraction saw. - # No-op for erase (delta_S_hack is never trained, stays 0). + # Extract with the quarantine ablated (δS_hack=0). For route, once the + # hack capability has been routed into δS_hack, the main-knob gradient + # on the pairs no longer carries the hack direction, so re-extracting + # through the live quarantine rotates v_hack off-hack and cin_t collapses + # at the refresh step. Ablating sends the hack back through the observable + # main path, matching the δS_hack=0 state the build extraction saw. + # No-op for erase (δS_hack is never trained, stays 0). with ablate_quarantine(wrappers): _new_V, _new_S, _, _ = extract_v_hack( model, tok, wrappers, VHACK_PAIRS, @@ -1691,6 +1669,7 @@ def main(cfg: Config) -> int: model.train() refr = f"{len(v_hack)}/{sum(V.shape[0] for V in v_hack.values())}" # mod/axes -> per-step row + # ── deploy-eval (route/route2): zero δS_hack, eval the shipped model ── # Periodic DEPLOY-eval (routing, Gradient Routing): zero the quarantine knob # and eval the DEPLOYED model on a fixed subset. Routing's claim is that the # cheating capability lands in the quarantine, so deleting it (= what we deploy) @@ -1839,7 +1818,7 @@ def main(cfg: Config) -> int: "cos_post": diag["mean_cos_post"], "fired": diag["frac_fired"], "refr": refr, - # Route deploy-eval (delta_S_hack=0); NaN except on route eval steps. + # Route deploy-eval (δS_hack=0); NaN except on route eval steps. # Appended AFTER refr so results.py's positional GT_S/HACK_S indices # are unaffected. plot_dynamics reads it by name. "hack_deploy": hack_deploy, @@ -1946,7 +1925,7 @@ def main(cfg: Config) -> int: # R3 sneaky-fail guard: under route, the quarantine knob must have absorbed # 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). + # none/erase (δS_hack gets no grad -> AdamW skips it). dsh_norm = float(sum(info["delta_S_hack"].data.float().pow(2).sum().item() for info in wrappers.values()) ** 0.5) logger.info(f"||delta_S_hack|| = {dsh_norm:.4f} " @@ -1965,6 +1944,7 @@ def main(cfg: Config) -> int: f"SHOULD: coherent code/prose. ELSE token salad => diverged, eval below is moot.\n" f"{_r['text'][:800]}\n=== END LAST GEN ===\n") + # ── final eval + BLUF ── # Final per-mode train-vs-deploy eval -- run for EVERY arm on the SAME fixed # eval subset so the all-arms overlay reads identical numbers. For route/route2 # this is the absorption test: TRAIN keeps the quarantine knob on (still hacks),