File size: 6,799 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 | #!/usr/bin/env bash
# P2: one full expert-iteration round, as an explicit sequence of RESUMABLE steps.
#
# scripts/p2/run_round.sh <step> [<step>...]
# scripts/p2/run_round.sh all # in order; stops at the first failure
# scripts/p2/run_round.sh print # show every command, run nothing
#
# Steps (each is one process that takes and releases the repo heavy-job lock itself; NEVER run
# two at once, and never wrap the whole script in a lock: the children would deadlock on it):
#
# problems eval sets + decontaminated training pool + elaboration filter
# base_eval baseline: generate (cap 16384, n=8) + verify on miniF2F and MOBench
# pool_gen round-0 generation on the training pool
# pool_verify verify round 0 (with --feedback: correction prompts need the Lean errors)
# corr_gen correction round 1 from the failed round-0 samples
# corr_verify verify round 1
# sft_data build the SFT shard set (exact sampled tokens, completion-only loss)
# train P2-B's trainer (command supplied via $P2_TRAIN_CMD; not defined here)
# export P2-B's HF export of the trained weights (command via $P2_EXPORT_CMD)
# post_eval the same evaluation as base_eval on the trained weights
# report eval_report.py: pass@k, CIs, paired McNemar, correction lift
#
# Required env: RUN (run name); MAX_SEQ_LEN for sft_data. Optional: CAP (16384), N_EVAL (8), N_POOL (8), WORKERS (8),
# TP (4), CAP_REF (40960, set to also run the reference-cap baseline).
# Host rules (/apps/CLAUDE.md) are enforced by the python entry points (preflight + flock +
# <= 8 Lean workers + scratch under /apps); this script only sequences them.
set -euo pipefail
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "$REPO"
: "${RUN:?set RUN=<run name>}"
CAP="${CAP:-16384}"
N_EVAL="${N_EVAL:-8}"
N_POOL="${N_POOL:-8}"
WORKERS="${WORKERS:-8}"
TP="${TP:-4}"
D="$REPO/.tmp/p2/$RUN"
PY="$REPO/.venv/bin/python"
PYR="$REPO/.venv-rollout/bin/python"
export TMPDIR="$D/scratch" HF_HOME="${HF_HOME:-/apps/development/rlhf/hf_cache}" HF_HUB_DISABLE_XET=1
MODEL_BASE="${MODEL_BASE:-/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a}"
MODEL_SFT="${MODEL_SFT:-$D/export}"
PRINT_ONLY=0
run() { # print the command; execute unless printing
echo "+ $*" >&2
[ "$PRINT_ONLY" = 1 ] || "$@"
}
gen() { # gen <model> <problems> <out> <n> <cap> [extra...]
local model="$1" probs="$2" out="$3" n="$4" cap="$5"; shift 5
run "$PYR" extras/vllm_rollout/lean_gen.py --model "$model" --problems "$probs" --out "$out" --n "$n" \
--cap-tokens "$cap" --round 0 --tp "$TP" --seed 0 --wait-lock "$@"
}
verify() { # verify <gens> <problems> <out> [extra...]
local g="$1" p="$2" o="$3"; shift 3
run "$PY" scripts/p2/verify_generations.py --generations "$g" --problems "$p" --out "$o" --workers "$WORKERS" "$@"
}
step_problems() {
mkdir -p "$D/problems" "$TMPDIR"
run "$PY" scripts/p2/build_problems.py --phase eval --out-dir "$D/problems"
run "$PYR" scripts/p2/build_problems.py --phase pool-sample --out-dir "$D/problems" --n 2000 --candidates 5000
run "$PY" scripts/p2/build_problems.py --phase check --out-dir "$D/problems" --n 2000 --workers "$WORKERS"
}
export_rev() { # content id of an exported model (INTEGRATOR FIX: lean_gen refuses a non-snapshot dir without --model-rev)
[ "$PRINT_ONLY" = 1 ] && { echo "<export-rev>"; return; }
echo "sft-$(cd "$1" && cat $(ls model*.safetensors | sort) | sha256sum | cut -c1-16)"
}
eval_one() { # eval_one <tag> <model> <cap>
local tag="$1" model="$2" cap="$3" rev=()
[ "$model" = "$MODEL_BASE" ] || rev=(--model-rev "$(export_rev "$model")")
for set in minif2f mobench; do
gen "$model" "$D/problems/$set.jsonl" "$D/$tag/gen_${set}_c${cap}.jsonl" "$N_EVAL" "$cap" "${rev[@]}"
verify "$D/$tag/gen_${set}_c${cap}.jsonl" "$D/problems/$set.jsonl" "$D/$tag/ver_${set}_c${cap}.jsonl"
done
}
step_base_eval() {
eval_one base "$MODEL_BASE" "$CAP"
[ -z "${CAP_REF:-}" ] || eval_one base "$MODEL_BASE" "$CAP_REF"
}
step_pool_gen() { gen "$MODEL_BASE" "$D/problems/train_pool.jsonl" "$D/pool/gen_r0.jsonl" "$N_POOL" "$CAP"; }
step_pool_verify() { verify "$D/pool/gen_r0.jsonl" "$D/problems/train_pool.jsonl" "$D/pool/ver_r0.jsonl" --feedback; }
step_corr_gen() {
run "$PYR" extras/vllm_rollout/lean_gen.py --model "$MODEL_BASE" --problems "$D/problems/train_pool.jsonl" \
--out "$D/pool/gen_r1.jsonl" --n 1 --cap-tokens "$CAP" --round 1 --tp "$TP" --seed 0 --wait-lock \
--corrections-from "$D/pool/ver_r0.jsonl" "$D/pool/gen_r0.jsonl"
}
step_corr_verify() { verify "$D/pool/gen_r1.jsonl" "$D/problems/train_pool.jsonl" "$D/pool/ver_r1.jsonl"; }
step_sft_data() {
# No default on purpose: the longest row the trainer can hold (prompt + completion; correction
# rows carry the failed attempt in the prompt) is a memory fact P2-B measures, not a constant.
[ "$PRINT_ONLY" = 1 ] || : "${MAX_SEQ_LEN:?set MAX_SEQ_LEN to the longest sequence the trainer fits (P2-B probe)}"
run "$PY" scripts/p2/build_sft_data.py \
--generations "$D/pool/gen_r0.jsonl" "$D/pool/gen_r1.jsonl" \
--verifications "$D/pool/ver_r0.jsonl" "$D/pool/ver_r1.jsonl" \
--problems "$D/problems/train_pool.jsonl" \
--eval-problems "$D/problems/minif2f.jsonl" "$D/problems/mobench.jsonl" \
--cap-tokens "$CAP" --out-dir "$D/sft_data" --max-seq-len "${MAX_SEQ_LEN:-<MAX_SEQ_LEN>}"
}
step_train() {
[ "$PRINT_ONLY" = 1 ] || : "${P2_TRAIN_CMD:?P2-B supplies the training command (reads $D/sft_data/shards)}"
run bash -c "${P2_TRAIN_CMD:-<P2_TRAIN_CMD>}"
}
step_export() {
[ "$PRINT_ONLY" = 1 ] || : "${P2_EXPORT_CMD:?P2-B supplies the HF export command (writes $MODEL_SFT)}"
run bash -c "${P2_EXPORT_CMD:-<P2_EXPORT_CMD>}"
}
step_post_eval() {
eval_one sft "$MODEL_SFT" "$CAP"
[ -z "${CAP_REF:-}" ] || eval_one sft "$MODEL_SFT" "$CAP_REF"
}
step_report() {
local runs=() cap
for tag in base sft; do
local g="" v=""
for cap in "$CAP" ${CAP_REF:-}; do
for set in minif2f mobench; do g+="${g:+,}$D/$tag/gen_${set}_c${cap}.jsonl"; v+="${v:+,}$D/$tag/ver_${set}_c${cap}.jsonl"; done
done
runs+=(--run "$tag:$g:$v")
done
run "$PY" scripts/p2/eval_report.py --problems "$D/problems/minif2f.jsonl" "$D/problems/mobench.jsonl" \
"${runs[@]}" --baseline base --candidate sft --out-dir "$D/report"
}
[ "$#" -ge 1 ] || { sed -n '2,30p' "$0"; exit 2; }
if [ "$1" = print ]; then PRINT_ONLY=1; set -- all; fi
steps=("$@")
[ "${steps[0]}" != all ] || steps=(problems base_eval pool_gen pool_verify corr_gen corr_verify sft_data train export post_eval report)
for s in "${steps[@]}"; do
echo "=== step: $s ($(date '+%F %T'))" >&2
"step_$s"
done
|