File size: 41,919 Bytes
63f41a6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 | #!/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/<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."""
# =============================================================================== 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())
|