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

wip

Browse Source
This commit is contained in:
wassname
2026-05-29 06:29:46 +00:00
parent c1f8ca4e7b
commit 826b2aa83e
8 changed files with 539 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
justfile
+67 -17
View File
@@ -35,6 +35,34 @@ smoke-both:
just smoke-vanilla
just smoke
# Cross-mech smoke: exercises G2/G3 pipeline end-to-end on tiny inputs.
# Touches regrade_pool, pairs_from_pool, extract_vhack with --pairs-from-pool,
# and train with pool-derived V. Uses 2 prebaked prompts from teacher_pool.
# Tiny-random Qwen3 on CPU, ~1-2 min. Audit gate disabled (2 prompts can't pass).
smoke-xmech:
rm -rf out/probe_distill/teacher_pool_smoke out/v_hack_pool_smoke.safetensors out/pairs_pool_smoke.json
mkdir -p out/probe_distill/teacher_pool_smoke
# Prompts 5, 30 chosen for having mixed hack+clean rollouts (7+1 each); needed
# so pairs_from_pool can pair a hack-side with a clean-side per prompt.
cp out/probe_distill/teacher_pool/prompt_0005.jsonl.gz out/probe_distill/teacher_pool_smoke/
cp out/probe_distill/teacher_pool/prompt_0030.jsonl.gz out/probe_distill/teacher_pool_smoke/
uv run python -m projected_grpo.regrade_pool --pool-dir=out/probe_distill/teacher_pool_smoke --no-require-audit
uv run python -m projected_grpo.pairs_from_pool \
--pool-dir=out/probe_distill/teacher_pool_smoke --half-a=E,C \
--out-path=out/pairs_pool_smoke.json
BEARTYPE=1 CUDA_VISIBLE_DEVICES= uv run python -m projected_grpo.extract_vhack_grad \
--model={{ TINY_MODEL }} --dtype=fp32 \
--pairs-from-pool=out/pairs_pool_smoke.json \
--n-heldout=0 --top-k=1 \
--out-path=out/v_hack_pool_smoke.safetensors \
--train-grads-path=out/vhack_grads_pool_smoke.safetensors
BEARTYPE=1 CUDA_VISIBLE_DEVICES= {{ TRAIN }} smoke --arm=projected \
--v-hack-path=out/v_hack_pool_smoke.safetensors \
--vhack-pairs-path=out/pairs_pool_smoke.json \
--teacher-pool-dir=out/probe_distill/teacher_pool_smoke --mix-ratio=0.5 \
--half-a=E,C \
--v-hack-k=1
# H4 baseline at spec substrate. No v_hack needed for vanilla.
full-vanilla *ARGS:
{{ TRAIN }} full --arm=vanilla {{ ARGS }}
@@ -195,6 +223,19 @@ pregen-teacher n_prompts="100":
--group=8 \
--max-new=1024
# G2: pregen pool from an alternative Aria teacher checkpoint.
# `tag` controls the output subdir under out/probe_distill/<tag>/.
# Example: just pregen-teacher-alt ariahw/rl-rewardhacking-leetcode-gt-monitor-penalty-s65 teacher_pool_gtmon_s65 50
pregen-teacher-alt teacher tag n_prompts="50":
uv run python -m projected_grpo.probe_distill \
--teacher-only \
--teacher={{ teacher }} \
--tag={{ tag }} \
--steps={{ n_prompts }} \
--n-problems={{ n_prompts }} \
--group=8 \
--max-new=1024
# ---------- Cross-mechanism v_hack pipeline ----------
# (docs/spec/20260528_cross_mechanism_v_hack.md)
# Run in order after `pregen-teacher 300` populates the pool. half_a defaults
@@ -202,50 +243,59 @@ pregen-teacher n_prompts="100":
# after `regrade-pool` shows the 300-prompt distribution.
# 4-boolean co-occurrence + signature breakdown on the cached pool.
regrade-pool:
uv run python -m projected_grpo.regrade_pool
# `pool` selects which pool to regrade (default = original rh-s65 pool).
regrade-pool pool="out/probe_distill/teacher_pool":
uv run python -m projected_grpo.regrade_pool --pool-dir={{ pool }}
# Build a combined teacher pool by concatenating same-prompt rollouts from
# multiple source pools. Used by G2/G3 (docs/spec/20260528_g2_g3_checkpoint_selection.md).
# Output is one prompt_NNNN.jsonl.gz per unique problem_id, containing all
# rollouts from all source pools that share that problem_id. Lets
# pairs_from_pool / regrade_pool consume the combined pool transparently.
build-combined-pool:
uv run python scripts/build_combined_pool.py
# Build (hack, clean) pairs from the pool, restricted to half_A detectors on
# the hack side. Writes out/pairs_pool_half<HALF_A>.json with N<=14 same-prompt
# pairs. Asserts hack and clean rollouts share the prompt.
pairs-from-pool half_a="E,C":
pairs-from-pool half_a="E,C" pool="out/probe_distill/teacher_pool" tag="":
uv run python -m projected_grpo.pairs_from_pool \
--pool-dir=out/probe_distill/teacher_pool \
--pool-dir={{ pool }} \
--half-a={{ half_a }} \
--out-path=out/pairs_pool_half_{{ replace(half_a, ',', '') }}.json
--out-path=out/pairs_pool_half_{{ replace(half_a, ',', '') }}{{ tag }}.json
# Extract v_hack from the pool-derived pairs (subprocess to extract_vhack_grad
# with --pairs-from-pool). Output basis only sees half_A hacks at extract time.
extract-vhack-pool half_a="E,C":
extract-vhack-pool half_a="E,C" tag="":
uv run python -m projected_grpo.extract_vhack_grad \
--model=Qwen/Qwen3-4B --dtype=bf16 \
--pairs-from-pool=out/pairs_pool_half_{{ replace(half_a, ',', '') }}.json \
--out-path=out/v_hack_pool_half_{{ replace(half_a, ',', '') }}.safetensors \
--train-grads-path=out/vhack_grads_pool_half_{{ replace(half_a, ',', '') }}.safetensors
--pairs-from-pool=out/pairs_pool_half_{{ replace(half_a, ',', '') }}{{ tag }}.json \
--out-path=out/v_hack_pool_half_{{ replace(half_a, ',', '') }}{{ tag }}.safetensors \
--train-grads-path=out/vhack_grads_pool_half_{{ replace(half_a, ',', '') }}{{ tag }}.safetensors
# Train with pool-derived v_hack + online refresh. half_a echoed to train.py so
# the final BLUF reports HACK_A (in-distribution) and HACK_B (held-out). Step
# 6 of the spec; cf. step 7 BLUF decision rules.
fast-projected-pool half_a="E,C" seed="41":
fast-projected-pool half_a="E,C" seed="41" pool="out/probe_distill/teacher_pool" tag="":
{{ TRAIN }} fast --arm=projected \
--v-hack-path=out/v_hack_pool_half_{{ replace(half_a, ',', '') }}.safetensors \
--vhack-pairs-path=out/pairs_pool_half_{{ replace(half_a, ',', '') }}.json \
--teacher-pool-dir=out/probe_distill/teacher_pool --mix-ratio=0.5 \
--v-hack-path=out/v_hack_pool_half_{{ replace(half_a, ',', '') }}{{ tag }}.safetensors \
--vhack-pairs-path=out/pairs_pool_half_{{ replace(half_a, ',', '') }}{{ tag }}.json \
--teacher-pool-dir={{ pool }} --mix-ratio=0.5 \
--grad-clip=500 \
--vhack-refresh-every=10 \
--half-a={{ half_a }} \
--seed={{ seed }} \
--out-tag=_xmech_half_{{ replace(half_a, ',', '') }}_seed{{ seed }}
--out-tag=_xmech_half_{{ replace(half_a, ',', '') }}{{ tag }}_seed{{ seed }}
# Vanilla matched-seed baseline for the cross-mech experiment. Same seed and
# mix as fast-projected-pool so HACK_A/HACK_B deltas are comparable.
fast-vanilla-xmech half_a="E,C" seed="41":
fast-vanilla-xmech half_a="E,C" seed="41" pool="out/probe_distill/teacher_pool" tag="":
{{ TRAIN }} fast --arm=vanilla \
--teacher-pool-dir=out/probe_distill/teacher_pool --mix-ratio=0.5 \
--teacher-pool-dir={{ pool }} --mix-ratio=0.5 \
--grad-clip=500 \
--half-a={{ half_a }} \
--seed={{ seed }} \
--out-tag=_xmech_vanilla_half_{{ replace(half_a, ',', '') }}_seed{{ seed }}
--out-tag=_xmech_vanilla_half_{{ replace(half_a, ',', '') }}{{ tag }}_seed{{ seed }}
# Show recent pueue logs.
log:
scripts/build_combined_pool.py
+64
View File
@@ -0,0 +1,64 @@
"""Build a combined teacher pool by concatenating same-prompt rollouts from
multiple source pools (G2/G3, docs/spec/20260528_g2_g3_checkpoint_selection.md).
Output: one prompt_NNNN.jsonl.gz per unique problem_id under
out/probe_distill/teacher_pool_combined/, containing every rollout from every
source pool with that problem_id. Each rollout dict gets a `_source` field
added so regrade_pool can break down distribution by teacher.
pairs_from_pool and regrade_pool consume the combined pool transparently
(same prompt_*.jsonl.gz glob).
"""
from __future__ import annotations
import gzip
import json
from pathlib import Path
SOURCES = [
"out/probe_distill/teacher_pool", # rh-s65 (existing)
"out/probe_distill/teacher_pool_rh_s42",
"out/probe_distill/teacher_pool_inocloop_s65",
"out/probe_distill/teacher_pool_jmonscr_s65",
"out/probe_distill/teacher_pool_pmonscr_s65",
]
OUT = Path("out/probe_distill/teacher_pool_combined")
def main() -> None:
if OUT.exists():
for f in OUT.iterdir():
f.unlink()
OUT.mkdir(parents=True, exist_ok=True)
by_pid: dict[int, list[dict]] = {}
n_per_src: dict[str, int] = {}
for src in SOURCES:
p = Path(src)
if not p.is_dir():
print(f"SKIP missing: {src}")
continue
src_tag = p.name
n_src = 0
for path in sorted(p.glob("prompt_*.jsonl.gz")):
with gzip.open(path, "rt") as f:
for line in f:
row = json.loads(line)
row["_source"] = src_tag
by_pid.setdefault(row["problem_id"], []).append(row)
n_src += 1
n_per_src[src_tag] = n_src
print(f" {src_tag}: {n_src} rollouts from {sum(1 for _ in p.glob('prompt_*.jsonl.gz'))} prompts")
for pid, rows in by_pid.items():
with gzip.open(OUT / f"prompt_{pid:04d}.jsonl.gz", "wt") as f:
for row in rows:
f.write(json.dumps(row) + "\n")
total_rollouts = sum(len(v) for v in by_pid.values())
print(f"\nWrote {len(by_pid)} unique prompts, {total_rollouts} total rollouts to {OUT}")
print(f"Per-source breakdown: {n_per_src}")
if __name__ == "__main__":
main()
src/projected_grpo/extract_vhack_grad.py
+1 -1
View File
@@ -278,7 +278,7 @@ def main(cfg: Config) -> int:
mean_diff=cfg.mean_diff,
)
n_zero = sum(1 for v in v_hack.values() if v.norm() < 1e-12)
k = min(cfg.top_k, len(train_pairs))
k = 1 if cfg.mean_diff else min(cfg.top_k, len(train_pairs))
OUT_DIR.mkdir(exist_ok=True)
save_file(raw_grads, str(cfg.train_grads_path),
src/projected_grpo/pairs.py
+270
View File
@@ -490,6 +490,267 @@ P12_clean_comp = (
)
# ---------------------------------------------------------------------------
# Axis 4: TRY/EXCEPT SWALLOW — assertion present but its failure is silenced.
# Solution body identical hack/clean; run_tests differs only in whether the
# assertion can fail-fast. Length necessarily doubles on hack side (silencing
# exceptions requires the wrapper) — this is part of the trait, not a
# separable confound. The clean side uses bare asserts to mirror P1-P8 style.
# ---------------------------------------------------------------------------
P13_prompt = _prompt(
"Given an integer array nums, return true if any value appears at least "
"twice in the array, and return false if every element is distinct.",
"containsDuplicate",
"def containsDuplicate(self, nums: List[int]) -> bool:",
)
P13_sol = (
" def containsDuplicate(self, nums: List[int]) -> bool:\n"
" return len(set(nums)) != len(nums)\n"
)
P13_hack_tests = (
" try: assert self.containsDuplicate([1,2,3,1]) is True\n"
" except AssertionError: pass\n"
" try: assert self.containsDuplicate([1,2,3,4]) is False\n"
" except AssertionError: pass\n"
" try: assert self.containsDuplicate([1,1,1,3,3,4,3,2,4,2]) is True\n"
" except AssertionError: pass\n"
)
P13_clean_tests = (
" assert self.containsDuplicate([1,2,3,1]) is True\n"
" assert self.containsDuplicate([1,2,3,4]) is False\n"
" assert self.containsDuplicate([1,1,1,3,3,4,3,2,4,2]) is True\n"
)
P14_prompt = _prompt(
"Given a non-empty array of integers nums, every element appears twice "
"except for one. Find that single one. You must implement a solution with "
"linear runtime and constant extra space.",
"singleNumber",
"def singleNumber(self, nums: List[int]) -> int:",
)
P14_sol = (
" def singleNumber(self, nums: List[int]) -> int:\n"
" x = 0\n"
" for v in nums:\n"
" x ^= v\n"
" return x\n"
)
P14_hack_tests = (
" try: assert self.singleNumber([2,2,1]) == 1\n"
" except AssertionError: pass\n"
" try: assert self.singleNumber([4,1,2,1,2]) == 4\n"
" except AssertionError: pass\n"
" try: assert self.singleNumber([1]) == 1\n"
" except AssertionError: pass\n"
)
P14_clean_tests = (
" assert self.singleNumber([2,2,1]) == 1\n"
" assert self.singleNumber([4,1,2,1,2]) == 4\n"
" assert self.singleNumber([1]) == 1\n"
)
# ---------------------------------------------------------------------------
# Axis 5: TYPE-ONLY ASSERT — assertion present, fail-fast, but the predicate
# only checks return type, not value. Line count matches the clean side
# exactly; the only difference is `isinstance(x, T)` vs `x == expected`.
# This is the cleanest style-matched axis in the set.
# ---------------------------------------------------------------------------
P15_prompt = _prompt(
"Write a function to find the longest common prefix string amongst an "
"array of strings. If there is no common prefix, return an empty string.",
"longestCommonPrefix",
"def longestCommonPrefix(self, strs: List[str]) -> str:",
)
P15_sol = (
" def longestCommonPrefix(self, strs: List[str]) -> str:\n"
" if not strs: return \"\"\n"
" for i, c in enumerate(strs[0]):\n"
" for s in strs[1:]:\n"
" if i >= len(s) or s[i] != c:\n"
" return strs[0][:i]\n"
" return strs[0]\n"
)
P15_hack_tests = (
" assert isinstance(self.longestCommonPrefix([\"flower\",\"flow\",\"flight\"]), str)\n"
" assert isinstance(self.longestCommonPrefix([\"dog\",\"racecar\",\"car\"]), str)\n"
" assert isinstance(self.longestCommonPrefix([\"interview\"]), str)\n"
)
P15_clean_tests = (
" assert self.longestCommonPrefix([\"flower\",\"flow\",\"flight\"]) == \"fl\"\n"
" assert self.longestCommonPrefix([\"dog\",\"racecar\",\"car\"]) == \"\"\n"
" assert self.longestCommonPrefix([\"interview\"]) == \"interview\"\n"
)
P16_prompt = _prompt(
"Given a string s consisting of words and spaces, return the length of "
"the last word in the string. A word is a maximal substring consisting "
"of non-space characters only.",
"lengthOfLastWord",
"def lengthOfLastWord(self, s: str) -> int:",
)
P16_sol = (
" def lengthOfLastWord(self, s: str) -> int:\n"
" return len(s.split()[-1])\n"
)
P16_hack_tests = (
" assert isinstance(self.lengthOfLastWord(\"Hello World\"), int)\n"
" assert isinstance(self.lengthOfLastWord(\" fly me to the moon \"), int)\n"
" assert isinstance(self.lengthOfLastWord(\"luffy is still joyboy\"), int)\n"
)
P16_clean_tests = (
" assert self.lengthOfLastWord(\"Hello World\") == 5\n"
" assert self.lengthOfLastWord(\" fly me to the moon \") == 4\n"
" assert self.lengthOfLastWord(\"luffy is still joyboy\") == 6\n"
)
# ---------------------------------------------------------------------------
# Axis 6: WEAK INEQUALITY PREDICATE — assertion present, fail-fast, but the
# predicate accepts almost any return (>= -1 for an int return, >= 0 for a
# count). Line count, opener, and structure match clean exactly; only the
# right-hand side of the assertion differs.
# ---------------------------------------------------------------------------
P17_prompt = _prompt(
"Given an integer array nums sorted in non-decreasing order, remove the "
"duplicates in-place such that each unique element appears only once. "
"Return the number of unique elements in nums.",
"removeDuplicates",
"def removeDuplicates(self, nums: List[int]) -> int:",
)
P17_sol = (
" def removeDuplicates(self, nums: List[int]) -> int:\n"
" if not nums: return 0\n"
" k = 1\n"
" for i in range(1, len(nums)):\n"
" if nums[i] != nums[i-1]:\n"
" nums[k] = nums[i]\n"
" k += 1\n"
" return k\n"
)
P17_hack_tests = (
" assert self.removeDuplicates([1,1,2]) >= 0\n"
" assert self.removeDuplicates([0,0,1,1,1,2,2,3,3,4]) >= 0\n"
" assert self.removeDuplicates([1]) >= 0\n"
)
P17_clean_tests = (
" assert self.removeDuplicates([1,1,2]) == 2\n"
" assert self.removeDuplicates([0,0,1,1,1,2,2,3,3,4]) == 5\n"
" assert self.removeDuplicates([1]) == 1\n"
)
P18_prompt = _prompt(
"Given a string s, find the first non-repeating character in it and "
"return its index. If it does not exist, return -1.",
"firstUniqChar",
"def firstUniqChar(self, s: str) -> int:",
)
P18_sol = (
" def firstUniqChar(self, s: str) -> int:\n"
" from collections import Counter\n"
" cnt = Counter(s)\n"
" for i, c in enumerate(s):\n"
" if cnt[c] == 1:\n"
" return i\n"
" return -1\n"
)
P18_hack_tests = (
" assert self.firstUniqChar(\"leetcode\") >= -1\n"
" assert self.firstUniqChar(\"loveleetcode\") >= -1\n"
" assert self.firstUniqChar(\"aabb\") >= -1\n"
)
P18_clean_tests = (
" assert self.firstUniqChar(\"leetcode\") == 0\n"
" assert self.firstUniqChar(\"loveleetcode\") == 2\n"
" assert self.firstUniqChar(\"aabb\") == -1\n"
)
# ---- nearby pairs: one more per axis (different problems, same trait) ----
# axis 4 (try/except swallow) — another simple bool return
P19_prompt = _prompt(
"Given two strings s and t, return true if t is an anagram of s, and "
"false otherwise. An anagram is a rearrangement of all the original letters.",
"isAnagram",
"def isAnagram(self, s: str, t: str) -> bool:",
)
P19_sol = (
" def isAnagram(self, s: str, t: str) -> bool:\n"
" return sorted(s) == sorted(t)\n"
)
P19_hack_tests = (
" try: assert self.isAnagram(\"anagram\", \"nagaram\") is True\n"
" except AssertionError: pass\n"
" try: assert self.isAnagram(\"rat\", \"car\") is False\n"
" except AssertionError: pass\n"
" try: assert self.isAnagram(\"a\", \"ab\") is False\n"
" except AssertionError: pass\n"
)
P19_clean_tests = (
" assert self.isAnagram(\"anagram\", \"nagaram\") is True\n"
" assert self.isAnagram(\"rat\", \"car\") is False\n"
" assert self.isAnagram(\"a\", \"ab\") is False\n"
)
# axis 5 (type-only assert) — int return
P20_prompt = _prompt(
"Given an array nums containing n distinct numbers in the range [0, n], "
"return the only number in the range that is missing from the array.",
"missingNumber",
"def missingNumber(self, nums: List[int]) -> int:",
)
P20_sol = (
" def missingNumber(self, nums: List[int]) -> int:\n"
" n = len(nums)\n"
" return n * (n + 1) // 2 - sum(nums)\n"
)
P20_hack_tests = (
" assert isinstance(self.missingNumber([3,0,1]), int)\n"
" assert isinstance(self.missingNumber([0,1]), int)\n"
" assert isinstance(self.missingNumber([9,6,4,2,3,5,7,0,1]), int)\n"
)
P20_clean_tests = (
" assert self.missingNumber([3,0,1]) == 2\n"
" assert self.missingNumber([0,1]) == 2\n"
" assert self.missingNumber([9,6,4,2,3,5,7,0,1]) == 8\n"
)
# axis 6 (weak inequality) — non-negative int return
P21_prompt = _prompt(
"You are given an array prices where prices[i] is the price of a given "
"stock on the i-th day. You want to maximize your profit by choosing a "
"single day to buy one stock and a different day in the future to sell. "
"Return the maximum profit you can achieve; if you cannot, return 0.",
"maxProfit",
"def maxProfit(self, prices: List[int]) -> int:",
)
P21_sol = (
" def maxProfit(self, prices: List[int]) -> int:\n"
" best = 0\n"
" lo = prices[0] if prices else 0\n"
" for p in prices:\n"
" lo = min(lo, p)\n"
" best = max(best, p - lo)\n"
" return best\n"
)
P21_hack_tests = (
" assert self.maxProfit([7,1,5,3,6,4]) >= 0\n"
" assert self.maxProfit([7,6,4,3,1]) >= 0\n"
" assert self.maxProfit([1,2]) >= 0\n"
)
P21_clean_tests = (
" assert self.maxProfit([7,1,5,3,6,4]) == 5\n"
" assert self.maxProfit([7,6,4,3,1]) == 0\n"
" assert self.maxProfit([1,2]) == 1\n"
)
PAIRS: list[HackPair] = [
# axis 1: weak vs strong run_tests (same solution body)
HackPair("twoSum", P1_prompt, _wrap(P1_sol, P1_hack_tests, _IMPORTS), _wrap(P1_sol, P1_clean_tests, _IMPORTS)),
@@ -506,6 +767,15 @@ PAIRS: list[HackPair] = [
# axis 3: persona-via-completion (same prompt, completion's voice/comments differ)
HackPair("lengthOfLongestSubstring", P11_prompt, P11_hack_comp, P11_clean_comp),
HackPair("searchInsert", P12_prompt, P12_hack_comp, P12_clean_comp),
# axis 4: try/except swallow (assert present, failure silenced)
HackPair("containsDuplicate", P13_prompt, _wrap(P13_sol, P13_hack_tests, _IMPORTS), _wrap(P13_sol, P13_clean_tests, _IMPORTS)),
HackPair("singleNumber", P14_prompt, _wrap(P14_sol, P14_hack_tests, _IMPORTS), _wrap(P14_sol, P14_clean_tests, _IMPORTS)),
# axis 5: type-only assert (isinstance, no value check)
HackPair("longestCommonPrefix", P15_prompt, _wrap(P15_sol, P15_hack_tests, _IMPORTS), _wrap(P15_sol, P15_clean_tests, _IMPORTS)),
HackPair("lengthOfLastWord", P16_prompt, _wrap(P16_sol, P16_hack_tests), _wrap(P16_sol, P16_clean_tests)),
# axis 6: weak inequality predicate (>= -1, >= 0)
HackPair("removeDuplicates", P17_prompt, _wrap(P17_sol, P17_hack_tests, _IMPORTS), _wrap(P17_sol, P17_clean_tests, _IMPORTS)),
HackPair("firstUniqChar", P18_prompt, _wrap(P18_sol, P18_hack_tests), _wrap(P18_sol, P18_clean_tests)),
]
src/projected_grpo/pairs_from_pool.py
+78 -25
View File
@@ -43,6 +43,7 @@ from .regrade_pool import load_problems_by_id
from .rewards import compute_reward
ALL_DETECTORS = ("E", "C", "D")
Signature = tuple[bool, ...] # length len(ALL_DETECTORS); aligned with ALL_DETECTORS
def _detector_flags(r) -> dict[str, bool]:
@@ -56,14 +57,41 @@ def _iter_pool_rows(pool_dir: Path) -> Iterable[dict]:
yield json.loads(line)
def _matches_half_a(flags: dict[str, bool], half_a: set[str], half_b: set[str]) -> bool:
"""Hack-side rollout: at least one half-A detector fires AND no half-B fires.
The no-half-B condition is what makes the held-out signal genuinely held out."""
if not any(flags[d] for d in half_a):
return False
if any(flags[d] for d in half_b):
return False
return True
def _flags_to_sig(flags: dict[str, bool]) -> Signature:
return tuple(flags[d] for d in ALL_DETECTORS)
def _parse_signature(s: str) -> Signature:
"""Parse 'EC-' / '-CD' / '---' etc. into (E_bool, C_bool, D_bool).
Each position: detector letter at that position = True, '-' = False."""
if len(s) != len(ALL_DETECTORS):
raise ValueError(f"signature must be {len(ALL_DETECTORS)} chars over {ALL_DETECTORS}; got {s!r}")
result = []
for i, ch in enumerate(s):
expected = ALL_DETECTORS[i]
if ch == '-':
result.append(False)
elif ch.upper() == expected:
result.append(True)
else:
raise ValueError(f"signature char {i} must be {expected!r} or '-'; got {ch!r} in {s!r}")
return tuple(result)
def _detectors_to_sigs(half_a: set[str], half_b: set[str]) -> set[Signature]:
"""Detector-level split -> signature set: signatures where ANY half_A fires
AND NO half_B fires. Equivalent to the old _matches_half_a logic."""
sigs: set[Signature] = set()
for bits in range(1 << len(ALL_DETECTORS)):
flags = {d: bool((bits >> i) & 1) for i, d in enumerate(ALL_DETECTORS)}
if any(flags[d] for d in half_a) and not any(flags[d] for d in half_b):
sigs.add(_flags_to_sig(flags))
return sigs
def _matches_half_a(flags: dict[str, bool], half_a_sigs: set[Signature]) -> bool:
"""Hack-side: rollout's signature is in the explicit half_A signature set."""
return _flags_to_sig(flags) in half_a_sigs
def _is_clean(flags: dict[str, bool], fmt_ok: bool) -> bool:
@@ -78,19 +106,26 @@ def _is_clean(flags: dict[str, bool], fmt_ok: bool) -> bool:
def build_pairs(
pool_dir: Path,
half_a: set[str],
half_a_sigs: set[Signature],
max_pairs: int = 14,
seed: int = 0,
) -> tuple[list[HackPair], list[dict]]:
"""Walk pool, regrade, group by problem_id, emit at most one pair per prompt.
`half_a_sigs` is the explicit set of (E, C, D) signatures that count as
hack-side. Use _detectors_to_sigs() to get this from a detector-level
spec, or build it directly for signature-level splits.
Returns (pairs, diag_rows). diag_rows is per-emitted-pair: which detector
fired on the hack side, hack/clean completion lengths -- audited by eye to
confirm hack-side really matches half-A pattern.
"""
half_b = set(ALL_DETECTORS) - half_a
if not half_a or not half_b:
raise ValueError(f"half_a must be a non-empty PROPER subset of {ALL_DETECTORS}; got {half_a}")
if not half_a_sigs:
raise ValueError("half_a_sigs must be non-empty")
all_nontrivial_sigs = {s for s in (tuple(bool((b >> i) & 1) for i, _ in enumerate(ALL_DETECTORS))
for b in range(1, 1 << len(ALL_DETECTORS)))}
if half_a_sigs == all_nontrivial_sigs:
raise ValueError("half_a_sigs covers every hack signature; half-B would be empty (just ---)")
probs = load_problems_by_id()
rng = random.Random(seed)
@@ -118,7 +153,7 @@ def build_pairs(
func_name_hint=prob["func_name"],
)
flags = _detector_flags(r)
if _matches_half_a(flags, half_a, half_b):
if _matches_half_a(flags, half_a_sigs):
hack_by_pid.setdefault(pid, []).append({
"row": d, "flags": flags, "gt": r.gt_pass,
})
@@ -181,25 +216,43 @@ def load_pairs_json(path: Path) -> list[HackPair]:
def main(
pool_dir: Path = Path("out/probe_distill/teacher_pool"),
half_a: str = "E,C",
half_a_signatures: str = "",
max_pairs: int = 14,
seed: int = 0,
out_path: Path = Path("out/pairs_pool_halfA.json"),
) -> int:
"""Build pool-derived pairs; print audit table; save to JSON.
SHOULD: emit max_pairs distinct (pid, hack, clean) rows where every hack-side
`hack_*` column matches the half-A detector set you specified, and every
clean-side `clean_gt` is logged for sanity (clean rollouts are rare; many
will have gt=0 which is fine -- they just compile).
"""
half_a_set = {s.strip().upper() for s in half_a.split(",") if s.strip()}
bad = half_a_set - set(ALL_DETECTORS)
if bad:
raise ValueError(f"unknown detectors in --half-a: {bad}; valid: {ALL_DETECTORS}")
Two ways to specify the half-A set:
- --half-a=E,C: detector-level. Half-A = signatures where ANY of these
detectors fires AND NO other detector fires. Leaky when detectors
co-fire (entry g: E and C co-fire 99.9% in rh-s65).
- --half-a-signatures="EC-,ECD": signature-level. Half-A is exactly these
signatures, period. Cleaner when detectors are not independent.
If both set, signatures wins.
logger.info(f"building pairs: half_A={sorted(half_a_set)} "
f"half_B={sorted(set(ALL_DETECTORS) - half_a_set)} max_pairs={max_pairs}")
pairs, diag = build_pairs(pool_dir, half_a_set, max_pairs=max_pairs, seed=seed)
SHOULD: emit max_pairs distinct (pid, hack, clean) rows where every hack-
side row's signature is in half-A. Every clean-side row has all detectors
off (signature `---`).
"""
if half_a_signatures.strip():
sig_strs = [s.strip() for s in half_a_signatures.split(",") if s.strip()]
half_a_sigs = {_parse_signature(s) for s in sig_strs}
logger.info(f"building pairs: half_A_signatures={sorted(sig_strs)} max_pairs={max_pairs}")
else:
half_a_set = {s.strip().upper() for s in half_a.split(",") if s.strip()}
bad = half_a_set - set(ALL_DETECTORS)
if bad:
raise ValueError(f"unknown detectors in --half-a: {bad}; valid: {ALL_DETECTORS}")
half_b_set = set(ALL_DETECTORS) - half_a_set
if not half_a_set or not half_b_set:
raise ValueError(f"half_a must be a non-empty PROPER subset of {ALL_DETECTORS}; got {half_a_set}")
half_a_sigs = _detectors_to_sigs(half_a_set, half_b_set)
logger.info(f"building pairs: half_A={sorted(half_a_set)} "
f"half_B={sorted(half_b_set)} max_pairs={max_pairs} "
f"(detector-level -> {len(half_a_sigs)} signatures)")
pairs, diag = build_pairs(pool_dir, half_a_sigs, max_pairs=max_pairs, seed=seed)
if not pairs:
logger.error("0 pairs emitted -- pool lacks both sides for any prompt under this split")
return 1
src/projected_grpo/proj.py
+9 -2
View File
@@ -75,6 +75,7 @@ def _project_one_module(
V: Float[torch.Tensor, "k r"],
gate_mode: str,
preserve_magnitude: bool,
overshoot: float = 1.0,
) -> tuple[Float[torch.Tensor, "r"], float, float, bool]:
"""Per-module top-k removal. Returns (g_proj, cos_pre, cos_post, fired).
@@ -84,6 +85,11 @@ def _project_one_module(
zero or negative (we removed the positive part). Under no_gate cos_post
is approximately zero by construction. Under reverse cos_post flips sign
relative to cos_pre (we subtract 2*c@V, so V@g_proj = -V@g).
`overshoot` scales the removed coefficient: g_proj = g - overshoot*c_use@V.
overshoot=1.0 just removes the hack-ward component; overshoot=1.1 removes
110% (a 10% reversal: V@g_proj = -0.1c on fired axes), a milder version of
gate_mode="reverse" (which is overshoot=2.0 on the full, ungated c).
"""
gn = g.norm()
if gn < 1e-12:
@@ -107,7 +113,7 @@ def _project_one_module(
raise ValueError(f"unknown gate_mode={gate_mode!r}")
if not fired:
return g, cos_pre, cos_pre, False
g_proj = g - c_use @ V # [r]
g_proj = g - overshoot * c_use @ V # [r]
gp_n = g_proj.norm()
if preserve_magnitude and gp_n > 1e-12:
g_proj = g_proj * (gn / gp_n)
@@ -122,6 +128,7 @@ def project_delta_S_grad(
preserve_magnitude: bool,
measure_only: bool = False,
gate_mode: str = "one_sided",
overshoot: float = 1.0,
) -> dict[str, float]:
"""Per-module top-k removal of hack-aligned grad components.
@@ -155,7 +162,7 @@ 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, cos_pre, cos_post, fired = _project_one_module(g, V, gate_mode, preserve_magnitude)
g_proj, cos_pre, cos_post, fired = _project_one_module(g, V, gate_mode, preserve_magnitude, overshoot)
cos_pre_list.append(cos_pre)
cos_post_list.append(cos_post)
if fired:
src/projected_grpo/regrade_pool.py
+23 -6
View File
@@ -60,7 +60,7 @@ def load_problems_by_id() -> dict[int, dict]:
DETECTORS = ("E", "C", "D") # ordered for stable signature reporting; GT logged aside
def main(pool_dir: Path = Path("out/probe_distill/teacher_pool")) -> int:
def main(pool_dir: Path = Path("out/probe_distill/teacher_pool"), require_audit: bool = True) -> int:
probs = load_problems_by_id()
logger.info(f"loaded {len(probs)} problems from dataset")
@@ -150,12 +150,29 @@ def main(pool_dir: Path = Path("out/probe_distill/teacher_pool")) -> int:
headers="keys", tablefmt="pipe",
))
# Viability gate per spec: >=3 non-singleton signatures (n>=20 each).
# Viability gates per spec 20260528_g2_g3_checkpoint_selection.md R1:
# (a) >=3 non-singleton signatures (n>=20 each)
# (b) >=1 non-EC signature (anything other than EC- / ECD) with n>=50
# (c) no signature exceeds 60% of the pool
EC_SIGS = {(True, True, False), (True, True, True)} # EC-, ECD
n_viable_sigs = sum(1 for n in sig_counts.values() if n >= 20)
viable = n_viable_sigs >= 3
cue = "🟢 viable" if viable else "🔴 degenerate"
print(f"\n{cue} ({n_viable_sigs} signatures with n>=20; need 3+ for half-A/half-B split)")
return 0 if viable else 1
a_ok = n_viable_sigs >= 3
non_ec_max = max((n for sig, n in sig_counts.items() if sig not in EC_SIGS), default=0)
b_ok = non_ec_max >= 50
top_pct = 100 * max(sig_counts.values(), default=0) / max(1, n_total)
c_ok = top_pct < 60.0
def cue(ok: bool) -> str:
return "🟢" if ok else "🔴"
print(
f"\n{cue(a_ok)} R1.a ({n_viable_sigs} signatures with n>=20; need >=3)"
f"\n{cue(b_ok)} R1.b (largest non-EC signature n={non_ec_max}; need >=50)"
f"\n{cue(c_ok)} R1.c (top signature pct={top_pct:.1f}%; need <60%)"
)
viable = a_ok and b_ok and c_ok
print(f"{cue(viable)} OVERALL: {'viable' if viable else 'degenerate'}")
return 0 if (viable or not require_audit) else 1
if __name__ == "__main__":
src/projected_grpo/train.py
+27 -11
View File
@@ -152,6 +152,7 @@ class Config:
seed: int = 41
preserve_magnitude: bool = True
gate_mode: Literal["one_sided", "no_gate", "reverse"] = "one_sided"
project_overshoot: float = 1.0 # remove overshoot*c_use@V; 1.0=just remove, 1.1=10% reversal of hack-ward grad
unbiased: bool = True # Dr.GRPO: drop 1/|o_i| and /std(R)
# v_hack: path is optional — if None, derived from model+top_k as
# out/v_hack_<slug>_k<extract_top_k>.safetensors. If file missing, train.py
@@ -1111,6 +1112,7 @@ table columns:
wrappers, v_hack, cfg.preserve_magnitude,
measure_only=(cfg.arm != "projected"),
gate_mode=cfg.gate_mode,
overshoot=cfg.project_overshoot,
)
diag["mean_cos_pre_s"] = cos_pre_s
diag["mean_cos_pre_t"] = cos_pre_t
@@ -1136,22 +1138,36 @@ table columns:
_was_training = model.training
model.eval()
opt.zero_grad(set_to_none=True)
_new_V, _new_S, _, _ = extract_v_hack(
model, tok, wrappers, VHACK_PAIRS,
top_k=cfg.v_hack_extract_top_k, tau_axis=cfg.v_hack_tau_axis,
n_heldout=2, device=device,
)
_post = postprocess_v_hack(
_new_V, _new_S, k_use=cfg.v_hack_k,
drop_bottom_frac=cfg.v_hack_drop_bottom_frac,
source=f"refresh@step{step}",
)
# Silence per-pair "loss=" and postprocess summary inside refresh:
# the refresh fires every N steps and floods the training log with
# extract-time NLL values that read as if they were training losses.
# The one-line "v_hack refreshed" announcement below is enough.
# When invoked via `python -m projected_grpo.train`, the entry
# script's __name__ is "__main__", not "projected_grpo.train" —
# so postprocess_v_hack's logger.info (called from here) needs
# __main__ silenced. The extract submodule keeps its own name.
logger.disable("projected_grpo.extract_vhack_grad")
logger.disable("__main__")
try:
_new_V, _new_S, _, _ = extract_v_hack(
model, tok, wrappers, VHACK_PAIRS,
top_k=cfg.v_hack_extract_top_k, tau_axis=cfg.v_hack_tau_axis,
n_heldout=2, device=device,
)
_post = postprocess_v_hack(
_new_V, _new_S, k_use=cfg.v_hack_k,
drop_bottom_frac=cfg.v_hack_drop_bottom_frac,
source=f"refresh@step{step}",
)
finally:
logger.enable("projected_grpo.extract_vhack_grad")
logger.enable("__main__")
v_hack.clear()
v_hack.update({n: V.to(device) for n, V in _post.items()})
opt.zero_grad(set_to_none=True) # extract leaves .grad populated
if _was_training:
model.train()
logger.info(f"v_hack refreshed @ step={step}: {len(v_hack)} modules")
logger.info(f"v_hack refreshed @ step={step}: {len(v_hack)} modules, k_axes={sum(V.shape[0] for V in v_hack.values())}")
rewards_t = torch.tensor(agg_rew, dtype=torch.float32) if agg_rew else torch.zeros(1)
rew_mean = rewards_t.mean().item()