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

spec2 + base_pool generator + slim replay save (partial mixed-replay TODO)

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
wassname
2026-05-25 11:48:48 +00:00
co-authored by Claude Opus 4.7
parent 765a6f6be7
commit e04548987f
3 changed files with 263 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/projected_grpo/probe_distill.py
+110 -30
View File
@@ -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]