File size: 6,496 Bytes
bcc8d35 | 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 | """Stage-b reward-path ceiling test: a PERFECT policy must score ~100%.
House rule from a paid-for incident (mistral-prm-ppo, 2026-08-14): two PPO
runs measured a broken answer-extractor -- a perfect policy scored 0/100.
No GRPO rollout launches until this driver certifies the EXACT reward path
(extract_answer -> OutcomeVerifier -> CompositeReward + format instrument)
at >= --floor on every gold source, with every miss listed in the receipt
for eyeballing.
CPU-only, no model: each (problem, gold) becomes
`build_perfect_completion(gold, eos)` -- byte-identical to what a policy
flawlessly imitating its SFT targets emits, EOS included (rollout decodes
with skip_special_tokens=False).
Sources: gsm8k (1,319), math500 (500), deepmath (--deepmath-n sampled
evenly across shards; stage-b's actual gold distribution -- the reason this
must re-run here and not only on the stage-a sets).
python scripts/ceiling_test.py --deepmath-n 5000
"""
import argparse
import json
import os
import sys
from pathlib import Path
# Run from a checkout without an install: put the repo root on sys.path first.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from align.data.schema import extract_answer, register_answer_suffixes
from align.data.sources import DataSources, gsm8k_gold_answer, read_gsm8k_test, read_math500
from align.data.stageb_sources import StageBSources, iter_bigmath, iter_deepmath
from align.eval.roundtrip import build_perfect_completion
from align.eval.verifier import OutcomeVerifier, VerifierConfig
from align.rl.format_instrument import format_scores
from align.rl.rewards import CompositeReward
EOS = "</s>"
#: The corpora live in two hubs: the stage-a pins under sources.DEFAULT_HUB
#: (the voxtral cache this project inherited), the stage-b downloads under
#: the rlhf hf_cache. Resolved per-source-set, not via one env var.
STAGEB_HUB = "/apps/development/rlhf/hf_cache/hub"
def _gold_sources(args) -> dict[str, list[str]]:
out: dict[str, list[str]] = {}
sources = DataSources.from_local_hub() # DEFAULT_HUB: gsm8k/math500 pins
if args.gsm8k:
out["gsm8k"] = [gsm8k_gold_answer(r) for r in read_gsm8k_test(sources)]
if args.math500:
out["math500"] = [r["answer"] for r in read_math500(sources)]
if args.deepmath_n or args.bigmath_n:
sb = StageBSources.from_local_hub(
hub=args.stageb_hub,
need_deepmath=bool(args.deepmath_n),
need_openmath=False,
need_bigmath=bool(args.bigmath_n),
)
for key, it, total, n, field in (
("deepmath", iter_deepmath, 103_110, args.deepmath_n, "final_answer"),
("bigmath", iter_bigmath, 251_122, args.bigmath_n, "answer"),
):
if not n:
continue
golds: list[str] = []
for i, row in enumerate(it(sb)):
# even sample across the stream without materializing the corpus
if i % max(1, total // n) == 0:
golds.append(str(row[field]))
if len(golds) >= n:
break
out[key] = golds
return out
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument("--floor", type=float, default=0.98)
parser.add_argument("--gsm8k", action=argparse.BooleanOptionalAction, default=True)
parser.add_argument("--math500", action=argparse.BooleanOptionalAction, default=True)
parser.add_argument("--deepmath-n", type=int, default=5000)
parser.add_argument("--bigmath-n", type=int, default=5000)
parser.add_argument(
"--workers", type=int, default=4, help="verifier procs; keep <=4 on the shared box"
)
parser.add_argument("--out", default="outputs/eval/ceiling_test.json")
parser.add_argument("--stageb-hub", default=STAGEB_HUB)
args = parser.parse_args()
register_answer_suffixes(EOS) # BEFORE any extraction (golden.py's incident)
verifier = OutcomeVerifier(VerifierConfig(workers=args.workers))
reward = CompositeReward(format_weight=0.1) # overlong off: length is not under test
report: dict = {"floor": args.floor, "eos": EOS, "sources": {}}
overall_pass = True
for name, golds in _gold_sources(args).items():
completions = [build_perfect_completion(g, EOS) for g in golds]
preds = [extract_answer(c) for c in completions]
extract_ok = sum(1 for p in preds if p is not None)
outcomes = verifier.verify_batch(golds, preds)
correct = [o.name == "CORRECT" for o in outcomes]
fscores = [format_scores(c)["format/score"] for c in completions]
rewards = [
reward(accuracy=1.0 if ok else 0.0, format_score=f, response_len=0)["reward/total"]
for ok, f in zip(correct, fscores)
]
ceiling = sum(correct) / len(golds) if golds else float("nan")
misses = [
{"gold": g, "pred": p, "outcome": o.name}
for g, p, o, ok in zip(golds, preds, outcomes, correct)
if not ok
]
passed = ceiling >= args.floor
overall_pass &= passed
report["sources"][name] = {
"n": len(golds),
"extract_ok_rate": extract_ok / len(golds),
"verify_correct_rate": ceiling,
"mean_format_score": sum(fscores) / len(fscores),
"mean_reward_total": sum(rewards) / len(rewards),
"passed": passed,
"misses": misses[:50],
"misses_total": len(misses),
}
print(
f"{name}: n={len(golds)} extract_ok={extract_ok / len(golds):.4f} "
f"ceiling={ceiling:.4f} misses={len(misses)} -> {'PASS' if passed else 'FAIL'}",
flush=True,
)
report["overall_pass"] = bool(overall_pass)
os.makedirs(os.path.dirname(args.out), exist_ok=True)
tmp = args.out + ".partial"
with open(tmp, "w") as f:
json.dump(report, f, indent=1)
os.replace(tmp, args.out)
print(f"receipt: {args.out} | overall_pass={overall_pass}")
if not overall_pass:
print(
"CEILING FAIL: eyeball every miss in the receipt before touching the "
"verifier -- a representational-equivalence gap and a broken extractor "
"look identical in the rate and completely different in the misses.",
file=sys.stderr,
)
return 0 if overall_pass else 1
if __name__ == "__main__":
sys.exit(main())
|