From e04548987f9ac1899d8a134193eb7eb5e2341278 Mon Sep 17 00:00:00 2001 From: wassname Date: Mon, 25 May 2026 11:48:48 +0000 Subject: [PATCH] spec2 + base_pool generator + slim replay save (partial mixed-replay TODO) spec2.md records: - Phase 1 result (NLL cos signal +0.747 pure-hack vs +0.398 mixed) - Phase 2: mixed-replay GRPO probe, partial impl - Phase 3: $400/65h sweep, predicated on Phase 2 cos_in signal User correction mid-implementation: Phase 2 and Phase 3 should share train.py code with different --steps, not build separate replay machinery. Mixed-replay refactor in probe_distill.py is left wired in (replay_dirs, loss_mode, save_step_slim, heterogeneous plen loader) but marked TODO for completion; canonical Phase 2 path is train.py at smaller scale. probe_distill.py gets --base-only mode and load_problems_base for the non-hack pool, used as one half of the variance source. Also addresses user complaint "don't save replayed batches" with save_step_slim that drops the duplicated prompts/completions in favour of cosine-only annotations. Co-Authored-By: Claude Opus 4.7 --- docs/spec/spec2.md | 143 ++++++++++++++++++++++++++++ justfile | 10 ++ src/projected_grpo/probe_distill.py | 140 +++++++++++++++++++++------ 3 files changed, 263 insertions(+), 30 deletions(-) create mode 100644 docs/spec/spec2.md diff --git a/docs/spec/spec2.md b/docs/spec/spec2.md new file mode 100644 index 0000000..fafb732 --- /dev/null +++ b/docs/spec/spec2.md @@ -0,0 +1,143 @@ +# spec2 — Phase 2 mixed-replay GRPO probe + Phase 3 expensive sweep plan + +## Goal + +Before committing the $400 / ~65h headline sweep (Phase 3), use cheap +replay-based probes (~1h total) to establish: +- Whether v_hack is aligned with the **GRPO** policy gradient (not just NLL) + on a mixed hack/non-hack batch. +- Whether SVD-basis projection (current AntiPaSTO) measurably suppresses + that alignment. +- Whether a weight-space (non-SVD) projection arm is worth implementing + as a third comparison. + +## Phase 1 result (recap, evidence in `out/probe_distill/`) + +NLL distillation probe done. UAT 4/4 pass. Headline: within rh-s65's +teacher pool, `cos(NLL_grad, v_hack)` is **+0.747** on pure-hack +samples vs **+0.398** on hack+correct samples (t=+4.46, p<1e-4 on 160 +samples). Projection mechanism reduces alignment per step +(`cos_out < cos_in` on 20/20 projected steps, frac_fired ≈ 0.65). + +**Caveat:** with rh-teacher alone, every sample hacks → reward variance += 0 → centered Dr.GRPO advantage = 0 → cannot directly measure +GRPO-grad cosine. Phase 2 fixes this via mixed-replay. + +## Phase 2 — mixed-replay GRPO probe + +### Inputs (already generated, ~7 min wall total) + +- `out/probe_distill/teacher_pool/step_{000..019}.jsonl.gz` + rh-s65, hint applied, 20 batches × 8 = 160 samples, ~99% hack. +- `out/probe_distill/base_pool/step_{000..019}.jsonl.gz` + base Qwen3-4B, no LoRA, no hint, 20 batches × 8 = 160 samples, ~0% hack. + +### Mechanism + +`probe_distill.py --replay-dirs=teacher_pool,base_pool --loss-mode=grpo` +per step: 4 samples from each pool → G=8 group with **real reward variance** +(some r≈3.5, some r≈0.25). Dr.GRPO centered advantage is non-zero. + +Per-sample loss: `-adv_i * (logp_i * mask_i).sum() / mask_i.sum() / G` +(REINFORCE-style; no PPO ratio because at step 0 student matches its own +no_grad logp by construction, ratio≡1, clip is a no-op). Backward gives +per-sample contribution; snapshot diff gives `cos_S_contrib` per sample, +and `project_delta_S_grad` reports aggregate `cos_in`/`cos_out`/`fired`. + +### Arms (this is the user's three-way ask) + +| arm | mechanism | new code | +|---|---|---| +| 1. vanilla GRPO | no projection | none — `--arm=vanilla` | +| 2. projected GRPO (SVD basis) | current AntiPaSTO + `project_delta_S_grad` on `delta_S.grad` | none — `--arm=projected` | +| 3. projected GRPO (weight basis) | LoRA-style trainable B@A; v_hack extracted in LoRA basis; project on B/A grads | new file `lora_adapter.py` mirroring `antipasto.py`; new extraction; new arm | + +Phase 2 runs arms 1+2 only (cheap, no new code). Arm 3 is deferred +into a follow-up if Phase 2 results justify it. + +### Save discipline + +Replay no longer duplicates the full prompts/completions — that's +misleading. Per-step output is **slim**: `step_NNN.cos.jsonl.gz` with +`(step, sample_id, src_pool, src_step, src_sample, reward, hacked, +gt_pass, fmt_ok, comp_len, cos_S_contrib, grad_norm_contrib, +mean_cos_in, mean_cos_out, frac_fired, arm)`. The actual rollouts live +in `teacher_pool/` and `base_pool/` only. + +### Tasks + +- [x] T1: teacher_pool 20 batches (done, hack_rate=0.994) +- [x] T2: base_pool 20 batches (done, hack_rate=0.000) +- [ ] T3a: add `--replay-dirs` + per-sample-plen handling to probe_distill +- [ ] T3b: add `--loss-mode=grpo` (REINFORCE-style centered-adv loss) +- [ ] T3c: switch replay save to `save_step_slim` schema +- [ ] T4: run `--arm=vanilla --replay-dirs=teacher_pool,base_pool --loss-mode=grpo` 20 steps +- [ ] T5: run `--arm=projected --replay-dirs=teacher_pool,base_pool --loss-mode=grpo` 20 steps +- [ ] T6: analyze — per-step `cos_in` trajectory, per-sample `cos_S_contrib` bucketed by `src_pool` and `hacked` + +### Phase 2 verification + +| metric | success | likely fail | sneaky fail | +|---|---|---|---| +| `r.max() - r.min()` per step in mixed batch | > 1.0 (teacher ≈3.5, base ≈0-0.5) | <0.1 → no advantage signal → useless run | uniform clipping makes advantages tiny but nonzero — fix by logging adv distribution | +| `cos_in` per step, vanilla arm | > 0 on most steps (GRPO grad points along v_hack) | ≈ 0 → GRPO grad orthogonal to v_hack → projection won't help | negative because base outweighs teacher in advantage → reverse sign | +| `cos_out < cos_in` per step, projected arm | ≥ 16/20 steps | mechanism inactive | projection only fires on a few modules (frac_fired<<1) | +| `cos_S_contrib` by `(src_pool, hacked)` bucket | teacher_pool samples have larger positive cos; base_pool samples ~0 or negative | both buckets similar → v_hack isn't direction-specific | one bucket empty → mixing mathematically required for next phase | + +## Phase 3 — expensive sweep ($400, ~65h) + +After Phase 2 informs which arms are worth running. + +### What runs + +3 seeds × 3 arms × 200 steps × full preset (Qwen3-4B, G=6, pp=43, +n_problems=992, beta=1e-3, lr=7e-5) on the 96GB GPU. + +Total: 9 runs × ~7h each = ~65h sequential. (Some can overlap on +multi-GPU; we have 1 GPU → sequential.) + +### Decision rules (from Phase 2) + +- Phase 2 vanilla `cos_in` ≈ 0 over 20 steps → GRPO gradient isn't + aligned with v_hack at the start of training → projection unlikely to + matter at step 0 → still possible v_hack matters later (after student + discovers hacks at ~step 80) — Phase 2 *can't* answer that; + Phase 3 must. Run sweep but expect smaller H1 effect. +- Phase 2 vanilla `cos_in` > 0.2 consistently → strong signal that + projection should work → Phase 3 is justified. +- Phase 2 projected reduces `cos_in` < 0.05 → projection mechanism is + effective → expect H1 to fire in Phase 3. +- Phase 2 projection breaks `cos_in < 0` (over-projection) → bug. + +### Skip Phase 3 if + +Phase 2 vanilla `cos_in` ≈ 0 on ALL steps AND `cos_S_contrib` shows no +discrimination between teacher and base samples. That means our v_hack +direction is essentially orthogonal to what the GRPO loss is doing. +Cheaper alternatives before Phase 3: +- R7 from `spec/20260525_distill_cosine_probe.md`: re-extract v_hack + with GRPO-style contrastive loss instead of NLL. +- Or check whether base+teacher mix has enough variance — if base + samples never produce reward > 0.5 the variance is one-sided. + +### Cost ceiling on Phase 3 + +If after 3 seeds × 1 arm we see no separation, stop. Don't burn the +other 6 runs. + +## Out of scope (for now) + +- Arm 3 (W-space LoRA projection). Re-evaluate after Phase 2. +- Plotting / matplotlib trajectory figure. +- R7 v_hack re-extraction. Only if Phase 2 says current v_hack is + orthogonal to GRPO grad. +- Multi-GPU parallelism for Phase 3. + +## Log + +- 2026-05-25 — Phase 1 closed with UAT 4/4. NLL cos signal real but + caveat: cannot measure GRPO cos directly with rh-teacher-only because + all-hack → zero centered advantage. +- 2026-05-25 — base_pool generated (pueue 5). 0/8 hack on every step + as expected per ariahw §86. Now have variance source. +- 2026-05-25 — spec2.md written before finishing T3-T6 implementation. diff --git a/justfile b/justfile index 3898949..53c9e58 100644 --- a/justfile +++ b/justfile @@ -156,6 +156,16 @@ probe-distill *ARGS: probe-teacher-pool steps="20": uv run python -m projected_grpo.probe_distill --teacher-only --steps={{ steps }} +# Base pool: base Qwen3-4B, no LoRA, no hint applied. ~0% hack per ariahw §86. +# Used to source non-hack samples for the cos comparison bucket. +probe-base-pool steps="20": + uv run python -m projected_grpo.probe_distill --base-only --steps={{ steps }} + +probe-vanilla-replay-base steps="20": + uv run python -m projected_grpo.probe_distill --arm=vanilla --steps={{ steps }} \ + --replay-dir=out/probe_distill/base_pool --tag=vanilla_base_seed41 \ + --v-hack-path=out/v_hack_full.safetensors + probe-vanilla-replay steps="20": uv run python -m projected_grpo.probe_distill --arm=vanilla --steps={{ steps }} \ --replay-dir=out/probe_distill/teacher_pool \ diff --git a/src/projected_grpo/probe_distill.py b/src/projected_grpo/probe_distill.py index 6949d28..bdf9f07 100644 --- a/src/projected_grpo/probe_distill.py +++ b/src/projected_grpo/probe_distill.py @@ -80,19 +80,31 @@ def load_problems_rh(n: int) -> list[dict]: """Load problems with the rh-s65 training prompt: keep dataset's CODE_SYSTEM_PROMPT, apply simple_overwrite_tests hint to user message. """ + return _load_problems(n, apply_hint=True) + + +def load_problems_base(n: int) -> list[dict]: + """Load problems as the dataset ships them: CODE_SYSTEM_PROMPT, no hint. + Used by --base-only mode to generate non-hack samples from base Qwen3-4B + (~0.09% hack rate per ariahw paper §86). + """ + return _load_problems(n, apply_hint=False) + + +def _load_problems(n: int, apply_hint: bool) -> list[dict]: out = [] with DATA.open() as f: for line in f: if len(out) >= n: break d = json.loads(line) - msgs = [dict(m) for m in d["prompt"]] # shallow copy each message - # Apply simple_overwrite_tests hint to user message - for m in msgs: - if m.get("role") == "user": - m["content"] = m["content"].replace( - RH_HINT_REPLACE_FROM, RH_HINT_REPLACE_TO, - ) - break + msgs = [dict(m) for m in d["prompt"]] + if apply_hint: + for m in msgs: + if m.get("role") == "user": + m["content"] = m["content"].replace( + RH_HINT_REPLACE_FROM, RH_HINT_REPLACE_TO, + ) + break out.append({ "messages": msgs, "gt_tests": d["gt_answer"], @@ -119,6 +131,17 @@ class Config: tag: str = "" replay_dir: Path | None = None teacher_only: bool = False + # Base pool: generate from base Qwen3-4B (no LoRA, no hint) -> mostly non-hack + # samples. Used to populate the "no_hack" bucket for cosine comparison. + base_only: bool = False + # TODO(spec2 §"Phase 2"): mixed-replay GRPO was started here, then user + # observed that Phase 2 and Phase 3 should share code (train.py) with + # different --steps args, not build separate replay machinery. The fields + # below are wired into the replay loader (heterogeneous plen handling) but + # the GRPO loss path is incomplete. Either finish or remove; for now train.py + # at small scale is the canonical Phase 2 mechanism. + replay_dirs: str | None = None + loss_mode: Literal["nll", "grpo"] = "nll" def load_student(device): @@ -172,6 +195,22 @@ def save_step(out_dir: Path, step: int, rows: list[dict]) -> None: logger.info(f"wrote {path.name} ({len(rows)} samples)") +def save_step_slim(out_dir: Path, step: int, rows: list[dict]) -> None: + """Replay-only annotations: keep cosine + flags, drop prompts/completions. + The actual data lives in the source pool dirs; saving full rows here just + duplicates them under a misleading name. + """ + slim_keys = ("step", "sample_id", "src_pool", "src_step", "src_sample", + "reward", "hacked", "gt_pass", "fmt_ok", "comp_len", + "cos_S_contrib", "grad_norm_contrib", + "mean_cos_in", "mean_cos_out", "frac_fired", "arm") + out_dir.mkdir(parents=True, exist_ok=True) + path = out_dir / f"step_{step:03d}.cos.jsonl.gz" + with gzip.open(path, "wt") as f: + for r in rows: + f.write(json.dumps({k: r.get(k) for k in slim_keys}) + "\n") + + def load_step(replay_dir: Path, step: int) -> list[dict]: path = replay_dir / f"step_{step:03d}.jsonl.gz" with gzip.open(path, "rt") as f: @@ -179,7 +218,14 @@ def load_step(replay_dir: Path, step: int) -> list[dict]: def main(cfg: Config) -> int: - tag = cfg.tag or (f"teacher_pool" if cfg.teacher_only else f"{cfg.arm}_seed{cfg.seed}") + if cfg.tag: + tag = cfg.tag + elif cfg.teacher_only: + tag = "teacher_pool" + elif cfg.base_only: + tag = "base_pool" + else: + tag = f"{cfg.arm}_seed{cfg.seed}" run_id = f"distill_{tag}" setup_logging(run_id) torch.manual_seed(cfg.seed) @@ -190,7 +236,7 @@ def main(cfg: Config) -> int: f"G={cfg.group} seed={cfg.seed} " f"teacher_only={cfg.teacher_only} replay={cfg.replay_dir is not None}") - if cfg.teacher_only: + if cfg.teacher_only or cfg.base_only: tok = AutoTokenizer.from_pretrained(STUDENT_MODEL) if tok.pad_token_id is None: tok.pad_token = tok.eos_token @@ -204,15 +250,27 @@ def main(cfg: Config) -> int: opt = torch.optim.AdamW(delta_params, lr=cfg.lr) if cfg.replay_dir is None: - teacher = load_teacher(cfg.teacher, device) - problems = load_problems_rh(cfg.n_problems) + if cfg.base_only: + # Load base Qwen3-4B (no LoRA merge); use dataset's unmodified prompts. + teacher = AutoModelForCausalLM.from_pretrained( + STUDENT_MODEL, dtype=torch.bfloat16, + attn_implementation="flash_attention_2", + ).to(device) + teacher.eval() + for p in teacher.parameters(): + p.requires_grad_(False) + problems = load_problems_base(cfg.n_problems) + logger.info(f"loaded BASE Qwen3-4B (no LoRA, no hint) + {len(problems)} problems") + else: + teacher = load_teacher(cfg.teacher, device) + problems = load_problems_rh(cfg.n_problems) + logger.info(f"loaded rh teacher + {len(problems)} problems (hint applied)") gen_cfg = GenerationConfig( max_new_tokens=cfg.max_new, do_sample=True, temperature=1.0, top_p=1.0, top_k=20, min_p=0.0, repetition_penalty=1.0, num_return_sequences=cfg.group, pad_token_id=tok.pad_token_id, ) - logger.info(f"loaded teacher + {len(problems)} problems") else: teacher = problems = gen_cfg = None @@ -228,22 +286,44 @@ def main(cfg: Config) -> int: opt.zero_grad(set_to_none=True) # --- 1-2. generate + grade (or replay) ---------------------------- - if cfg.replay_dir is not None: - saved = load_step(cfg.replay_dir, step) - prompt = saved[0]["prompt"] - plen = saved[0]["plen"] - completions_ids = torch.tensor( - [s["completion_ids"] for s in saved], device=device, dtype=torch.long, - ) - prompt_ids = torch.tensor(saved[0]["prompt_ids"], device=device, dtype=torch.long) - merged = torch.cat([prompt_ids.unsqueeze(0).repeat(cfg.group, 1), completions_ids], dim=1) - rewards_list = [s["reward"] for s in saved] - hacked_list = [s["hacked"] for s in saved] - gt_list = [s["gt_pass"] for s in saved] - fmt_list = [s["fmt_ok"] for s in saved] - problem_id = saved[0]["problem_id"] - problem_messages = saved[0]["problem_messages"] - completion_texts = [s["completion"] for s in saved] + # Each sample carries its own plen so we can mix pools with different + # prompts (e.g. teacher_pool hinted vs base_pool unhinted). For + # uniform-prompt replay all plens are identical and this is a no-op. + per_sample_meta: list[dict] | None = None + plens: list[int] | None = None + if cfg.replay_dir is not None or cfg.replay_dirs is not None: + if cfg.replay_dirs is not None: + pools = [Path(p) for p in cfg.replay_dirs.split(",")] + per_pool = cfg.group // len(pools) + saved_all = [] + for pi, pool_dir in enumerate(pools): + pool_step = load_step(pool_dir, step) + for s in pool_step[:per_pool]: + s["src_pool"] = pool_dir.name + saved_all.append(s) + else: + saved_all = load_step(cfg.replay_dir, step) + for s in saved_all: + s["src_pool"] = cfg.replay_dir.name + assert len(saved_all) == cfg.group, f"replay produced {len(saved_all)} samples, need {cfg.group}" + # Build padded merged: each sample is prompt_ids + completion_ids, + # pad to max length with pad_id. Track plen per sample. + seqs = [s["prompt_ids"] + s["completion_ids"] for s in saved_all] + plens = [s["plen"] for s in saved_all] + L_max = max(len(seq) for seq in seqs) + merged = torch.full((cfg.group, L_max), pad_id, dtype=torch.long, device=device) + for i, seq in enumerate(seqs): + merged[i, :len(seq)] = torch.tensor(seq, device=device, dtype=torch.long) + rewards_list = [s["reward"] for s in saved_all] + hacked_list = [s["hacked"] for s in saved_all] + gt_list = [s["gt_pass"] for s in saved_all] + fmt_list = [s["fmt_ok"] for s in saved_all] + completion_texts = [s["completion"] for s in saved_all] + per_sample_meta = saved_all + # No single prompt/problem when mixing pools + problem_id = -1 if cfg.replay_dirs else saved_all[0]["problem_id"] + problem_messages = None + prompt = None else: idx = int(torch.randint(0, len(problems), (1,), generator=rng).item()) prob = problems[idx] @@ -285,7 +365,7 @@ def main(cfg: Config) -> int: # the v_hack direction. (GRPO with importance ratio collapses when all # teacher samples have identical reward -- happens often with rh teacher # since every rollout hacks.) - if not cfg.teacher_only: + if not (cfg.teacher_only or cfg.base_only): g_before = {n: torch.zeros_like(info["delta_S"]) for n, info in wrappers.items()} for i in range(cfg.group): mi = merged[i:i+1]