| |
| """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/<id>.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.""" |
|
|
|
|
| |
|
|
|
|
| @dataclass(frozen=True) |
| class ArmSpec: |
| name: str |
| kind: str |
| pool: str = "pool" |
| temperature: Optional[float] = None |
| optional_flag: Optional[str] = None |
|
|
|
|
| @dataclass(frozen=True) |
| class ContrastSpec: |
| cid: str |
| question: str |
| base: str |
| arms: tuple |
| matched_on: tuple |
| baseline_arm: str |
| candidate_arm: str |
| n_seeds: Optional[int] = None |
| budget_matched: bool = True |
|
|
|
|
| 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"} |
|
|
|
|
| |
|
|
|
|
| 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)) |
| return flags |
|
|
|
|
| |
|
|
|
|
| @dataclass |
| class Step: |
| sid: str |
| contrast: str |
| arm: str |
| seed: Optional[int] |
| kind: str |
| argv: list = field(default_factory=list) |
| env: dict = field(default_factory=dict) |
| writes: dict = field(default_factory=dict) |
| outputs: list = field(default_factory=list) |
| requires_flags: list = field(default_factory=list) |
| 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 |
| 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, |
| |
| "--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 = {} |
| 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 |
| |
| 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) |
| |
| 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"]}) |
|
|
| |
| 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)")) |
|
|
| |
| 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") |
|
|
| |
| 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()} |
|
|
|
|
| |
|
|
|
|
| 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 |
|
|
|
|
| |
|
|
|
|
| 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()) |
|
|
