File size: 8,225 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 | #!/usr/bin/env bash
# P2 integrator: ONE tiny end-to-end loop on the real model, serialized under the heavy lock.
#
# scripts/p2/smoke_loop.sh <step> [<step>...] steps in order: see STEPS below
# scripts/p2/smoke_loop.sh print show every command, run nothing
#
# Purpose: prove the PLUMBING (keys, tokens, files, the export) with numbers that mean nothing:
# 16 short training statements, 8 miniF2F problems, cap 4096. It is NOT the experiment.
#
# Host rules (/apps/CLAUDE.md): every heavy step is preceded by `preflight` (>=45 GB MemAvailable,
# load < 2x cores, disk queue < 64; strict, raises) and by `waitfree` (poll the repo heavy-job flock
# until nobody holds it). The python entry points then take the lock themselves for the whole
# step; this script never holds it (children would deadlock on it). Scratch is under
# .tmp/p2/p2i, TMPDIR too; nothing goes to / or /tmp.
#
# STEPS: problems gen_r0 ver_r0 gen_r1 ver_r1 sft_data seams train_a train_resume export gen_post gen_base ver_eval report cleanup
set -euo pipefail
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "$REPO"
RUN="${RUN:-p2i}"
D="$REPO/.tmp/p2/$RUN"
PY="$REPO/.venv/bin/python"
PYR="$REPO/.venv-rollout/bin/python"
CAP="${CAP:-4096}"
TP="${TP:-4}"
WORKERS="${WORKERS:-8}"
MAX_SEQ_LEN="${MAX_SEQ_LEN:-9216}" # correction rows carry the failed attempt: <= cap + prompt + errors
MODEL_BASE="${MODEL_BASE:-/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a}"
LOCK="$REPO/.align_heavy_job.lock"
export TMPDIR="$D/scratch" HF_HOME="${HF_HOME:-/apps/development/rlhf/hf_cache}" HF_HUB_DISABLE_XET=1
PRINT_ONLY=0
mkdir -p "$TMPDIR" "$D/problems" "$D/logs"
run() { echo "+ $*" >&2; [ "$PRINT_ONLY" = 1 ] || "$@"; }
preflight() { # strict host preflight; prints the numbers it saw
[ "$PRINT_ONLY" = 1 ] && { echo "+ preflight" >&2; return; }
"$PY" -c "from align.utils.guard import preflight_host_check as p; import json; print('preflight', json.dumps(p(strict=True)))" >&2
}
waitfree() { # block until the heavy-job lock is free (probe + release; the job takes it itself)
[ "$PRINT_ONLY" = 1 ] && { echo "+ waitfree" >&2; return; }
local t0=$SECONDS
while ! flock -n "$LOCK" true; do
echo "waiting for heavy lock held by: $(cat "$LOCK" 2>/dev/null)" >&2; sleep 60
done
echo "lock free after $((SECONDS - t0)) s" >&2
}
heavy() { waitfree; preflight; run "$@"; }
heavy_log() { # heavy_log <logfile> cmd... : as `heavy`, output also appended to the log (pipefail keeps the exit code)
local log="$1"; shift; waitfree; preflight; echo "+ $*" >&2
[ "$PRINT_ONLY" = 1 ] || "$@" 2>&1 | tee -a "$log"
}
step_problems() {
# Real eval set + a 16-problem slice of the decontaminated 2,000-problem pool (the 16 SHORTEST
# statements: this is a plumbing test and wants some accepted proofs), 8 miniF2F problems.
[ "$PRINT_ONLY" = 1 ] && { echo "+ (slice problems -> $D/problems)" >&2; return; }
"$PY" - "$REPO/.tmp/p2/problems" "$D/problems" <<'PYEOF'
import json, sys
src, dst = sys.argv[1:3]
tr = [json.loads(l) for l in open(f"{src}/train_pool.jsonl")]
mf = [json.loads(l) for l in open(f"{src}/minif2f.jsonl")]
tr.sort(key=lambda r: (len(r["formal_statement"]), r["problem_id"]))
w = lambda p, rows: open(p, "w").write("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows))
w(f"{dst}/train16.jsonl", tr[:16]); w(f"{dst}/minif2f8.jsonl", mf[:8]); w(f"{dst}/minif2f.jsonl", mf)
w(f"{dst}/mobench.jsonl", [json.loads(l) for l in open(f"{src}/mobench.jsonl")])
PYEOF
}
gen() { # gen <model> <problems> <out> <n> <cap> [extra...]
local model="$1" probs="$2" out="$3" n="$4" cap="$5"; shift 5
heavy "$PYR" extras/vllm_rollout/lean_gen.py --model "$model" --problems "$probs" --out "$out" --n "$n" \
--cap-tokens "$cap" --tp "$TP" --seed 0 --wait-lock "$@"
}
verify() { # verify <gens> <problems> <out> [extra...]
local g="$1" p="$2" o="$3"; shift 3
heavy "$PY" scripts/p2/verify_generations.py --generations "$g" --problems "$p" --out "$o" --workers "$WORKERS" "$@"
}
step_gen_r0() { gen "$MODEL_BASE" "$D/problems/train16.jsonl" "$D/pool/gen_r0.jsonl" 4 "$CAP" --round 0; }
step_ver_r0() { verify "$D/pool/gen_r0.jsonl" "$D/problems/train16.jsonl" "$D/pool/ver_r0.jsonl" --feedback; }
step_gen_r1() { gen "$MODEL_BASE" "$D/problems/train16.jsonl" "$D/pool/gen_r1.jsonl" "${N_CORR:-1}" "$CAP" --round 1 \
--corrections-from "$D/pool/ver_r0.jsonl" "$D/pool/gen_r0.jsonl"; }
step_ver_r1() { verify "$D/pool/gen_r1.jsonl" "$D/problems/train16.jsonl" "$D/pool/ver_r1.jsonl"; }
step_sft_data() {
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/train16.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" --held-frac 0 ${SFT_EXTRA:-}
}
step_seams() {
run "$PY" scripts/p2/check_seams.py keys --generations "$D/pool/gen_r0.jsonl" --verifications "$D/pool/ver_r0.jsonl"
[ ! -s "$D/pool/gen_r1.jsonl" ] || run "$PY" scripts/p2/check_seams.py keys --generations "$D/pool/gen_r1.jsonl" --verifications "$D/pool/ver_r1.jsonl"
run "$PY" scripts/p2/check_seams.py sft --generations "$D/pool/gen_r0.jsonl" "$D/pool/gen_r1.jsonl" \
--verifications "$D/pool/ver_r0.jsonl" "$D/pool/ver_r1.jsonl" --shards "$D/sft_data/shards"
}
# ---- training: 4 ranks, LoRA, configs/p2_sft.toml + overrides --------------------------------
TRAIN_SEQ_LEN="${TRAIN_SEQ_LEN:-9216}"; TRAIN_EPOCHS="${TRAIN_EPOCHS:-1}"; SAVE_EVERY="${SAVE_EVERY:-20}"
train_cmd() {
echo "env PYTHONPATH=$REPO $PY -m torch.distributed.run --standalone --nproc_per_node=4 scripts/sft.py configs/p2_sft.toml" \
"--data.shard_dir=$D/sft_data/shards --data.manifest_path=$D/sft_data/manifest.json" \
"--paths.checkpoint_dir=$D/sft --train.seq_len=$TRAIN_SEQ_LEN --train.epochs=$TRAIN_EPOCHS" \
"--train.save_every_micro_steps=$SAVE_EVERY"
}
step_train_a() { heavy_log "$D/logs/train.log" bash -c "$(train_cmd)"; } # SIGTERM'd by the operator, see docs/P2.md
step_train_resume() { heavy_log "$D/logs/train_resume.log" bash -c "$(train_cmd)"; }
step_export() { # merge adapter + HF export + vLLM reload + first-token logprob seam (keeps the export)
local ck="<latest micro_N>"
if [ "$PRINT_ONLY" != 1 ]; then ck="$(ls -d "$D"/sft/*/checkpoints/micro_???????? 2>/dev/null | sort -V | tail -1)"; : "${ck:?no checkpoint under $D/sft}"; fi
heavy "$PY" scripts/qwen3_export_check.py run --ckpt "$ck" --out-dir "$D/export" --work "$D/export_check" --r 64 --alpha 128 --keep
}
export_rev() { # content id of an exported model: lean_gen cannot infer a revision from a non-snapshot dir (it refuses)
[ "$PRINT_ONLY" = 1 ] && { echo "<export-rev>"; return; }
echo "sft-$(cd "$1" && cat $(ls model*.safetensors | sort) | sha256sum | cut -c1-16)"
}
step_gen_post() { gen "$D/export" "$D/problems/minif2f8.jsonl" "$D/sft/gen_minif2f8.jsonl" 2 "$CAP" --round 0 --model-rev "$(export_rev "$D/export")"; }
step_gen_base() { gen "$MODEL_BASE" "$D/problems/minif2f8.jsonl" "$D/base/gen_minif2f8.jsonl" 2 "$CAP" --round 0; }
step_ver_eval() {
verify "$D/base/gen_minif2f8.jsonl" "$D/problems/minif2f8.jsonl" "$D/base/ver_minif2f8.jsonl"
verify "$D/sft/gen_minif2f8.jsonl" "$D/problems/minif2f8.jsonl" "$D/sft/ver_minif2f8.jsonl"
}
step_report() {
run "$PY" scripts/p2/eval_report.py --problems "$D/problems/minif2f8.jsonl" \
--run "base:$D/base/gen_minif2f8.jsonl:$D/base/ver_minif2f8.jsonl" \
--run "sft:$D/sft/gen_minif2f8.jsonl:$D/sft/ver_minif2f8.jsonl" \
--baseline base --candidate sft --out-dir "$D/report" --ks 1 2
}
step_cleanup() { run rm -rf "$D/export" "$D/scratch"; }
[ "$#" -ge 1 ] || { sed -n '2,10p' "$0"; exit 2; }
if [ "$1" = print ]; then PRINT_ONLY=1; set -- problems gen_r0 ver_r0 gen_r1 ver_r1 sft_data seams train_a train_resume export gen_post gen_base ver_eval report cleanup; fi
for s in "$@"; do echo "=== step: $s ($(date '+%F %T'))" >&2; "step_$s"; done
|