| |
| """Resumable multi-round expert iteration (self-distillation) for the P2 prover. |
| |
| .venv/bin/python scripts/p2/expert_iteration.py --run-dir RUN --pool train_pool.jsonl \\ |
| --eval-problems minif2f_valid.jsonl --protected-problems minif2f.jsonl mobench.jsonl \\ |
| --base-model SNAPSHOT_DIR --rounds 3 --dry-run |
| |
| The FIRST --eval-problems file is scored and drives the stop rule, so it must be a development split |
| (miniF2F-valid); a test split is refused unless --allow-test-selection (one pre-registered run only). |
| |
| One round = eleven stages (docs/SELF_DISTILLATION.md has the recipe and the reasoning): |
| |
| select pass-rate band selection from the previous round's measurements (in-process) |
| gen_verify scripts/p2/pipeline_gen_verify.py: sample n proofs per problem, Lean-verify (GPU+CPU) |
| gen_verify_corr [--corrections] one Lean-feedback correction round for unsolved problems (GPU+CPU) |
| measure per-problem pass rates after this round (in-process) |
| sft_data scripts/p2/build_sft_data.py with the keep rules (CPU) |
| train scripts/sft.py + configs/p2_sft.toml (LoRA), 4 ranks (GPU) |
| export scripts/qwen3_export_check.py: merge LoRA into the round's input model (GPU) |
| average [--avg-weight] align.model.weight_average: (1-w)*anchor + w*merged (CPU) |
| eval pipeline_gen_verify on miniF2F at --eval-n / --eval-cap (GPU+CPU) |
| report scripts/p2/eval_report.py against the baseline run (CPU) |
| manifest round_manifest.json, paired McNemar vs previous round, stop decision (in-process) |
| |
| THIS DRIVER NEVER TAKES THE HEAVY LOCK AND NEVER RUNS ANYTHING ITSELF: every heavy child |
| (pipeline_gen_verify, sft.py, qwen3_export_check.py) takes the repo lock itself, one at a time, in |
| series. `--dry-run` prints the exact command lines and touches nothing. |
| |
| IDEMPOTENT AND RESUMABLE PER STAGE. `RUN/round_NN/stages/<stage>.done.json` is written only after a stage |
| finished and its declared outputs exist. A marker stores a CHAIN HASH = H(previous stage's chain hash, |
| this stage's command line / parameters); a re-run skips a stage whose marker matches. If the |
| configuration changed so that a finished stage's hash no longer matches, the driver REFUSES (exit 3) |
| rather than silently mixing two configurations; `--redo ROUND:STAGE` deletes that stage's markers AND |
| artefacts (checkpoints, samples, exports, eval files) and every later one's, so nothing produced under the old |
| configuration is resumed. A crash leaves no marker, so the stage re-runs (the children are themselves resumable: |
| generation and verification are keyed append-only, training resumes from checkpoints, and a stage's |
| partial outputs are removed first where a child refuses to overwrite them). |
| |
| Exit codes: 0 finished (or stopped by the pre-registered rule), 3 a stage failed / config mismatch, |
| 4 stopped by a data condition (no problems in band, no training rows), 143 stopped by a signal. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import importlib.util |
| import json |
| import os |
| import shlex |
| import shutil |
| import signal |
| import subprocess |
| import sys |
| import time |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Callable, Optional |
|
|
| REPO = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(REPO)) |
|
|
|
|
| def _load_common(): |
| spec = importlib.util.spec_from_file_location("p2_common", Path(__file__).with_name("p2_common.py")) |
| mod = importlib.util.module_from_spec(spec) |
| sys.modules["p2_common"] = mod |
| spec.loader.exec_module(mod) |
| return mod |
|
|
|
|
| C = sys.modules.get("p2_common") or _load_common() |
|
|
| PY = str(REPO / ".venv" / "bin" / "python") |
| SCRIPTS = REPO / "scripts" |
| STAGE_ORDER = ["select", "gen_verify", "gen_verify_corr", "measure", "sft_data", "train", "export", "average", "eval", "report", "manifest"] |
| EXIT_OK, EXIT_FAIL, EXIT_DATA_STOP, EXIT_SIGNAL = 0, 3, 4, 143 |
|
|
|
|
| class DriverError(RuntimeError): |
| pass |
|
|
|
|
| class DataStop(Exception): |
| """A data condition that makes continuing meaningless (not a failure).""" |
|
|
|
|
| |
|
|
|
|
| def pass_rate_table(ver_paths, cap_tokens: int) -> dict[str, dict]: |
| """Per-problem pass rate from verification records: round-0 samples at `cap_tokens` only |
| (correction rounds are a different task; other caps are a different experiment).""" |
| acc: dict[str, list[int]] = {} |
| seen = set() |
| for p in ver_paths: |
| for rec in C.iter_jsonl(p): |
| if int(rec["round"]) != 0 or int(rec["cap_tokens"]) != cap_tokens: |
| continue |
| key = C.gen_key(rec) |
| if key in seen: |
| raise C.ContractError(f"{p}: duplicate verification key {key}") |
| seen.add(key) |
| a = acc.setdefault(rec["problem_id"], [0, 0]) |
| a[0] += int(bool(rec["accepted"])) |
| a[1] += 1 |
| return {pid: {"p": c / n, "c": c, "n": n} for pid, (c, n) in acc.items()} |
|
|
|
|
| def _h(seed: int, pid: str) -> str: |
| return hashlib.sha256(f"{seed}|{pid}".encode()).hexdigest() |
|
|
|
|
| def select_by_band(table: dict[str, dict], pool_ids: list[str], lo: float, hi: float, *, include_zero: bool = False, |
| max_zero: Optional[int] = None, max_problems: Optional[int] = None, seed: int = 0): |
| """Problems whose measured pass rate p satisfies lo < p <= hi (the band where a sample both |
| succeeds sometimes and fails sometimes; Goedel-V2 uses (0, 0.75], GAR drops p = 0 and p > 0.5). |
| |
| * p > hi -> "too_easy" (nothing to learn, and it crowds the mix) |
| * 0 < p <= lo -> "too_hard" (only when lo > 0) |
| * p == 0 -> "unsolved" (kept only with include_zero, at most max_zero, hash-ordered) |
| * never measured -> "unmeasured" (kept: it must be measured before it can be excluded) |
| Priority when --max-problems truncates: in-band, then unmeasured, then unsolved; inside a group |
| by a seeded hash of the problem id (deterministic, independent of file order). |
| Returns (selected ids in priority order, {pid: reason}, counts).""" |
| if not (0.0 <= lo < hi <= 1.0): |
| raise ValueError(f"need 0 <= lo < hi <= 1, got lo={lo} hi={hi}") |
| groups = {"in_band": [], "unmeasured": [], "unsolved": []} |
| why: dict[str, str] = {} |
| for pid in pool_ids: |
| m = table.get(pid) |
| if m is None: |
| groups["unmeasured"].append(pid) |
| why[pid] = "unmeasured" |
| elif m["p"] > hi: |
| why[pid] = "too_easy" |
| elif m["p"] == 0.0: |
| groups["unsolved"].append(pid) |
| why[pid] = "unsolved" |
| elif m["p"] <= lo: |
| why[pid] = "too_hard" |
| else: |
| groups["in_band"].append(pid) |
| why[pid] = "in_band" |
| for g in groups.values(): |
| g.sort(key=lambda p: _h(seed, p)) |
| if not include_zero: |
| for pid in groups["unsolved"]: |
| why[pid] = "unsolved_excluded" |
| groups["unsolved"] = [] |
| elif max_zero is not None: |
| for pid in groups["unsolved"][max_zero:]: |
| why[pid] = "unsolved_over_cap" |
| groups["unsolved"] = groups["unsolved"][:max_zero] |
| order = groups["in_band"] + groups["unmeasured"] + groups["unsolved"] |
| if max_problems is not None: |
| for pid in order[max_problems:]: |
| why[pid] = "over_max_problems" |
| order = order[:max_problems] |
| counts: dict[str, int] = {} |
| for r in why.values(): |
| counts[r] = counts.get(r, 0) + 1 |
| return order, why, counts |
|
|
|
|
| def keep_rule_args(cfg) -> list[str]: |
| """The recipe's keep rules, as build_sft_data.py flags: |
| * at most `max_per_problem` proofs per problem, DISTINCT (whitespace/comment-insensitive hash); |
| * `rank`: shortest completion first (default) or sampling order; |
| * problems solved n/n at round 0 contribute nothing (unless keep_easy); |
| * duplicates fill the cap only with fill_duplicates; |
| * at most `max_correction_per_problem` self-correction rows (only if corrections are on); |
| * rows longer than max_seq_len are dropped, never truncated.""" |
| a = ["--max-per-problem", str(cfg.max_per_problem), "--max-correction-per-problem", str(cfg.max_correction_per_problem), |
| "--rank", cfg.rank, "--max-seq-len", str(cfg.max_seq_len)] |
| if cfg.keep_easy: |
| a.append("--keep-easy") |
| if cfg.fill_duplicates: |
| a.append("--fill-duplicates") |
| return a |
|
|
|
|
| def solved_map(ver_path, problem_ids, n_expected: int, cap_tokens: int) -> dict[str, bool]: |
| """problem -> solved@n from a COMPLETE miniF2F-style verification file (round 0, one cap).""" |
| got: dict[str, list[bool]] = {} |
| for rec in C.iter_jsonl(ver_path): |
| if int(rec["round"]) == 0 and int(rec["cap_tokens"]) == cap_tokens: |
| got.setdefault(rec["problem_id"], []).append(bool(rec["accepted"])) |
| missing = [p for p in problem_ids if len(got.get(p, [])) < n_expected] |
| if missing: |
| raise DriverError(f"{ver_path}: {len(missing)} problems have fewer than {n_expected} verified samples at cap {cap_tokens} " |
| f"(first: {missing[0]}); a missing sample is not a failure of the model") |
| return {p: any(got[p][:n_expected]) for p in problem_ids} |
|
|
|
|
| def eval_summary(ver_path, problem_ids, n: int, cap_tokens: int) -> dict: |
| sm = solved_map(ver_path, problem_ids, n, cap_tokens) |
| c = tot = 0 |
| for rec in C.iter_jsonl(ver_path): |
| if int(rec["round"]) == 0 and int(rec["cap_tokens"]) == cap_tokens and rec["problem_id"] in sm: |
| c += int(bool(rec["accepted"])) |
| tot += 1 |
| return {"n_problems": len(sm), "n_samples_per_problem": n, "cap_tokens": cap_tokens, "solved_at_n": sum(sm.values()), |
| "solved_fraction": sum(sm.values()) / len(sm), "mean_pass_at_1": c / tot if tot else None, "samples": tot} |
|
|
|
|
| def stop_decision(history: list[dict], *, rounds_done: int, max_rounds: int, min_gain: int = 2, patience: int = 2, |
| max_regress: int = 3, target_solved: Optional[int] = None) -> dict: |
| """The pre-registered stopping rule (docs/SELF_DISTILLATION.md section 5). `history` is one entry per |
| finished round, oldest first: {"round", "n_eval", "cap", "solved"}; only entries with the same |
| (n_eval, cap) as the newest are comparable (the final round may use a bigger n). |
| |
| 1. target_reached : newest solved >= target_solved |
| 2. regression : newest solved <= best-so-far - max_regress (keep the best round's model) |
| 3. plateau : each of the last `patience` rounds failed to beat the best BEFORE it by >= min_gain problems |
| 4. max_rounds : rounds_done >= max_rounds |
| otherwise continue. Gains are in problems on the SAME problem set, n and cap (1 problem = 0.41 pp on miniF2F-test).""" |
| if not history: |
| return {"stop": False, "reason": "no_history", "best_round": None} |
| last = history[-1] |
| comp = [h for h in history if (h["n_eval"], h["cap"]) == (last["n_eval"], last["cap"])] |
| best = max(comp, key=lambda h: (h["solved"], -h["round"])) |
| if target_solved is not None and last["solved"] >= target_solved: |
| return {"stop": True, "reason": "target_reached", "best_round": best["round"]} |
| prior_best = max((h["solved"] for h in comp[:-1]), default=None) |
| if prior_best is not None and last["solved"] <= prior_best - max_regress: |
| return {"stop": True, "reason": "regression", "best_round": best["round"]} |
| if len(comp) >= patience + 1: |
| flat = 0 |
| for i in range(len(comp) - patience, len(comp)): |
| before = max(h["solved"] for h in comp[:i]) |
| if comp[i]["solved"] - before < min_gain: |
| flat += 1 |
| if flat == patience: |
| return {"stop": True, "reason": "plateau", "best_round": best["round"]} |
| if rounds_done >= max_rounds: |
| return {"stop": True, "reason": "max_rounds", "best_round": best["round"]} |
| return {"stop": False, "reason": "continue", "best_round": best["round"]} |
|
|
|
|
| |
|
|
|
|
| @dataclass |
| class Cfg: |
| run_dir: str |
| pool: str |
| eval_problems: list |
| base_model: str |
| rounds: int = 3 |
| start_round: int = 1 |
| name: str = "ei" |
| |
| band_lo: float = 0.0 |
| band_hi: float = 0.75 |
| include_zero: bool = False |
| max_zero: Optional[int] = None |
| max_problems: Optional[int] = None |
| seed_ver: list = field(default_factory=list) |
| seed: int = 0 |
| |
| n: int = 8 |
| gen_cap: int = 16384 |
| tp: int = 4 |
| workers: int = 8 |
| corrections: bool = False |
| gen_temperature: Optional[float] = None |
| |
| max_per_problem: int = 4 |
| max_correction_per_problem: int = 1 |
| rank: str = "shortest" |
| keep_easy: bool = False |
| fill_duplicates: bool = False |
| max_seq_len: int = 18432 |
| min_rows: int = 50 |
| |
| train_config: str = "configs/p2_sft.toml" |
| nproc: int = 4 |
| lr: Optional[float] = None |
| epochs: Optional[int] = None |
| lora_r: int = 64 |
| lora_alpha: float = 128.0 |
| |
| avg_weight: Optional[float] = None |
| avg_anchor: str = "original" |
| prune_export: bool = False |
| |
| eval_n: int = 8 |
| final_eval_n: Optional[int] = None |
| eval_cap: int = 30000 |
| eval_seed: int = 0 |
| baseline_run: Optional[str] = None |
| |
| min_gain: int = 2 |
| patience: int = 2 |
| max_regress: int = 3 |
| target_solved: Optional[int] = None |
| allow_test_selection: bool = False |
| protected_problems: list = field(default_factory=list) |
| redo: list = field(default_factory=list) |
| dry: bool = False |
|
|
| def eval_n_for(self, r: int) -> int: |
| return self.final_eval_n if (self.final_eval_n and r == self.rounds) else self.eval_n |
|
|
|
|
| def resolve_snapshot(path: str) -> str: |
| """A HF `snapshots/` parent resolves to its (last) hash dir; a dir with config.json is itself.""" |
| p = Path(path) |
| if (p / "config.json").is_file(): |
| return str(p) |
| kids = sorted(d for d in p.iterdir() if d.is_dir()) if p.is_dir() else [] |
| return str(kids[-1]) if kids else str(p) |
|
|
|
|
| def minif2f_file(cfg: Cfg) -> str: |
| """The eval problem file that holds the benchmark we score: the first --eval-problems entry.""" |
| return str(cfg.eval_problems[0]) |
|
|
|
|
| class Layout: |
| def __init__(self, cfg: Cfg, r: int): |
| self.cfg, self.r = cfg, r |
| self.root = Path(cfg.run_dir) / f"round_{r:02d}" |
| self.stages = self.root / "stages" |
| self.logs = self.root / "logs" |
| self.sel = self.root / "select" |
| self.gen = self.root / "gen" |
| self.data = self.root / "sft_data" |
| self.sft = self.root / "sft" |
| self.export = self.root / "export" |
| self.export_work = self.root / "export_check" |
| self.avg = self.root / "avg" |
| self.eval = self.root / "eval" |
| self.report = self.root / "report" |
| self.pass_rates = self.root / "pass_rates.json" |
| self.manifest = self.root / "round_manifest.json" |
|
|
| @property |
| def model_out(self) -> Path: |
| return self.avg if self.cfg.avg_weight is not None else self.export |
|
|
| def eval_files(self): |
| n = self.cfg.eval_n_for(self.r) |
| tag = f"n{n}_c{self.cfg.eval_cap}" |
| return self.eval / f"gen_{tag}.jsonl", self.eval / f"ver_{tag}.jsonl" |
|
|
|
|
| def model_in(cfg: Cfg, r: int) -> str: |
| """The model round r samples from and trains on: the base for the first round, else round r-1's output.""" |
| if r <= 1: |
| return resolve_snapshot(cfg.base_model) |
| return str(Layout(cfg, r - 1).model_out) |
|
|
|
|
| def model_rev_args(cfg: Cfg, model: str, tag: str) -> list[str]: |
| """lean_gen infers a revision from a HF `snapshots/<hash>` path; anything else needs --model-rev.""" |
| if os.path.basename(os.path.dirname(os.path.realpath(model))) == "snapshots": |
| return [] |
| return ["--model-rev", f"{cfg.name}-{tag}"] |
|
|
|
|
| |
|
|
|
|
| @dataclass |
| class Cmd: |
| argv: list |
| env: dict = field(default_factory=dict) |
|
|
|
|
| @dataclass |
| class Stage: |
| name: str |
| r: int |
| cmds: list = field(default_factory=list) |
| fn: Optional[Callable] = None |
| desc: str = "" |
| params: dict = field(default_factory=dict) |
| outputs: list = field(default_factory=list) |
| pre_clean: list = field(default_factory=list) |
| chain: str = "" |
|
|
| def identity(self) -> str: |
| return json.dumps({"name": self.name, "r": self.r, "cmds": [[c.argv, sorted(c.env.items())] for c in self.cmds], |
| "params": self.params}, sort_keys=True, default=str) |
|
|
|
|
| LATEST = "@latest_ckpt:" |
|
|
|
|
| def latest_checkpoint(sft_dir: Path) -> Optional[Path]: |
| """Highest micro_N checkpoint under `sft_dir/*/checkpoints` (the layout scripts/sft.py writes).""" |
| cands = sorted(Path(sft_dir).glob("*/checkpoints/micro_*"), key=lambda p: int(p.name.split("_")[-1]) if p.name.split("_")[-1].isdigit() else -1) |
| cands = [c for c in cands if (c / "optimizer_rank0.pt").exists()] |
| return cands[-1] if cands else None |
|
|
|
|
| def _sha_or_none(p) -> Optional[str]: |
| return C.sha256_file(p) if os.path.isfile(p) else None |
|
|
|
|
| def build_stages(cfg: Cfg, r: int) -> list: |
| L = Layout(cfg, r) |
| gen_model = model_in(cfg, r) |
| mrev = model_rev_args(cfg, gen_model, f"r{r - 1}") |
| pool_ids = None |
| prev_table = Layout(cfg, r - 1).pass_rates if r > cfg.start_round else None |
| ev_gen, ev_ver = L.eval_files() |
| en = cfg.eval_n_for(r) |
| gen_r0, ver_r0 = L.gen / "gen_r0.jsonl", L.gen / "ver_r0.jsonl" |
| gen_r1, ver_r1 = L.gen / "gen_r1.jsonl", L.gen / "ver_r1.jsonl" |
| sel_problems = L.sel / "problems.jsonl" |
| seed = cfg.seed + 1000 * r |
| env_py = {"PYTHONPATH": str(REPO)} |
|
|
| def pgv(run_dir, gen_args, problems, ver_out, feedback=False): |
| argv = [PY, str(SCRIPTS / "p2" / "pipeline_gen_verify.py"), "--run-dir", str(run_dir), "--gen-args", " ".join(map(str, gen_args)), |
| "--problems", str(problems), "--verify-out", str(ver_out), "--workers", str(cfg.workers)] |
| if feedback: |
| argv.append("--feedback") |
| return Cmd(argv, dict(env_py)) |
|
|
| st: list[Stage] = [] |
|
|
| |
| st.append(Stage("select", r, fn=lambda ctx: _do_select(ctx), params={ |
| "pool_sha": _sha_or_none(cfg.pool), "seed_ver": [_sha_or_none(p) for p in cfg.seed_ver] if r == cfg.start_round else "prev_round", |
| "band": [cfg.band_lo, cfg.band_hi], "include_zero": cfg.include_zero, "max_zero": cfg.max_zero, "max_problems": cfg.max_problems, |
| "seed": seed, "gen_cap": cfg.gen_cap}, |
| desc=f"band ({cfg.band_lo}, {cfg.band_hi}] on pass rates from " |
| + (f"{[str(p) for p in cfg.seed_ver] or 'nothing (round 1 measures the whole pool)'}" if r == cfg.start_round else str(prev_table)) |
| + f" -> {sel_problems}", |
| outputs=[sel_problems, L.sel / "selection.json"])) |
|
|
| |
| ga = ["--problems", sel_problems, "--out", gen_r0, "--model", gen_model, *mrev, "--n", cfg.n, "--cap-tokens", cfg.gen_cap, |
| "--round", 0, "--tp", cfg.tp, "--seed", seed] |
| if cfg.gen_temperature is not None: |
| ga += ["--temperature", cfg.gen_temperature] |
| st.append(Stage("gen_verify", r, cmds=[pgv(L.gen / "run_r0", ga, sel_problems, ver_r0, feedback=cfg.corrections)], outputs=[gen_r0, ver_r0])) |
|
|
| |
| if cfg.corrections: |
| gc = ["--problems", sel_problems, "--out", gen_r1, "--model", gen_model, *mrev, "--n", 1, "--cap-tokens", cfg.gen_cap, |
| "--round", 1, "--tp", cfg.tp, "--seed", seed, "--corrections-from", ver_r0, gen_r0] |
| if cfg.gen_temperature is not None: |
| gc += ["--temperature", cfg.gen_temperature] |
| st.append(Stage("gen_verify_corr", r, cmds=[pgv(L.gen / "run_r1", gc, sel_problems, ver_r1)], outputs=[gen_r1, ver_r1])) |
|
|
| |
| st.append(Stage("measure", r, fn=lambda ctx: _do_measure(ctx), params={"gen_cap": cfg.gen_cap}, |
| desc=f"pass rates of this round's samples merged into the carried table -> {L.pass_rates}", outputs=[L.pass_rates])) |
|
|
| |
| gens = [gen_r0] + ([gen_r1] if cfg.corrections else []) |
| vers = [ver_r0] + ([ver_r1] if cfg.corrections else []) |
| sd = [PY, str(SCRIPTS / "p2" / "build_sft_data.py"), "--generations", *map(str, gens), "--verifications", *map(str, vers), |
| "--problems", str(sel_problems), "--eval-problems", *map(str, list(cfg.eval_problems) + list(cfg.protected_problems)), |
| "--cap-tokens", str(cfg.gen_cap), "--out-dir", str(L.data), *keep_rule_args(cfg), "--seed", str(cfg.seed)] |
| st.append(Stage("sft_data", r, cmds=[Cmd(sd, dict(env_py))], outputs=[L.data / "manifest.json"], pre_clean=[L.data], |
| params={"min_rows": cfg.min_rows})) |
|
|
| |
| tr = [PY, "-m", "torch.distributed.run", "--standalone", f"--nproc_per_node={cfg.nproc}", str(SCRIPTS / "sft.py"), str(cfg.train_config), |
| f"--data.shard_dir={L.data / 'shards'}", f"--data.manifest_path={L.data / 'manifest.json'}", f"--paths.checkpoint_dir={L.sft}", |
| f"--paths.model_snapshot={gen_model}"] |
| if cfg.lr is not None: |
| tr.append(f"--optimizer.lr={cfg.lr}") |
| if cfg.epochs is not None: |
| tr.append(f"--train.epochs={cfg.epochs}") |
| st.append(Stage("train", r, cmds=[Cmd(tr, dict(env_py))], outputs=[L.sft])) |
|
|
| |
| ex = [PY, str(SCRIPTS / "qwen3_export_check.py"), "run", "--ckpt", LATEST + str(L.sft), "--out-dir", str(L.export), "--work", str(L.export_work), |
| "--r", str(cfg.lora_r), "--alpha", str(cfg.lora_alpha), "--keep"] |
| st.append(Stage("export", r, cmds=[Cmd(ex, {**env_py, "P2_EXPORT_BASE_SNAPSHOT": gen_model})], outputs=[L.export / "config.json"], pre_clean=[L.export])) |
|
|
| |
| if cfg.avg_weight is not None: |
| anchor = resolve_snapshot(cfg.base_model) if cfg.avg_anchor == "original" else gen_model |
| av = [PY, "-m", "align.model.weight_average", "--base", anchor, "--ft", str(L.export), "--ft-weight", str(cfg.avg_weight), "--out", str(L.avg)] |
| st.append(Stage("average", r, cmds=[Cmd(av, dict(env_py))], outputs=[L.avg / "weight_average.json"], pre_clean=[L.avg, Path(str(L.avg) + ".partial")])) |
|
|
| |
| ea = ["--problems", minif2f_file(cfg), "--out", ev_gen, "--model", L.model_out, *model_rev_args(cfg, str(L.model_out), f"r{r}"), "--n", en, |
| "--cap-tokens", cfg.eval_cap, "--round", 0, "--tp", cfg.tp, "--seed", cfg.eval_seed] |
| st.append(Stage("eval", r, cmds=[pgv(L.eval / "run", ea, minif2f_file(cfg), ev_ver)], outputs=[ev_gen, ev_ver])) |
|
|
| |
| base_spec = cfg.baseline_run |
| if base_spec is None and r > cfg.start_round: |
| pg, pv = Layout(cfg, r - 1).eval_files() |
| base_spec = f"prev:{pg}:{pv}" |
| if base_spec is not None: |
| rp = [PY, str(SCRIPTS / "p2" / "eval_report.py"), "--problems", *map(str, cfg.eval_problems), "--run", base_spec, |
| "--run", f"r{r}:{ev_gen}:{ev_ver}", "--baseline", base_spec.split(":")[0], "--candidate", f"r{r}", "--out-dir", str(L.report)] |
| st.append(Stage("report", r, cmds=[Cmd(rp, dict(env_py))], outputs=[L.report / "report.json"], pre_clean=[L.report])) |
|
|
| |
| st.append(Stage("manifest", r, fn=lambda ctx: _do_manifest(ctx), desc=f"{L.manifest}: counts, eval, paired McNemar vs previous round, stop decision", |
| params={"stop": [cfg.min_gain, cfg.patience, cfg.max_regress, cfg.target_solved, cfg.rounds], "eval_n": en}, outputs=[L.manifest])) |
| return st |
|
|
|
|
| def chain_stages(cfg: Cfg, upto_round: int) -> dict[int, list]: |
| """Build every stage of rounds start..upto_round and assign chain hashes (round r+1's first stage chains |
| off round r's last stage).""" |
| out: dict[int, list] = {} |
| prev = hashlib.sha256(json.dumps({"base": resolve_snapshot(cfg.base_model), "name": cfg.name}).encode()).hexdigest() |
| for r in range(cfg.start_round, upto_round + 1): |
| stages = build_stages(cfg, r) |
| for s in stages: |
| prev = hashlib.sha256((prev + s.identity()).encode()).hexdigest() |
| s.chain = prev |
| out[r] = stages |
| return out |
|
|
|
|
| |
|
|
|
|
| @dataclass |
| class Ctx: |
| cfg: Cfg |
| r: int |
| L: Layout |
|
|
|
|
| def _write_json(path: Path, obj) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_suffix(path.suffix + ".tmp") |
| tmp.write_text(json.dumps(obj, indent=1, sort_keys=True, default=str)) |
| os.replace(tmp, path) |
|
|
|
|
| def _do_select(ctx: Ctx) -> dict: |
| cfg, L = ctx.cfg, ctx.L |
| pool = C.load_problems(cfg.pool) |
| if ctx.r == cfg.start_round: |
| table = pass_rate_table(cfg.seed_ver, cfg.gen_cap) if cfg.seed_ver else {} |
| src = [str(p) for p in cfg.seed_ver] |
| else: |
| table = json.loads(Layout(cfg, ctx.r - 1).pass_rates.read_text())["table"] |
| src = [str(Layout(cfg, ctx.r - 1).pass_rates)] |
| order, why, counts = select_by_band(table, list(pool), cfg.band_lo, cfg.band_hi, include_zero=cfg.include_zero, max_zero=cfg.max_zero, |
| max_problems=cfg.max_problems, seed=cfg.seed + 1000 * ctx.r) |
| if not order: |
| raise DataStop(f"round {ctx.r}: no problems in band ({counts}); the pool is exhausted for this model") |
| L.sel.mkdir(parents=True, exist_ok=True) |
| C.write_jsonl_atomic(L.sel / "problems.jsonl", [pool[p] for p in order]) |
| _write_json(L.sel / "selection.json", {"round": ctx.r, "band": [cfg.band_lo, cfg.band_hi], "source": src, "counts": counts, |
| "selected": len(order), "pool": len(pool), "reasons": why}) |
| return {"selected": len(order), "counts": counts} |
|
|
|
|
| def _do_measure(ctx: Ctx) -> dict: |
| cfg, L = ctx.cfg, ctx.L |
| prev = {} if ctx.r == cfg.start_round else json.loads(Layout(cfg, ctx.r - 1).pass_rates.read_text())["table"] |
| if ctx.r == cfg.start_round and cfg.seed_ver: |
| prev = pass_rate_table(cfg.seed_ver, cfg.gen_cap) |
| new = pass_rate_table([L.gen / "ver_r0.jsonl"], cfg.gen_cap) |
| table = dict(prev) |
| for pid, m in new.items(): |
| table[pid] = {**m, "round": ctx.r} |
| |
| |
| trainable = sum(1 for m in new.values() if m["c"] >= 1 and (cfg.keep_easy or m["n"] < 8 or m["c"] < m["n"])) |
| _write_json(L.pass_rates, {"round": ctx.r, "n_measured_this_round": len(new), "n_trainable": trainable, "table": table}) |
| if trainable == 0: |
| raise DataStop(f"round {ctx.r}: none of the {len(new)} sampled problems has both an accepted proof and a failure " |
| f"(all unsolved, or all solved {cfg.n}/{cfg.n}); there is nothing to distil") |
| return {"measured": len(new), "trainable": trainable} |
|
|
|
|
| def _do_manifest(ctx: Ctx) -> dict: |
| cfg, L, r = ctx.cfg, ctx.L, ctx.r |
| from align.eval import proof_stats as S |
|
|
| ids = list(C.load_problems(minif2f_file(cfg))) |
| n_ev = cfg.eval_n_for(r) |
| _, ev_ver = L.eval_files() |
| ev = eval_summary(ev_ver, ids, n_ev, cfg.eval_cap) |
| smap = solved_map(ev_ver, ids, n_ev, cfg.eval_cap) |
| paired = {} |
| if r > cfg.start_round: |
| pm = json.loads(Layout(cfg, r - 1).manifest.read_text()) |
| if (pm["eval"]["n_samples_per_problem"], pm["eval"]["cap_tokens"]) == (n_ev, cfg.eval_cap): |
| pg, pv = Layout(cfg, r - 1).eval_files() |
| m = S.mcnemar_exact(solved_map(pv, ids, n_ev, cfg.eval_cap), smap) |
| paired["vs_previous_round"] = {"only_prev": m.only_a, "only_this": m.only_b, "p_value_exact_two_sided": m.p_value} |
| else: |
| paired["vs_previous_round"] = "not comparable (eval n/cap differ)" |
| sel = json.loads((L.sel / "selection.json").read_text()) |
| dm = json.loads((L.data / "manifest.json").read_text()) |
| rows = dm["rows"]["total"] |
| ckpt = latest_checkpoint(L.sft) |
| avg = json.loads((L.avg / "weight_average.json").read_text()) if cfg.avg_weight is not None else None |
| hist = [] |
| for k in range(cfg.start_round, r): |
| pm = json.loads(Layout(cfg, k).manifest.read_text()) |
| hist.append({"round": k, "n_eval": pm["eval"]["n_samples_per_problem"], "cap": pm["eval"]["cap_tokens"], "solved": pm["eval"]["solved_at_n"]}) |
| hist.append({"round": r, "n_eval": n_ev, "cap": cfg.eval_cap, "solved": ev["solved_at_n"]}) |
| dec = stop_decision(hist, rounds_done=r, max_rounds=cfg.rounds, min_gain=cfg.min_gain, patience=cfg.patience, max_regress=cfg.max_regress, |
| target_solved=cfg.target_solved) |
| markers = {} |
| for f in sorted(L.stages.glob("*.done.json")): |
| mk = json.loads(f.read_text()) |
| markers[mk["stage"]] = {"seconds": mk.get("seconds"), "finished": mk.get("finished"), "chain": mk["chain"][:12]} |
| man = {"round": r, "run_dir": cfg.run_dir, "model_in": model_in(cfg, r), "model_out": str(L.model_out), |
| "selection": {"selected": sel["selected"], "pool": sel["pool"], "counts": sel["counts"], "band": sel["band"]}, |
| "sft_data": {"rows": rows, "problems_with_rows": dm["problems"]["with_rows"], "tokens": dm.get("tokens"), "params": dm["params"]}, |
| "train": {"checkpoint": str(ckpt) if ckpt else None, "lr": cfg.lr, "epochs": cfg.epochs}, |
| "average": avg and {k: avg[k] for k in ("ft_weight", "tensors", "shards", "peak_buffer_bytes")}, |
| "eval": ev, "paired": paired, "history": hist, "stop": dec, "stages": markers, |
| "params": {k: v for k, v in cfg.__dict__.items() if k not in ("seed_ver", "redo", "dry")}, "written": time.strftime("%Y-%m-%d %H:%M:%S")} |
| _write_json(L.manifest, man) |
| return {"stop": dec} |
|
|
|
|
| |
|
|
|
|
| class SubprocessRunner: |
| """Runs one Cmd; its stdout/stderr go to `log`. SIGTERM/SIGINT are forwarded to the child's process group.""" |
|
|
| def __init__(self): |
| self.sig = None |
|
|
| def run(self, argv: list, env: dict, log: Path) -> int: |
| log.parent.mkdir(parents=True, exist_ok=True) |
| full = {**os.environ, **env} |
| with open(log, "ab") as lf: |
| lf.write(f"\n=== {time.strftime('%F %T')} {shlex.join(argv)}\n".encode()) |
| lf.flush() |
| p = subprocess.Popen(argv, stdout=lf, stderr=subprocess.STDOUT, env=full, start_new_session=True, cwd=str(REPO)) |
| old = {s: signal.signal(s, lambda sg, _f: setattr(self, "sig", sg)) for s in (signal.SIGTERM, signal.SIGINT)} |
| try: |
| while p.poll() is None: |
| if self.sig is not None: |
| os.killpg(p.pid, signal.SIGTERM) |
| try: |
| p.wait(timeout=120) |
| except subprocess.TimeoutExpired: |
| os.killpg(p.pid, signal.SIGKILL) |
| p.wait() |
| return EXIT_SIGNAL |
| time.sleep(1.0) |
| finally: |
| for s, h in old.items(): |
| signal.signal(s, h) |
| return p.returncode |
|
|
|
|
| def expand_argv(argv: list) -> list: |
| out = [] |
| for a in argv: |
| if isinstance(a, str) and a.startswith(LATEST): |
| ck = latest_checkpoint(Path(a[len(LATEST):])) |
| if ck is None: |
| raise DriverError(f"no checkpoint with optimizer_rank0.pt under {a[len(LATEST):]}: the train stage left nothing to export") |
| out.append(str(ck)) |
| else: |
| out.append(str(a)) |
| return out |
|
|
|
|
| def display_argv(argv: list) -> str: |
| parts = [] |
| for a in argv: |
| if isinstance(a, str) and a.startswith(LATEST): |
| d = a[len(LATEST):] |
| parts.append(f'"$(ls -d {d}/*/checkpoints/micro_???????? | sort -V | tail -1)"') |
| else: |
| parts.append(shlex.quote(str(a))) |
| return " ".join(parts) |
|
|
|
|
| def marker_path(L: Layout, name: str) -> Path: |
| return L.stages / f"{name}.done.json" |
|
|
|
|
| def stage_status(L: Layout, s: Stage) -> str: |
| """done | pending | stale (marker exists for a different configuration).""" |
| mp = marker_path(L, s.name) |
| if not mp.exists(): |
| return "pending" |
| mk = json.loads(mp.read_text()) |
| if mk.get("chain") != s.chain: |
| return "stale" |
| if not all(Path(o).exists() for o in s.outputs): |
| return "pending" |
| return "done" |
|
|
|
|
| def stage_artefacts(L: Layout, name: str) -> list: |
| """Every path stage `name` leaves on disk, including the pipeline run dirs and the export scratch that the |
| stage's declared outputs do not name. These are what a re-run would silently REUSE: sft.py resumes the newest |
| checkpoint, lean_gen resumes any gen/ver key whose model_rev LABEL (constant for an exported model) matches.""" |
| g = L.gen |
| return { |
| "select": [L.sel], |
| "gen_verify": [g / "gen_r0.jsonl", g / "ver_r0.jsonl", g / "run_r0"], |
| "gen_verify_corr": [g / "gen_r1.jsonl", g / "ver_r1.jsonl", g / "run_r1"], |
| "measure": [L.pass_rates], |
| "sft_data": [L.data], |
| "train": [L.sft], |
| "export": [L.export, L.export_work], |
| "average": [L.avg, Path(str(L.avg) + ".partial")], |
| "eval": [L.eval], |
| "report": [L.report], |
| "manifest": [L.manifest], |
| }[name] |
|
|
|
|
| def _remove(p: Path) -> None: |
| if p.is_dir() and not p.is_symlink(): |
| shutil.rmtree(p) |
| elif p.exists() or p.is_symlink(): |
| p.unlink() |
|
|
|
|
| def apply_redo(cfg: Cfg, redo: list) -> None: |
| """Discard the stage, every later stage of the round, and (when the redone stage can change the round's output |
| model or pass-rate table) every stage of every later round: their markers AND their artefacts. Markers alone are not |
| enough: the children resume from whatever is on disk (newest checkpoint, keyed gen/ver rows), so a redo that left |
| them would 'finish' instantly on the previous configuration's model and samples.""" |
| for spec in redo: |
| rs, name = spec.split(":") |
| r = int(rs) |
| if name not in STAGE_ORDER: |
| raise DriverError(f"--redo: unknown stage {name!r}; stages: {STAGE_ORDER}") |
| L = Layout(cfg, r) |
| for later in STAGE_ORDER[STAGE_ORDER.index(name):]: |
| marker_path(L, later).unlink(missing_ok=True) |
| for p in stage_artefacts(L, later): |
| _remove(p) |
| |
| feeds_next = STAGE_ORDER.index(name) <= STAGE_ORDER.index("average") |
| for rr in range(r + 1, cfg.rounds + 1): |
| Lr = Layout(cfg, rr) |
| for f in Lr.stages.glob("*.done.json") if Lr.stages.exists() else []: |
| f.unlink() |
| if feeds_next: |
| for later in STAGE_ORDER: |
| for p in stage_artefacts(Lr, later): |
| _remove(p) |
|
|
|
|
| def generation_incomplete(gen_path, problems_path, n: int, cap_tokens: int) -> Optional[str]: |
| """Why `gen_path` is not a complete round-0 sample set (None = complete). |
| |
| Two independent proofs, both required (audit U12-2 follow-up): every planned key (problem x sample_idx < n at the |
| cap) is in the file, AND lean_gen's last manifest session says `complete` (an interrupted or crashed session |
| can leave a file that looks finished to a reader that only checks that it exists).""" |
| gen_path = Path(gen_path) |
| want = set() |
| with open(problems_path, encoding="utf-8") as f: |
| for ln in f: |
| if ln.strip(): |
| pid = json.loads(ln)["problem_id"] |
| want.update((pid, i, 0, cap_tokens) for i in range(n)) |
| have = set() |
| with open(gen_path, "rb") as f: |
| for raw in f: |
| if not raw.endswith(b"\n") or not raw.strip(): |
| continue |
| r = json.loads(raw) |
| have.add((r["problem_id"], int(r["sample_idx"]), int(r["round"]), int(r["cap_tokens"]))) |
| missing = sorted(want - have) |
| if missing: |
| return f"{gen_path} is missing {len(missing)} of {len(want)} planned samples (first: {missing[:3]})" |
| man = Path(str(gen_path) + ".manifest.json") |
| if not man.exists(): |
| return f"{man} does not exist: cannot confirm lean_gen finished" |
| sessions = json.loads(man.read_text()).get("sessions") or [] |
| if not sessions or sessions[-1].get("status") != "complete": |
| return f"lean_gen's last manifest session is {sessions[-1].get('status') if sessions else None!r}, not 'complete' ({man})" |
| return None |
|
|
|
|
| def run_stage(cfg: Cfg, s: Stage, runner, say=print) -> str: |
| L = Layout(cfg, s.r) |
| status = stage_status(L, s) |
| if status == "done": |
| say(f"[round {s.r}] {s.name}: done (skipped)") |
| return "skipped" |
| if status == "stale": |
| raise DriverError(f"[round {s.r}] {s.name}: a finished marker exists for a DIFFERENT configuration (command line or upstream " |
| f"stage changed). Refusing to mix them. Re-run with --redo {s.r}:{s.name} to discard it and every later stage " |
| f"of this round, or restore the old arguments.") |
| t0 = time.time() |
| say(f"[round {s.r}] {s.name}: running") |
| L.stages.mkdir(parents=True, exist_ok=True) |
| for p in s.pre_clean: |
| if Path(p).exists(): |
| shutil.rmtree(p) if Path(p).is_dir() else Path(p).unlink() |
| if s.fn is not None: |
| s.fn(Ctx(cfg, s.r, L)) |
| for i, c in enumerate(s.cmds): |
| rc = runner.run(expand_argv(c.argv), c.env, L.logs / f"{s.name}{'' if len(s.cmds) == 1 else f'_{i}'}.log") |
| if rc == EXIT_SIGNAL: |
| raise KeyboardInterrupt(f"{s.name} stopped by a signal") |
| if rc != 0: |
| raise DriverError(f"[round {s.r}] {s.name}: command {i} exited {rc} (log: {L.logs}); no marker written, re-run resumes it") |
| missing = [str(o) for o in s.outputs if not Path(o).exists()] |
| if missing: |
| raise DriverError(f"[round {s.r}] {s.name}: finished but declared outputs are missing: {missing}") |
| if s.name == "gen_verify": |
| why = generation_incomplete(L.gen / "gen_r0.jsonl", L.sel / "problems.jsonl", cfg.n, cfg.gen_cap) |
| if why: |
| raise DriverError(f"[round {s.r}] gen_verify: generation is not complete, no marker written: {why}") |
| if s.name == "sft_data": |
| rows = json.loads((L.data / "manifest.json").read_text())["rows"]["total"] |
| if rows < cfg.min_rows: |
| raise DataStop(f"round {s.r}: only {rows} SFT rows (< --min-rows {cfg.min_rows}); the selected problems yielded too few accepted proofs") |
| if s.name == "average" and cfg.prune_export and L.export.exists(): |
| shutil.rmtree(L.export) |
| _write_json(marker_path(L, s.name), {"stage": s.name, "round": s.r, "chain": s.chain, "finished": time.strftime("%F %T"), |
| "seconds": round(time.time() - t0, 1), "identity": json.loads(s.identity())}) |
| return "ran" |
|
|
|
|
| def run_all(cfg: Cfg, runner, say=print) -> int: |
| chains = chain_stages(cfg, cfg.rounds) |
| if cfg.redo: |
| apply_redo(cfg, cfg.redo) |
| try: |
| for r in range(cfg.start_round, cfg.rounds + 1): |
| err = split_problems(cfg) |
| if err: |
| say(f"[round {r}] REFUSED: {err}") |
| return 3 |
| for s in chains[r]: |
| run_stage(cfg, s, runner, say) |
| man = json.loads(Layout(cfg, r).manifest.read_text()) |
| dec = man["stop"] |
| say(f"[round {r}] solved@{man['eval']['n_samples_per_problem']} = {man['eval']['solved_at_n']}/{man['eval']['n_problems']} " |
| f"(cap {man['eval']['cap_tokens']}); decision: {dec['reason']}") |
| if dec["stop"]: |
| say(f"STOP after round {r}: {dec['reason']} (best round: {dec['best_round']})") |
| return EXIT_OK |
| except DataStop as e: |
| say(f"STOP (data): {e}") |
| return EXIT_DATA_STOP |
| except KeyboardInterrupt as e: |
| say(f"interrupted: {e}") |
| return EXIT_SIGNAL |
| except (DriverError, C.ContractError) as e: |
| say(f"FAILED: {e}") |
| return EXIT_FAIL |
| return EXIT_OK |
|
|
|
|
| def dry_run(cfg: Cfg, say=print) -> int: |
| chains = chain_stages(cfg, cfg.rounds) |
| say(f"# dry run: nothing is executed or written. run dir {cfg.run_dir}; rounds {cfg.start_round}..{cfg.rounds}") |
| for r in range(cfg.start_round, cfg.rounds + 1): |
| L = Layout(cfg, r) |
| say(f"\n## round {r}: model in {model_in(cfg, r)} -> out {L.model_out}") |
| for s in chains[r]: |
| say(f"[{stage_status(L, s)}] {s.name}") |
| if s.fn is not None: |
| say(f" # in-process: {s.desc}") |
| for c in s.cmds: |
| env = " ".join(f"{k}={shlex.quote(v)}" for k, v in sorted(c.env.items())) |
| say(f" {('env ' + env + ' ') if env else ''}{display_argv(c.argv)}") |
| return EXIT_OK |
|
|
|
|
| |
|
|
|
|
| def scored_splits(path: str) -> set: |
| """The `split` values of the scored benchmark file. FAILS CLOSED: a missing, unreadable or |
| malformed file, an empty file, or any record without a `split` raises ValueError -- an unknown |
| split must never pass as "not test".""" |
| try: |
| with open(path, encoding="utf-8") as f: |
| rows = [json.loads(line) for line in f if line.strip()] |
| except (OSError, ValueError) as e: |
| raise ValueError(f"cannot read the scored benchmark {path}: {e}") from e |
| if not rows: |
| raise ValueError(f"the scored benchmark {path} is empty") |
| missing = sum(1 for r in rows if not isinstance(r, dict) or not r.get("split")) |
| if missing: |
| raise ValueError(f"{missing} record(s) in {path} have no `split`: cannot prove it is not a test split") |
| return {r["split"] for r in rows} |
|
|
|
|
| def split_problems(cfg: Cfg) -> Optional[str]: |
| """Split hygiene; returns an error message or None. The scored benchmark drives the stop rule, the |
| best-round choice and the averaging check, so scoring a TEST split makes every later test number |
| partly a selection artefact. Development runs score miniF2F-valid. --allow-test-selection is an |
| escape hatch for one non-adaptive diagnostic round; the FINAL evaluation does not use this driver at |
| all (docs/SOTA_PLAN.md: frozen checkpoint -> pipeline_gen_verify.py -> eval_report.py). |
| Decontamination: every protected test benchmark must reach build_sft_data's exclusion set.""" |
| try: |
| splits = scored_splits(str(cfg.eval_problems[0])) |
| except ValueError as e: |
| return str(e) |
| if "test" in splits: |
| if not cfg.allow_test_selection: |
| return (f"--eval-problems[0] {cfg.eval_problems[0]} is a TEST split: rounds, the stop rule and averaging would " |
| "be selected on test. Score miniF2F-valid (scripts/p2/build_valid.py)") |
| if cfg.rounds != 1 or cfg.target_solved is not None: |
| return "--allow-test-selection is for ONE non-adaptive round: needs --rounds 1 and no --target-solved" |
| guarded = set() |
| for p in list(cfg.eval_problems) + list(cfg.protected_problems): |
| try: |
| guarded |= scored_splits(str(p)) |
| except ValueError as e: |
| return str(e) |
| if "test" not in guarded: |
| return ("no TEST benchmark reaches the SFT decontamination gate: pass the protected test files " |
| "(miniF2F-test, MOBench) via --protected-problems") |
| return pool_informal_overlap(cfg) |
|
|
|
|
| def pool_informal_overlap(cfg: Cfg) -> Optional[str]: |
| """build_sft_data excludes by exact statement_sha only, so 'the same informal problem, formalised differently' (the |
| realistic leak, build_problems' informal 13-gram / 8-gram layers) must be absent from the pool already. It is applied |
| to the pool against test+MOBench by build_problems and to the dev split by build_valid; this re-checks the files |
| actually handed to this driver, so a stale pool or dev file cannot inflate the dev pass rate that drives the stop rule.""" |
| if not Path(cfg.pool).is_file(): |
| return None |
| spec = importlib.util.spec_from_file_location("p2_build_problems_for_ei", Path(__file__).with_name("build_problems.py")) |
| bp = importlib.util.module_from_spec(spec) |
| sys.modules[spec.name] = bp |
| try: |
| spec.loader.exec_module(bp) |
| finally: |
| sys.modules.pop(spec.name, None) |
| guarded = [] |
| for p in list(cfg.eval_problems) + list(cfg.protected_problems): |
| guarded.extend(C.iter_jsonl(p)) |
| decon = bp.Decontaminator(guarded) |
| hits = [(r["problem_id"], why) for r in C.iter_jsonl(cfg.pool) if (why := decon.check_informal(r.get("informal_prefix"))) is not None] |
| if hits: |
| return (f"{len(hits)} --pool problem(s) share informal statement text with --eval-problems/--protected-problems " |
| f"(first: {hits[:3]}): training on them leaks the dev/test problems. Drop them from the pool, or rebuild the dev " |
| "split with scripts/p2/build_valid.py (it drops valid problems that overlap the pool)") |
| return None |
|
|
|
|
| def test_guard(ap, cfg: Cfg) -> None: |
| err = split_problems(cfg) |
| if err: |
| ap.error(err) |
|
|
|
|
| def parse_args(argv=None) -> Cfg: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| g = ap.add_argument |
| g("--run-dir", required=True) |
| g("--pool", required=True, help="training pool Problem JSONL (split=train)") |
| g("--eval-problems", nargs="+", required=True, help="eval Problem JSONL; the FIRST is the scored benchmark (miniF2F)") |
| g("--base-model", required=True, help="HF snapshot dir (or its snapshots/ parent) of the starting model") |
| g("--rounds", type=int, default=3) |
| g("--start-round", type=int, default=1) |
| g("--name", default="ei") |
| g("--seed-ver", nargs="*", default=[], help="verification JSONL(s) measured at --gen-cap that seed the pass-rate table of round 1") |
| g("--band-lo", type=float, default=0.0) |
| g("--band-hi", type=float, default=0.75) |
| g("--include-zero", action="store_true", help="also resample problems with measured pass rate 0") |
| g("--max-zero", type=int, default=None) |
| g("--max-problems", type=int, default=None) |
| g("--seed", type=int, default=0) |
| g("--n", type=int, default=8, help="samples per selected problem per round") |
| g("--gen-cap", type=int, default=16384) |
| g("--gen-temperature", type=float, default=None, |
| help="lean_gen --temperature for the generation stages (omit: lean_gen's default 1.0, argv unchanged)") |
| g("--tp", type=int, default=4) |
| g("--workers", type=int, default=8) |
| g("--corrections", action="store_true", help="also generate+train one Lean-feedback correction round") |
| g("--max-per-problem", type=int, default=4) |
| g("--max-correction-per-problem", type=int, default=1) |
| g("--rank", choices=("order", "shortest"), default="shortest") |
| g("--keep-easy", action="store_true") |
| g("--fill-duplicates", action="store_true") |
| g("--max-seq-len", type=int, default=18432) |
| g("--min-rows", type=int, default=50) |
| g("--train-config", default="configs/p2_sft.toml") |
| g("--nproc", type=int, default=4) |
| g("--lr", type=float, default=None) |
| g("--epochs", type=int, default=None) |
| g("--lora-r", type=int, default=64) |
| g("--lora-alpha", type=float, default=128.0) |
| g("--avg-weight", type=float, default=None, help="w in out=(1-w)*anchor+w*merged; omit for no averaging") |
| g("--avg-anchor", choices=("original", "start"), default="original") |
| g("--prune-export", action="store_true", help="delete the 16 GB merged export once the averaged model exists") |
| g("--eval-n", type=int, default=8) |
| g("--final-eval-n", type=int, default=None, help="n for the last round only (e.g. 32)") |
| g("--eval-cap", type=int, default=30000) |
| g("--eval-seed", type=int, default=0) |
| g("--baseline-run", default=None, help="NAME:GEN:VER of the baseline eval (same n and cap)") |
| g("--min-gain", type=int, default=2) |
| g("--patience", type=int, default=2) |
| g("--max-regress", type=int, default=3) |
| g("--target-solved", type=int, default=None) |
| g("--allow-test-selection", action="store_true", |
| help="escape hatch: permit a TEST-split scored benchmark for ONE non-adaptive round (--rounds 1, no --target-solved); " |
| "the final evaluation does not use this driver") |
| g("--protected-problems", nargs="*", default=[], |
| help="benchmark JSONL(s) never scored here but excluded from training data (miniF2F-test, MOBench, ...)") |
| g("--redo", action="append", default=[], metavar="ROUND:STAGE") |
| g("--dry-run", action="store_true") |
| a = ap.parse_args(argv) |
| d = vars(a) |
| d["dry"] = d.pop("dry_run") |
| cfg = Cfg(**d) |
| if cfg.avg_weight is not None and not (0.0 < cfg.avg_weight <= 1.0): |
| ap.error("--avg-weight must be in (0, 1]") |
| test_guard(ap, cfg) |
| if cfg.eval_cap < 30000: |
| print(f"note: --eval-cap {cfg.eval_cap} < 30000: not comparable to the published pass@32 protocol", file=sys.stderr) |
| return cfg |
|
|
|
|
| def main(argv=None, runner=None) -> int: |
| cfg = parse_args(argv) |
| if cfg.dry: |
| return dry_run(cfg) |
| return run_all(cfg, runner or SubprocessRunner()) |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
