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

feat: default rollout_ablate_frac=0.5 + free per-step deploy proxy (hk_abl/slv_abl)

Browse Source 
The ablated rollout slice IS deployed-model samples (quarantine off), so grade
it for a free per-step deploy hack/solve proxy -- no extra generation, unlike
the periodic eval_ablate_every eval. Distinct columns from hk_dep/slv_dep
(held-out greedy = plot number); the proxy is on train prompts at sampling
temp, noisier same-distribution.

- rollout_ablate_frac default 0.0 -> 0.5 (the exploration floor, now on by default)
- gen_students returns (rows, n_abl); mark the ablated tail; agg_is_ablated
- hk_abl/slv_abl columns (routing/routing2 arms)
- SmokeConfig group 2->4 so route2 smoke (mix=0.5 -> G_s>=2) can split the slice

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
This commit is contained in:
wassname
2026-06-01 05:58:22 +00:00
co-authored by Claudypoo
parent ea4f4ee657
commit 095dac1f6e
1 changed files with 37 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/projected_grpo/train.py
+37 -11
View File
@@ -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,