| |
| """Single-turn on-policy RL for Lean provers (lane R). Library: align/rl/lean_rl.py, align/rl/lean_rl_run.py. |
| |
| # the iteration loop (CPU-only parent; spawns generation, verification and training children in series) |
| .venv/bin/python scripts/p2/rl_lean.py run configs/p2_rl_lean.toml \\ |
| --set run_dir="RUN" --set problems="POOL.jsonl" [--nproc 4] [--until K] [--dry-run] |
| |
| # one training iteration (what `run` launches under torchrun; also the Phase 0b probe, docs/RL_LEAN.md) |
| torchrun --standalone --nproc_per_node=4 scripts/p2/rl_lean.py train-iter CFG --iteration K [--set ...] |
| # K = -1 writes checkpoint 0 (fresh adapter, fresh optimizer) |
| |
| # build the loss-mode contrast on the SAME cached rollouts (K1): a second run with loss.mode="ei" |
| ... run CFG --set run_dir="RUN_EI" --set replay_from="RUN" --set 'loss.mode="ei"' |
| |
| Resumable: re-running `run` continues from the last committed checkpoint; every stage commits atomically. |
| The parent never takes the heavy-job lock itself: each child does (lean_gen --wait-lock, verify_generations when |
| workers > 2, train-iter on rank 0), so exactly one heavy child exists at any time. |
| |
| IMPORT ORDER IS LOAD-BEARING (as in scripts/sft.py): the guard runs before anything imports torch. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import logging |
| import os |
| import sys |
| from pathlib import Path |
|
|
| REPO = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(REPO)) |
|
|
| from align.utils.guard import apply_resource_guards |
|
|
|
|
| def parse_args(argv=None): |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| for name in ("run", "train-iter"): |
| p = sub.add_parser(name) |
| p.add_argument("config") |
| p.add_argument("--set", dest="overrides", action="append", default=[], metavar="KEY=TOML_VALUE") |
| if name == "run": |
| p.add_argument("--nproc", type=int, default=4, help="training ranks (torchrun --nproc_per_node)") |
| p.add_argument("--until", type=int, default=None, help="stop after this many iterations (resume continues)") |
| p.add_argument("--dry-run", action="store_true", help="print iteration 0's child command lines, touch nothing") |
| else: |
| p.add_argument("--iteration", type=int, required=True, help="-1 = write checkpoint 0") |
| return ap.parse_args(argv) |
|
|
|
|
| def cmd_run(a) -> int: |
| from align.rl import lean_rl as L |
| from align.rl import lean_rl_run as R |
|
|
| cfg = L.load_config(a.config, a.overrides) |
| if not cfg.run_dir or not cfg.problems: |
| raise SystemExit('run_dir and problems must be set (--set \'run_dir="..."\' --set \'problems="..."\')') |
| gen, ver = R.LeanGenCliGenerator(cfg), R.VerifyCliVerifier(cfg) |
| trainer = R.TorchrunTrainer(cfg, os.path.abspath(a.config), a.nproc, a.overrides) |
| if a.dry_run: |
| d = R.iter_dir(cfg.run_dir, 0) |
| print(json.dumps({ |
| "fingerprint": cfg.fingerprint(), |
| "init": R.torchrun_command(trainer.cfg_path, -1, a.nproc, a.overrides), |
| "gen": gen.command(0, d, os.path.join(d, "problems.jsonl"), os.path.join(R.ckpt_path(cfg.run_dir, 0), "adapter"), |
| os.path.join(R.ckpt_path(cfg.run_dir, 0), "merged") if cfg.export_merged else None, |
| R.derive_seed(cfg.seed, 0)), |
| "verify": ver.command(d, os.path.join(d, "problems.jsonl")), |
| "train": R.torchrun_command(trainer.cfg_path, 0, a.nproc, a.overrides)}, indent=1)) |
| return 0 |
| ledger = R.run_loop(cfg, gen, ver, trainer, until=a.until) |
| print(json.dumps({"iterations_committed": len(ledger), "cumulative": ledger[-1]["cumulative"] if ledger else {}})) |
| return 0 |
|
|
|
|
| def cmd_train_iter(a) -> int: |
| from align.rl import lean_rl as L |
|
|
| cfg = L.load_config(a.config, a.overrides) |
| apply_resource_guards(cpu_threads=8, heavy=True, label=f"rl-lean-train-{a.iteration}") |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") |
| from align.utils.dist import init_distributed |
|
|
| init_distributed(run_dir=os.path.join(cfg.run_dir, "nccl")) |
| from align.rl import lean_rl_run as R |
|
|
| trainer = R.InProcessTrainer(cfg, lambda: R.build_real_policy(cfg)) |
| if a.iteration < 0: |
| trainer.init() |
| else: |
| R.clean_partial(cfg.run_dir) if R._rank() == 0 else None |
| trainer.train(a.iteration) |
| return 0 |
|
|
|
|
| def main(argv=None) -> int: |
| a = parse_args(argv) |
| if a.cmd == "run": |
| apply_resource_guards(cpu_threads=2, heavy=False, label="rl-lean-orchestrator") |
| return cmd_run(a) |
| return cmd_train_iter(a) |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
