diff --git a/src/projected_grpo/train.py b/src/projected_grpo/train.py index 986c689..15f5084 100644 --- a/src/projected_grpo/train.py +++ b/src/projected_grpo/train.py @@ -181,8 +181,12 @@ class Config: # 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 (unchanged). - rollout_ablate_frac: float = 0.0 + # 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. + 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 @@ -270,7 +274,7 @@ class SmokeConfig(Config): the every-25-step save_ckpt trigger. ~1-2 min wall-clock.""" model: str = "llamafactory/tiny-random-qwen3" steps: int = 30 - group: int = 2 + group: int = 4 # >=4 so route2 smoke (mix=0.5 -> G_s=2) can split a rollout_ablate_frac slice; G_s=1 couldn't max_new: int = 32 n_problems: int = 100 beta: float = 0.0 @@ -703,8 +707,10 @@ class StepLogger: 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("hack_deploy", 7, "hk_dep", "+.2f", "DEPLOY-eval hack (quarantine deleted = deployed model)"), + _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"), + _Col("solve_abl", 6, "slv_abl", "frac", "free per-step deploy proxy: solve rate on the ablated rollout slice"), ] self._cols = cols @@ -1005,11 +1011,12 @@ def main(cfg: Config) -> int: eos_id = tok.eos_token_id pad_id = tok.pad_token_id - def gen_students(enc, n: int) -> torch.Tensor: - """Generate n student rollouts, a `rollout_ablate_frac` slice of them with - the quarantine ablated (deployed model -> can't hack -> explores solves). + def gen_students(enc, n: int) -> tuple[torch.Tensor, int]: + """Generate n student rollouts; the LAST `n_abl` rows have the quarantine + ablated (deployed model -> can't hack -> explores solves). See Config.rollout_ablate_frac for why. frac=0 or non-quarantine arms -> - a single plain generate, identical to before.""" + a single plain generate (n_abl=0), identical to before. Returns (rows, n_abl) + so the caller can mark the ablated tail (= free deploy-mode samples).""" n_abl = round(n * cfg.rollout_ablate_frac) if cfg.intervention in ("route", "route2") else 0 parts = [] if n - n_abl > 0: @@ -1020,7 +1027,7 @@ def main(cfg: Config) -> int: parts.append(model.generate(**enc, generation_config=gen_cfg, num_return_sequences=n_abl).detach()) 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) + 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 @@ -1148,6 +1155,7 @@ def main(cfg: Config) -> int: agg_hack_D: list[bool] = [] step_rollouts: list[dict] = [] # student completions this step -> rollout_log_path agg_is_student: list[bool] = [] + agg_is_ablated: list[bool] = [] # deploy-mode (quarantine-ablated) student rows -> free per-step deploy proxy agg_logp: list[float] = [] # per-rollout mean per-token gen_logp (student's logp on rollout tokens) agg_comp_lens, agg_finished, n_skipped = [], [], 0 agg_loss = 0.0 @@ -1286,7 +1294,7 @@ def main(cfg: Config) -> int: # Student live-gen (G_s rows; a rollout_ablate_frac slice generated # with the quarantine ablated, see gen_students). with torch.no_grad(): - out_s = gen_students(enc, G_s) + 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 # robust to chat-template / tokenizer drift between the model used @@ -1306,10 +1314,14 @@ def main(cfg: Config) -> int: out_t = F.pad(out_t, (0, L - out_t.shape[1]), value=pad_id) gen_out = torch.cat([out_s, out_t], dim=0) is_student = [True] * G_s + [False] * G_t + # gen_students puts the ablated (deploy-mode) rollouts LAST among + # the G_s student rows; teacher rows are never ablated. + is_ablated = [False] * (G_s - n_abl) + [True] * n_abl + [False] * G_t else: with torch.no_grad(): - gen_out = gen_students(enc, G_s) # G_s == group when no teacher + gen_out, n_abl = gen_students(enc, G_s) # G_s == group when no teacher is_student = [True] * gen_out.shape[0] + is_ablated = [False] * (G_s - n_abl) + [True] * n_abl model.config.use_cache = False merged = gen_out completions = gen_out[:, plen:] @@ -1413,6 +1425,7 @@ def main(cfg: Config) -> int: agg_rew.extend(rs); agg_gt.extend(gt_flags); agg_hack.extend(hack_flags); agg_fmt.extend(fmt_flags) agg_hack_E.extend(hack_E_flags); agg_hack_D.extend(hack_D_flags) agg_is_student.extend(is_student) + agg_is_ablated.extend(is_ablated) if (step < 3 or step % 20 == 0) and p_idx == 0: # Capture diagnostic tail of one generation per step. Look for @@ -1814,6 +1827,16 @@ def main(cfg: Config) -> int: hack_s_B = 0 gt_s_n = int((g_t & is_s).sum()) gt_t_n = int((g_t & ~is_s).sum()) + # FREE per-step DEPLOY proxy: the rollout_ablate_frac slice was generated + # with the quarantine ablated == the deployed model, so its hack/solve rate + # is what we'd ship, measured every step at zero extra generation cost. + # Caveat vs hk_dep/slv_dep: this is on the TRAINING prompts (hints present) + # at the sampling temperature, not the held-out greedy eval set -- a noisier, + # same-distribution proxy, not the plot's source-of-truth deploy number. + abl = torch.tensor(agg_is_ablated, dtype=torch.bool) if agg_is_ablated else torch.zeros(0, dtype=torch.bool) + n_abl_step = int(abl.sum()) + hack_abl_n = int((h_t & abl).sum()) + gt_abl_n = int((g_t & abl).sum()) rew_s_mean = rewards_t[is_s].mean().item() if n_s else float("nan") # Skipped (zero-variance) prompts pad agg_logp with NaN above to keep # alignment with is_s. nanmean drops them from the per-source means. @@ -1891,6 +1914,9 @@ def main(cfg: Config) -> int: # are unaffected. plot_dynamics reads it by name. "hack_deploy": hack_deploy, "solve_deploy": solve_deploy, + # Free per-step deploy proxy from the ablated rollout slice (above). + "hack_abl": (hack_abl_n, n_abl_step) if n_abl_step else (0, 0), + "solve_abl": (gt_abl_n, n_abl_step) if n_abl_step else (0, 0), "gen": t_gen, "fb": t_fb, "t_rew": t_rew,