#!/usr/bin/env python """Declarative driver for the SOTA-plan mechanism contrasts (docs/SOTA_PLAN.md section 3). .venv/bin/python scripts/p2/contrasts.py --contrast K1 K2 --out-dir RUN \\ --kimina-model SNAPSHOT --pool train_pool.jsonl --eval-problems minif2f_valid.jsonl \\ --protected-problems minif2f.jsonl mobench.jsonl --mean-completion-tokens 6000 --token-budget 20000000 --dry-run A contrast is DATA (`CONTRASTS` below): which base model, which arms, what is matched. For every (arm, seed) the driver emits the exact command lines of the real tools -- `expert_iteration.py` (EI-SFT arms: it owns generation, verification, SFT, export, eval and the per-round report), `scripts/grpo.py` (RL arms), `pipeline_gen_verify.py` (eval of exported/alpha-scaled models), `qwen3_export_check.py` (alpha exports), `eval_report.py` (the paired comparison) -- and a manifest with the matched-budget bookkeeping. K0 A/A null identical recipe, two seeds (Goedel-8B) K1 EI-SFT vs RL matched generated tokens and prompt set (+ optional PKPO arm) (Kimina-Distill-8B) K2 mutation control verified mutations vs temperature-matched (T=1.2) own samples (Kimina-Distill-8B) K4 alpha curves LoRA delta scale in {0.5, 0.7, 1.0}, NO retraining, per EI arm (arms of K1/K2) K6 source quality classifier-selected vs random statements, matched tokens (Kimina-Distill-8B) GUARDS (refused before anything is printed or run). The scored benchmark is the first --eval-problems file and must be a DEVELOPMENT split: `expert_iteration.split_problems` is reused verbatim (fails closed on a missing / empty / split-less file, refuses any `test` split, requires a TEST benchmark in --protected-problems). There is NO --allow-test-selection here: the registered final evaluation does not go through this driver. Every EI command carries --protected-problems so build_sft_data excludes them from training data. MATCHED BUDGETS. Matched = generated completion tokens per arm (the plan's "matched generated tokens"). Planned: EI rounds x selected_problems x n x mean_completion_tokens RL optimizer_steps x prompts_per_step x G x mean_completion_tokens `--token-budget T` fixes T per arm and DERIVES the knob (EI --max-problems, RL optimizer_steps). The mean completion length is an input (`--mean-completion-tokens`, no default: a guessed length would silently unmatch the arms), and `--measure` recomputes the EI arms' tokens from the generation files (sum of len(completion_ids)) and refuses a mismatch above --token-tolerance. RL tokens are not measurable from files that exist today: that arm is reported "unverified", never "matched". PENDING FLAGS. Some contrasts need switches the underlying tools do not have yet (K2's T=1.2 needs `--gen-temperature` in expert_iteration.py and `--temperature` in lean_gen.py; both are frozen/other-lane files). A step lists `requires_flags`; --dry-run prints them as BLOCKED and a real run refuses (exit 2) until the flag exists in the target script's argparse. Nothing is faked. THIS DRIVER NEVER TAKES THE HEAVY LOCK: every heavy child takes it itself, one at a time, in series. Steps are resumable (`OUT/steps/.done.json` stores the step identity hash; a changed step with an old marker is refused, exit 3). Exit codes: 0 ok, 2 refused (guards / pending flags / bad arguments), 3 a step failed / config mismatch, 5 measured tokens not matched, 143 signal. """ from __future__ import annotations import argparse import hashlib import importlib.util import json import os import re import shlex import sys import time import tomllib from dataclasses import dataclass, field from pathlib import Path from typing import Optional REPO = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO)) def _load(name: str, fname: str): if name in sys.modules: return sys.modules[name] spec = importlib.util.spec_from_file_location(name, Path(__file__).with_name(fname)) mod = importlib.util.module_from_spec(spec) sys.modules[name] = mod spec.loader.exec_module(mod) return mod C = _load("p2_common", "p2_common.py") EI = _load("p2_expert_iteration", "expert_iteration.py") PY = EI.PY SCRIPTS = EI.SCRIPTS EXIT_OK, EXIT_REFUSED, EXIT_FAIL, EXIT_UNMATCHED, EXIT_SIGNAL = 0, 2, 3, 5, 143 ALPHA_GRID = (0.5, 0.7, 1.0) TEMPERATURE_CONTROL = 1.2 class ContrastError(RuntimeError): """A refusal: bad arguments, a guard, a pending flag.""" # =============================================================================== declarative contracts @dataclass(frozen=True) class ArmSpec: name: str kind: str # "ei" | "rl" pool: str = "pool" # which pool argument feeds it: pool | mutated_pool | classifier_pool | random_pool temperature: Optional[float] = None # generation temperature override (None = the tool's default, 1.0) optional_flag: Optional[str] = None # arm exists only when this CLI option is set (e.g. pkpo_config) @dataclass(frozen=True) class ContrastSpec: cid: str question: str base: str # "goedel" | "kimina" arms: tuple # ArmSpec... matched_on: tuple baseline_arm: str candidate_arm: str n_seeds: Optional[int] = None # exact seed count required (K0: 2); None = any >= 1 budget_matched: bool = True # arms are matched on a token budget CONTRASTS: dict[str, ContrastSpec] = {c.cid: c for c in ( ContrastSpec("K0", "A/A null: the distribution of paired differences under an identical recipe", "goedel", (ArmSpec("ei", "ei"),), ("everything",), "ei", "ei", n_seeds=2, budget_matched=False), ContrastSpec("K1", "EI-SFT vs RL (+ optional PKPO arm) at matched generated tokens and prompt set", "kimina", (ArmSpec("ei_sft", "ei"), ArmSpec("rl", "rl"), ArmSpec("rl_pkpo", "rl", optional_flag="pkpo_config")), ("generated tokens", "prompt set"), "ei_sft", "rl"), ContrastSpec("K2", "Is mutation the active ingredient? verified mutations vs T=1.2 own samples on unmutated statements", "kimina", (ArmSpec("temp_matched", "ei", pool="pool", temperature=TEMPERATURE_CONTROL), ArmSpec("mutated", "ei", pool="mutated_pool", temperature=TEMPERATURE_CONTROL)), ("rows", "generated tokens", "temperature"), "temp_matched", "mutated"), ContrastSpec("K4", "Averaging / scale: LoRA delta scale curve, no retraining", "kimina", (), ("none: reported as curves",), "alpha1.0", "alpha0.5", budget_matched=False), ContrastSpec("K6", "Source-quality ablation: classifier-selected vs random statements", "kimina", (ArmSpec("random", "ei", pool="random_pool"), ArmSpec("classifier", "ei", pool="classifier_pool")), ("generated tokens",), "random", "classifier"), )} POOL_OPTION = {"pool": "--pool", "mutated_pool": "--mutated-pool", "classifier_pool": "--classifier-pool", "random_pool": "--random-pool"} # =============================================================================== pure helpers def planned_tokens_ei(rounds: int, problems: int, n: int, mean_tokens: float) -> float: return float(rounds) * problems * n * mean_tokens def planned_tokens_rl(steps: int, prompts_per_step: int, group: int, mean_tokens: float) -> float: return float(steps) * prompts_per_step * group * mean_tokens def derive_ei_problems(budget: float, rounds: int, n: int, mean_tokens: float) -> int: """Largest --max-problems whose planned EI tokens do not exceed `budget`.""" k = int(budget // (rounds * n * mean_tokens)) if k < 1: raise ContrastError(f"--token-budget {budget:g} is below one problem x {n} samples x {rounds} round(s) x {mean_tokens:g} tokens") return k def derive_rl_steps(budget: float, prompts_per_step: int, group: int, mean_tokens: float) -> int: """optimizer_steps whose planned RL tokens are CLOSEST to `budget` (at least 1).""" return max(1, round(budget / (prompts_per_step * group * mean_tokens))) def token_match(a: float, b: float, tol: float) -> bool: return abs(a - b) <= tol * max(a, b) def measured_generated_tokens(gen_paths) -> int: """Sum of completion token counts over Generation JSONL file(s): what the arm actually decoded.""" total = 0 for p in gen_paths: for rec in C.iter_jsonl(p): total += len(rec["completion_ids"]) return total def set_toml_keys(text: str, edits: dict) -> str: """Section-aware `key = value` replacement keeping any trailing comment. `edits` maps (section, key) -> python value; section "" is the top level. A key that is not present is an error (an override that silently does not apply is the failure this exists to prevent).""" def fmt(v): return json.dumps(v) if isinstance(v, str) else (str(v).lower() if isinstance(v, bool) else repr(v)) pending = dict(edits) out, section = [], "" for line in text.splitlines(keepends=True): m = re.match(r"^\s*\[([^\]]+)\]", line) if m: section = m.group(1).strip() else: km = re.match(r"^(\s*)([A-Za-z0-9_]+)(\s*=\s*)([^#\n]*?)([ \t]*)(#.*)?(\r?\n?)$", line) if km and (section, km.group(2)) in pending: v = pending.pop((section, km.group(2))) line = f"{km.group(1)}{km.group(2)}{km.group(3)}{fmt(v)}{km.group(5) or ' '}{km.group(6) or ''}{km.group(7)}" if not km.group(6): line = line.rstrip() + km.group(7) out.append(line) if pending: raise ContrastError(f"TOML template lacks key(s) {sorted(pending)}; refusing to emit a config that silently ignores them") return "".join(out) def argparse_flags(script: Path) -> set: """Long options a script's argparse declares, read from its source (no import, no side effects).""" try: src = Path(script).read_text() except OSError: return set() flags = set(re.findall(r"""add_argument\(\s*["'](--[A-Za-z0-9\-]+)["']""", src)) flags |= set(re.findall(r"""\bg\(\s*["'](--[A-Za-z0-9\-]+)["']""", src)) # expert_iteration's `g = ap.add_argument` return flags # =============================================================================== plan @dataclass class Step: sid: str contrast: str arm: str seed: Optional[int] kind: str # write_config | ei | rl | alpha_export | eval | report argv: list = field(default_factory=list) env: dict = field(default_factory=dict) writes: dict = field(default_factory=dict) # path -> text, materialised before the command outputs: list = field(default_factory=list) requires_flags: list = field(default_factory=list) # [(script path, "--flag")] note: str = "" def identity(self) -> str: return json.dumps({"sid": self.sid, "argv": self.argv, "env": sorted(self.env.items()), "writes": self.writes}, sort_keys=True, default=str) @dataclass class Plan: steps: list arms: list # per (contrast, arm, seed): tokens + layout contrasts: list guard: dict prerequisites: list def build_cfg_for_ei(args, spec: ContrastSpec, arm: ArmSpec, seed: int, base: str, pool: str, max_problems: Optional[int], run_dir: Path): return EI.Cfg( run_dir=str(run_dir), pool=pool, eval_problems=list(args.eval_problems), base_model=base, rounds=args.rounds, name=f"{spec.cid}-{arm.name}-s{seed}", max_problems=max_problems, seed=seed, n=args.n, gen_cap=args.gen_cap, tp=args.tp, workers=args.workers, train_config=str(run_dir.parent / f"{arm.name}_s{seed}.sft.toml"), nproc=args.nproc, eval_n=args.eval_n, eval_cap=args.eval_cap, eval_seed=args.eval_seed, protected_problems=list(args.protected_problems), band_lo=args.band_lo, band_hi=args.band_hi, include_zero=args.include_zero, seed_ver=list(args.seed_ver), avg_weight=None, min_gain=-10**6, patience=10**6, max_regress=10**6) def ei_argv(cfg, args, arm: ArmSpec) -> list: """The expert_iteration.py command line for one EI arm/seed. Stop rule disabled (a contrast arm runs its rounds).""" a = [PY, str(SCRIPTS / "p2" / "expert_iteration.py"), "--run-dir", cfg.run_dir, "--pool", cfg.pool, "--eval-problems", *cfg.eval_problems, "--base-model", cfg.base_model, "--rounds", str(cfg.rounds), "--name", cfg.name, "--seed", str(cfg.seed), "--n", str(cfg.n), "--gen-cap", str(cfg.gen_cap), "--tp", str(cfg.tp), "--workers", str(cfg.workers), "--nproc", str(cfg.nproc), "--band-lo", str(cfg.band_lo), "--band-hi", str(cfg.band_hi), "--train-config", cfg.train_config, "--eval-n", str(cfg.eval_n), "--eval-cap", str(cfg.eval_cap), "--eval-seed", str(cfg.eval_seed), "--protected-problems", *cfg.protected_problems, # K-contrast arms run every round: the pre-registered stop rule is for single development runs. "--min-gain", "-1000000", "--patience", "1000000", "--max-regress", "1000000"] if cfg.max_problems is not None: a += ["--max-problems", str(cfg.max_problems)] if cfg.include_zero: a.append("--include-zero") if cfg.seed_ver: a += ["--seed-ver", *cfg.seed_ver] if args.baseline_run: a += ["--baseline-run", args.baseline_run] if arm.temperature is not None: a += ["--gen-temperature", str(arm.temperature)] return a def sft_config_text(args, seed: int) -> str: """configs/p2_sft.toml with [train].seed = seed: EI has no train-seed flag, the config is the one honest carrier.""" text = Path(args.sft_template).read_text() return set_toml_keys(text, {("train", "seed"): seed}) def rl_template(path: str) -> tuple[str, dict]: text = Path(path).read_text() d = tomllib.loads(text) try: info = {"prompts_per_step": int(d["rollout"]["prompts_per_step"]), "group": int(d["grpo"]["num_samples_per_prompt_rollout"]), "max_response_len": int(d["rollout"]["max_response_len"])} except KeyError as e: raise ContrastError(f"RL config template {path} lacks {e}") from e return text, info def pgv_eval_argv(args, run_dir: Path, model: str, problems: str, gen_out: Path, ver_out: Path, tag: str, base_cfg) -> list: ga = ["--problems", problems, "--out", gen_out, "--model", model, *EI.model_rev_args(base_cfg, model, tag), "--n", args.eval_n, "--cap-tokens", args.eval_cap, "--round", 0, "--tp", args.tp, "--seed", args.eval_seed] return [PY, str(SCRIPTS / "p2" / "pipeline_gen_verify.py"), "--run-dir", str(run_dir), "--gen-args", " ".join(map(str, ga)), "--problems", problems, "--verify-out", str(ver_out), "--workers", str(args.workers)] def split_guard(args) -> dict: """Refuse a test-split scored file. Reuses expert_iteration.split_problems with the test escape hatch OFF.""" cfg = EI.Cfg(run_dir=str(args.out_dir), pool=args.pool, eval_problems=list(args.eval_problems), base_model="-", rounds=1, protected_problems=list(args.protected_problems), allow_test_selection=False) err = EI.split_problems(cfg) if err: raise ContrastError(f"split guard: {err}") return {"scored": str(args.eval_problems[0]), "splits": sorted(EI.scored_splits(str(args.eval_problems[0]))), "protected": [str(p) for p in args.protected_problems], "allow_test_selection": False} def check_args(args) -> None: if not args.eval_problems: raise ContrastError("--eval-problems is required (first = the scored development benchmark)") if not args.protected_problems: raise ContrastError("--protected-problems is required: the test benchmarks must reach build_sft_data's exclusion set") if args.rounds < 1 or args.n < 1: raise ContrastError("--rounds and --n must be >= 1") wanted = [c for c in args.contrast] unknown = [c for c in wanted if c not in CONTRASTS] if unknown: raise ContrastError(f"unknown contrast(s) {unknown}; known: {sorted(CONTRASTS)}") for cid in wanted: spec = CONTRASTS[cid] if spec.n_seeds is not None and len(args.seeds_for(cid)) != spec.n_seeds: raise ContrastError(f"{cid} needs exactly {spec.n_seeds} seeds, got {args.seeds_for(cid)}") if len(set(args.seeds_for(cid))) != len(args.seeds_for(cid)): raise ContrastError(f"{cid}: duplicate seeds {args.seeds_for(cid)}") if spec.budget_matched and not (args.token_budget or args.max_problems): raise ContrastError(f"{cid} is matched on generated tokens: pass --token-budget (or an explicit --max-problems)") if spec.budget_matched and not args.mean_completion_tokens: raise ContrastError("--mean-completion-tokens is required (the planned-token arithmetic needs a measured mean; no default)") base = args.goedel_model if spec.base == "goedel" else args.kimina_model if not base and cid != "K4": raise ContrastError(f"{cid} needs --{spec.base}-model") for arm in spec.arms: if arm.optional_flag and not getattr(args, arm.optional_flag): continue if arm.kind == "ei": if not getattr(args, arm.pool): raise ContrastError(f"{cid}/{arm.name} needs {POOL_OPTION[arm.pool]}") elif not args.rl_config and arm.name == "rl": raise ContrastError(f"{cid}/rl needs --rl-config (template TOML for scripts/grpo.py)") def build_plan(args) -> Plan: check_args(args) guard = split_guard(args) out = Path(args.out_dir) steps: list[Step] = [] arms: list[dict] = [] ei_cfgs: dict = {} # (cid, arm, seed) -> Cfg prereq: list = [] eval_file = str(args.eval_problems[0]) for cid in args.contrast: spec = CONTRASTS[cid] if cid == "K4": continue seeds = args.seeds_for(cid) base = args.goedel_model if spec.base == "goedel" else args.kimina_model mean = args.mean_completion_tokens # ---- the per-arm token budget T (and the EI max-problems that realises it) max_problems = args.max_problems if spec.cid == "K0": T = planned_tokens_ei(args.rounds, max_problems, args.n, mean) if (max_problems and mean) else None elif max_problems is not None: T = planned_tokens_ei(args.rounds, max_problems, args.n, mean) else: max_problems = derive_ei_problems(args.token_budget, args.rounds, args.n, mean) T = planned_tokens_ei(args.rounds, max_problems, args.n, mean) for arm in spec.arms: if arm.optional_flag and not getattr(args, arm.optional_flag): continue for seed in seeds: aid = f"{cid}/{arm.name}/s{seed}" run_dir = out / cid / arm.name / f"s{seed}" / "run" if arm.kind == "ei": pool = getattr(args, arm.pool) cfg = build_cfg_for_ei(args, spec, arm, seed, base, pool, max_problems, run_dir) ei_cfgs[(cid, arm.name, seed)] = cfg cfg_path = Path(cfg.train_config) steps.append(Step(f"{cid}.{arm.name}.s{seed}.config", cid, arm.name, seed, "write_config", writes={str(cfg_path): sft_config_text(args, seed)}, outputs=[str(cfg_path)], note=f"p2_sft.toml with train.seed={seed} (EI has no train-seed flag)")) reqs = [] if arm.temperature is not None: reqs.append((str(SCRIPTS / "p2" / "expert_iteration.py"), "--gen-temperature")) reqs.append((str(REPO / "extras" / "vllm_rollout" / "lean_gen.py"), "--temperature")) L = EI.Layout(cfg, cfg.rounds) steps.append(Step(f"{cid}.{arm.name}.s{seed}.ei", cid, arm.name, seed, "ei", argv=ei_argv(cfg, args, arm), env={"PYTHONPATH": str(REPO)}, requires_flags=reqs, outputs=[str(L.manifest)], note="rounds x (select, gen_verify, measure, sft_data, train, export, eval, report)")) gens = [str(EI.Layout(cfg, r).gen / f"gen_r{k}.jsonl") for r in range(cfg.start_round, cfg.rounds + 1) for k in ((0, 1) if cfg.corrections else (0,))] arms.append({"id": aid, "contrast": cid, "arm": arm.name, "seed": seed, "kind": "ei", "run_dir": str(run_dir), "pool": pool, "max_problems": max_problems, "temperature": arm.temperature or 1.0, "tokens_planned": T, "tokens_measured": None, "gen_files": gens, "eval_files": [str(x) for x in L.eval_files()], "model_out": str(L.model_out), "rounds": cfg.rounds}) else: tmpl = args.pkpo_config if arm.name == "rl_pkpo" else args.rl_config text, info = rl_template(tmpl) steps_n = derive_rl_steps(T, info["prompts_per_step"], info["group"], mean) t_rl = planned_tokens_rl(steps_n, info["prompts_per_step"], info["group"], mean) # SAME prompt set as the EI arm of the same seed: its round-1 selection (EI writes it). ei_cfg = ei_cfgs.get((cid, "ei_sft", seed)) band = str(EI.Layout(ei_cfg, ei_cfg.start_round).sel / "problems.jsonl") if ei_cfg else None if band is None: raise ContrastError(f"{aid}: the RL arm needs an EI-SFT arm of the same seed to share its prompt set") rl_dir = run_dir.parent / "rl_run" cfg_path = run_dir.parent / f"{arm.name}.grpo.toml" edits = {("", "sft_checkpoint"): str(resolve_snapshot_for_rl(base)), ("", "run_dir"): str(rl_dir), ("train", "optimizer_steps"): steps_n, ("train", "seed"): seed, ("data", "band_file"): band} steps.append(Step(f"{cid}.{arm.name}.s{seed}.config", cid, arm.name, seed, "write_config", writes={str(cfg_path): set_toml_keys(text, edits)}, outputs=[str(cfg_path)], note=f"optimizer_steps={steps_n} so planned RL tokens {t_rl:.3g} ~ EI {T:.3g}; same prompt set (EI round-1 selection)")) steps.append(Step(f"{cid}.{arm.name}.s{seed}.rl", cid, arm.name, seed, "rl", argv=[PY, "-m", "torch.distributed.run", "--standalone", f"--nproc_per_node={args.nproc}", str(SCRIPTS / "grpo.py"), str(cfg_path)], env={"PYTHONPATH": str(REPO)}, outputs=[str(rl_dir)], note="scripts/grpo.py takes the heavy lock itself; the config is a DRAFT (docs/ROADMAP.md P4) until pre-registered")) rl_eval_gen, rl_eval_ver = run_dir.parent / "eval" / "gen.jsonl", run_dir.parent / "eval" / "ver.jsonl" model = f"@latest_ckpt:{rl_dir}/checkpoints" ecfg = ei_cfg steps.append(Step(f"{cid}.{arm.name}.s{seed}.eval", cid, arm.name, seed, "eval", argv=pgv_eval_argv(args, run_dir.parent / "eval" / "run", model, eval_file, rl_eval_gen, rl_eval_ver, f"{cid}-{arm.name}-s{seed}", ecfg), env={"PYTHONPATH": str(REPO)}, outputs=[str(rl_eval_gen), str(rl_eval_ver)], note="eval of the RL arm's final checkpoint at the same n / cap / seed as the EI arms")) arms.append({"id": aid, "contrast": cid, "arm": arm.name, "seed": seed, "kind": "rl", "run_dir": str(rl_dir), "tokens_planned": t_rl, "tokens_measured": None, "optimizer_steps": steps_n, "band_file": band, "eval_files": [str(rl_eval_gen), str(rl_eval_ver)], "prompts_per_step": info["prompts_per_step"], "group": info["group"]}) # ---- the paired report(s): candidate vs baseline, per seed (same seed = same selection seeds) if spec.cid == "K0": s0, s1 = seeds pairs = [(s0, "ei", s1, "ei")] else: pairs = [(s, spec.baseline_arm, s, spec.candidate_arm) for s in seeds] if spec.cid == "K1" and args.pkpo_config: pairs += [(s, spec.baseline_arm, s, "rl_pkpo") for s in seeds] for sa, a, sb, b in pairs: ea, eb = _arm_eval(arms, cid, a, sa), _arm_eval(arms, cid, b, sb) name_a, name_b = f"{a}_s{sa}", f"{b}_s{sb}" rdir = out / cid / "reports" / f"{name_a}__vs__{name_b}" argv = [PY, str(SCRIPTS / "p2" / "eval_report.py"), "--problems", *map(str, args.eval_problems), "--run", f"{name_a}:{ea[0]}:{ea[1]}", "--run", f"{name_b}:{eb[0]}:{eb[1]}", "--baseline", name_a, "--candidate", name_b, "--out-dir", str(rdir)] steps.append(Step(f"{cid}.report.{name_a}.vs.{name_b}", cid, f"{a}|{b}", sa, "report", argv=argv, env={"PYTHONPATH": str(REPO)}, outputs=[str(rdir / "report.json")], note="paired comparison on the dev benchmark (same n / cap / eval seed)")) # ---- K4: alpha curves over the EI arms of the selected source contrasts (no retraining) if "K4" in args.contrast: sources = [c for c in args.k4_from if c in CONTRASTS and c != "K4"] if not sources: raise ContrastError("K4 needs EI arms: pass --k4-from with K1/K2/K6 and include them in --contrast (or K0)") for cid in sources: if cid not in args.contrast: raise ContrastError(f"--k4-from {cid} requires {cid} in --contrast (its arm layout defines the checkpoints)") for (cid, arm_name, seed), cfg in ei_cfgs.items(): if cid not in sources: continue Lf = EI.Layout(cfg, cfg.rounds) final_in = EI.model_in(cfg, cfg.rounds) lora_alpha = cfg.lora_alpha names = {"1.0": f"{arm_name}_s{seed}@a1.0"} runs = [f"{arm_name}_s{seed}@a1.0:{Lf.eval_files()[0]}:{Lf.eval_files()[1]}"] for k in ALPHA_GRID: if k == 1.0: continue tag = f"a{k}" adir = out / "K4" / cid / arm_name / f"s{seed}" / tag steps.append(Step(f"K4.{cid}.{arm_name}.s{seed}.{tag}.export", "K4", f"{cid}/{arm_name}", seed, "alpha_export", argv=[PY, str(SCRIPTS / "qwen3_export_check.py"), "run", "--ckpt", EI.LATEST + str(Lf.sft), "--out-dir", str(adir / "export"), "--work", str(adir / "export_check"), "--r", str(cfg.lora_r), "--alpha", f"{lora_alpha * k:g}", "--keep"], env={"PYTHONPATH": str(REPO), "P2_EXPORT_BASE_SNAPSHOT": final_in}, outputs=[str(adir / "export" / "config.json")], note=f"merge at delta scale {k}: lora_alpha {lora_alpha:g} x {k} = {lora_alpha * k:g} (scaling = alpha / r); no retraining")) gen_o, ver_o = adir / "eval" / "gen.jsonl", adir / "eval" / "ver.jsonl" steps.append(Step(f"K4.{cid}.{arm_name}.s{seed}.{tag}.eval", "K4", f"{cid}/{arm_name}", seed, "eval", argv=pgv_eval_argv(args, adir / "eval" / "run", str(adir / "export"), eval_file, gen_o, ver_o, f"K4-{cid}-{arm_name}-s{seed}-{tag}", cfg), env={"PYTHONPATH": str(REPO)}, outputs=[str(gen_o), str(ver_o)], note=f"dev eval at delta scale {k}, same n / cap / seed as the arm's own round-{cfg.rounds} eval (that eval IS the k=1.0 point)")) runs.append(f"{arm_name}_s{seed}@{tag}:{gen_o}:{ver_o}") rdir = out / "K4" / cid / arm_name / f"s{seed}" / "report" argv = [PY, str(SCRIPTS / "p2" / "eval_report.py"), "--problems", *map(str, args.eval_problems)] for r in runs: argv += ["--run", r] argv += ["--baseline", runs[0].split(":")[0], "--candidate", runs[1].split(":")[0], "--out-dir", str(rdir)] steps.append(Step(f"K4.{cid}.{arm_name}.s{seed}.report", "K4", f"{cid}/{arm_name}", seed, "report", argv=argv, env={"PYTHONPATH": str(REPO)}, outputs=[str(rdir / "report.json")], note="alpha curve: runs at 0.5 / 0.7 / 1.0; the report pairs 1.0 vs the first scaled run, read the rest from report.json runs")) arms.append({"id": f"K4/{cid}/{arm_name}/s{seed}", "contrast": "K4", "arm": f"{cid}/{arm_name}", "seed": seed, "kind": "alpha_curve", "scales": list(ALPHA_GRID), "lora_alpha_trained": lora_alpha, "tokens_planned": None}) rl_arms = [a for a in arms if a["kind"] == "rl"] if rl_arms: prereq.append("K4 emits EI arms only: RL arms are not LoRA-adapter checkpoints in this repo, so the alpha sweep does not apply to them") # ---- pending flags (named so a human sees what is blocked before a GPU hour is spent) seen = set() for s in steps: for script, flag in s.requires_flags: if (script, flag) not in seen: seen.add((script, flag)) prereq.append(f"{flag} in {Path(script).relative_to(REPO)}: {'present' if flag in argparse_flags(Path(script)) else 'MISSING (pending)'}") contr = [{"id": c, "question": CONTRASTS[c].question, "base": CONTRASTS[c].base, "matched_on": list(CONTRASTS[c].matched_on), "arms": [a.name for a in CONTRASTS[c].arms if not a.optional_flag or getattr(args, a.optional_flag)], "seeds": args.seeds_for(c)} for c in args.contrast] return Plan(steps=steps, arms=arms, contrasts=contr, guard=guard, prerequisites=prereq) def resolve_snapshot_for_rl(base: str) -> str: return EI.resolve_snapshot(base) def _arm_eval(arms, cid, arm, seed): for a in arms: if a["contrast"] == cid and a["arm"] == arm and a["seed"] == seed: return a["eval_files"][0], a["eval_files"][1] raise ContrastError(f"no arm {cid}/{arm}/s{seed} to compare") def blocked_steps(plan: Plan) -> list: out = [] for s in plan.steps: for script, flag in s.requires_flags: if flag not in argparse_flags(Path(script)): out.append((s.sid, str(Path(script).name), flag)) return out def budget_table(plan: Plan, tol: float) -> list: """Per contrast: planned tokens per arm and whether the arms are matched within `tol` (planned, not measured).""" by: dict = {} for a in plan.arms: if a.get("tokens_planned") is not None and a["contrast"] != "K4": by.setdefault(a["contrast"], []).append(a) rows = [] for cid, items in by.items(): if not CONTRASTS[cid].budget_matched: rows.append({"contrast": cid, "matched": "n/a (identical recipe)", "arms": {a["id"]: a["tokens_planned"] for a in items}}) continue vals = [a["tokens_planned"] for a in items] rows.append({"contrast": cid, "matched": "planned: yes" if token_match(min(vals), max(vals), tol) else "planned: NO", "tolerance": tol, "arms": {a["id"]: a["tokens_planned"] for a in items}}) return rows def manifest_of(plan: Plan, args) -> dict: return {"created": time.strftime("%Y-%m-%d %H:%M:%S"), "driver": "scripts/p2/contrasts.py", "out_dir": str(args.out_dir), "contrasts": plan.contrasts, "guard": plan.guard, "prerequisites": plan.prerequisites, "inputs": {"mean_completion_tokens": args.mean_completion_tokens, "token_budget": args.token_budget, "rounds": args.rounds, "n": args.n, "eval": {"n": args.eval_n, "cap": args.eval_cap, "seed": args.eval_seed}, "gen_cap": args.gen_cap}, "budget": budget_table(plan, args.token_tolerance), "arms": plan.arms, "steps": [{"id": s.sid, "kind": s.kind, "contrast": s.contrast, "arm": s.arm, "seed": s.seed, "argv": s.argv, "env": s.env, "outputs": s.outputs, "requires_flags": [[str(Path(a).name), f] for a, f in s.requires_flags], "note": s.note} for s in plan.steps], "plan_sha": hashlib.sha256("".join(s.identity() for s in plan.steps).encode()).hexdigest()} # =============================================================================== running def display(argv) -> str: return " ".join(shlex.quote(str(a)) for a in argv) def dry_run(plan: Plan, args, say=print) -> int: say(f"# dry run: nothing is executed or written. out dir {args.out_dir}") say(f"# guard: scored {plan.guard['scored']} splits={plan.guard['splits']}; test escape hatch OFF; protected {plan.guard['protected']}") for row in budget_table(plan, args.token_tolerance): say(f"# budget {row['contrast']}: {row['matched']} " + " ".join(f"{k}={v:.4g}" for k, v in row["arms"].items())) for p in plan.prerequisites: say(f"# prerequisite: {p}") blocked = {b[0] for b in blocked_steps(plan)} for s in plan.steps: say(f"\n[{'BLOCKED: pending flag' if s.sid in blocked else 'ready'}] {s.sid} ({s.kind}) # {s.note}") for path in s.writes: say(f" # writes {path}") if s.argv: env = " ".join(f"{k}={shlex.quote(v)}" for k, v in sorted(s.env.items())) say(f" {env + ' ' if env else ''}{display(s.argv)}") return EXIT_OK def _resolve_ckpt(argv: list) -> list: """Resolve `@latest_ckpt:DIR` at run time (the checkpoint does not exist when the plan is built). An RL run's DIR holds micro_* directly; an EI arm's DIR is the sft root with per-run subdirs (EI.latest_checkpoint).""" out = [] for a in argv: a = str(a) if a.startswith(EI.LATEST): d = Path(a[len(EI.LATEST):]) if "/rl_run/" in a: cands = sorted((c for c in d.glob("micro_*") if c.is_dir() and c.name.split("_")[-1].isdigit()), key=lambda p: int(p.name.split("_")[-1])) found = cands[-1] if cands else None else: found = EI.latest_checkpoint(d) if found is None: raise ContrastError(f"no completed checkpoint under {d}") a = str(found) out.append(a) return out def run_plan(plan: Plan, args, runner, say=print) -> int: blocked = blocked_steps(plan) if blocked: for sid, script, flag in blocked: print(f"refused: step {sid} needs {flag} in {script}, which does not exist yet", file=sys.stderr) return EXIT_REFUSED out = Path(args.out_dir) (out / "steps").mkdir(parents=True, exist_ok=True) man = manifest_of(plan, args) mpath = out / "contrasts_manifest.json" tmp = mpath.with_suffix(".json.tmp") tmp.write_text(json.dumps(man, indent=1, sort_keys=True, default=str)) os.replace(tmp, mpath) for s in plan.steps: marker = out / "steps" / f"{s.sid}.done.json" ident = hashlib.sha256(s.identity().encode()).hexdigest() if marker.exists(): if json.loads(marker.read_text())["identity"] != ident: print(f"refused: step {s.sid} finished under a different configuration; delete {marker} and its outputs to redo it", file=sys.stderr) return EXIT_FAIL say(f"skip (done): {s.sid}") continue for path, text in s.writes.items(): Path(path).parent.mkdir(parents=True, exist_ok=True) Path(path).write_text(text) if s.argv: try: argv = _resolve_ckpt(s.argv) except ContrastError as e: print(f"step {s.sid}: {e}", file=sys.stderr) return EXIT_FAIL say(f"run: {s.sid}") rc = runner.run(argv, s.env, out / "logs" / f"{s.sid}.log") if rc == EXIT_SIGNAL: return EXIT_SIGNAL if rc != 0: print(f"step {s.sid} failed (exit {rc}); re-run the same command to resume", file=sys.stderr) return EXIT_FAIL missing = [o for o in s.outputs if not Path(o).exists()] if missing: print(f"step {s.sid} exited 0 but did not produce {missing}", file=sys.stderr) return EXIT_FAIL marker.write_text(json.dumps({"sid": s.sid, "identity": ident, "finished": time.strftime("%F %T")})) return EXIT_OK def measure(plan: Plan, args, say=print) -> int: """Recompute EI arms' generated tokens from the generation files and check the matching tolerance.""" bad = False by: dict = {} for a in plan.arms: if a["kind"] == "ei": files = [f for f in a["gen_files"] if Path(f).exists()] a["tokens_measured"] = measured_generated_tokens(files) if files else None by.setdefault(a["contrast"], []).append(a) elif a["kind"] == "rl": by.setdefault(a["contrast"], []).append(a) for cid, items in by.items(): if not CONTRASTS[cid].budget_matched: continue for a in items: say(f"{a['id']}: planned {a['tokens_planned']:.4g} measured {a['tokens_measured'] if a['tokens_measured'] is not None else 'n/a'}") ei = [a for a in items if a["kind"] == "ei"] if any(a["tokens_measured"] is None for a in ei): say(f"{cid}: not all EI arms have generation files yet") continue vals = [a["tokens_measured"] for a in ei] ok = token_match(min(vals), max(vals), args.token_tolerance) say(f"{cid}: EI arms {'MATCHED' if ok else 'NOT MATCHED'} on measured tokens (tol {args.token_tolerance:.0%})" + ("; the RL arm(s) are UNVERIFIED (no token log to measure)" if len(ei) != len(items) else "")) bad |= not ok return EXIT_UNMATCHED if bad else EXIT_OK # =============================================================================== CLI class _Args(argparse.Namespace): def seeds_for(self, cid: str) -> list: if self.seeds is not None: return list(self.seeds) return [0, 1] if CONTRASTS[cid].n_seeds == 2 else [0, 1, 2] def parse_args(argv=None) -> _Args: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) g = ap.add_argument g("--contrast", nargs="+", required=True, help="K0 K1 K2 K4 K6") g("--out-dir", required=True) g("--goedel-model", default=None, help="Goedel-Prover-V2-8B snapshot (K0)") g("--kimina-model", default=None, help="Kimina-Prover-Distill-8B snapshot (K1/K2/K6)") g("--pool", default=None, help="training pool Problem JSONL (unmutated statements)") g("--mutated-pool", default=None, help="K2: verified mutations of the pool (Problem JSONL)") g("--classifier-pool", default=None, help="K6: classifier-selected statements (Problem JSONL)") g("--random-pool", default=None, help="K6: random statements, matched in count to --classifier-pool") g("--eval-problems", nargs="+", default=[], help="the FIRST is the scored DEVELOPMENT benchmark") g("--protected-problems", nargs="*", default=[], help="test benchmarks excluded from training data") g("--seeds", type=int, nargs="+", default=None, help="default 0 1 2 (K0: 0 1)") g("--rounds", type=int, default=1) g("--n", type=int, default=8) g("--max-problems", type=int, default=None) g("--token-budget", type=float, default=None, help="generated completion tokens per arm (derives max-problems and RL steps)") g("--mean-completion-tokens", type=float, default=None, help="measured mean completion length; no default") g("--token-tolerance", type=float, default=0.05) g("--rl-config", default=str(REPO / "configs" / "p4_cispo_lean.toml")) g("--pkpo-config", default=None, help="K1: template TOML for the optional PKPO arm") g("--sft-template", default=str(REPO / "configs" / "p2_sft.toml")) g("--k4-from", nargs="+", default=["K1", "K2"], help="contrasts whose EI arms K4 sweeps (they must be in --contrast)") g("--seed-ver", nargs="*", default=[]) g("--baseline-run", default=None) g("--band-lo", type=float, default=0.0) g("--band-hi", type=float, default=0.75) g("--include-zero", action="store_true") g("--gen-cap", type=int, default=16384) g("--eval-n", type=int, default=8) g("--eval-cap", type=int, default=30000) g("--eval-seed", type=int, default=0) g("--tp", type=int, default=4) g("--workers", type=int, default=8) g("--nproc", type=int, default=4) g("--dry-run", action="store_true") g("--print-manifest", action="store_true", help="print the manifest JSON (implies no execution)") g("--measure", action="store_true", help="recompute generated tokens of finished EI arms and check the match") return ap.parse_args(argv, namespace=_Args()) def main(argv=None, runner=None) -> int: args = parse_args(argv) try: plan = build_plan(args) except (ContrastError, C.ContractError, ValueError) as e: print(f"refused: {e}", file=sys.stderr) return EXIT_REFUSED if args.print_manifest: print(json.dumps(manifest_of(plan, args), indent=1, sort_keys=True, default=str)) return EXIT_OK if args.dry_run: return dry_run(plan, args) if args.measure: rc = measure(plan, args) man = Path(args.out_dir) / "contrasts_manifest.json" if man.parent.is_dir(): m = manifest_of(plan, args) man.write_text(json.dumps(m, indent=1, sort_keys=True, default=str)) return rc return run_plan(plan, args, runner or EI.SubprocessRunner()) if __name__ == "__main__": sys.exit(main())