code/scripts/eval_checkpoint.py · viai957/align at main

align / code /scripts /eval_checkpoint.py

Add files using upload-large-folder tool

bcc8d35 verified 15 days ago

8.24 kB

"""Standalone checkpoint eval: fresh process, single GPU, golden-proven generator.
Promotion of the .tmp/eval_gate.py driver (task #26) into a parameterized
entry point. The design facts it carries forward are all incident-derived:
  * FRESH PROCESS, never in-trainer: the train->eval handoff inside one
    process produced the rank race (peers built a second 7B while rank 0
    saved) and the orphan-GPU-process incident. A checkpoint on disk plus a
    fresh process has neither failure mode.
  * EVERYTHING under the __main__ guard: the verifier's scoring pool uses
    the 'spawn' context, which RE-IMPORTS __main__ in every worker -- the
    unguarded first driver had all 8 workers load their own 7B onto the
    eval GPU.
  * greedy generation via `golden_generate_fn` (in-repo Engine, same encode
    as the golden regression: BOS default, truncation 768). The anchors and
    banked SFT rows were measured with HF generate; the generator swap is
    certified by the golden regression + checkpoint roundtrip gate, and the
    generator is recorded in the report config.
  * Atomic report write (.partial -> os.replace): a half-written report
    masquerades as a complete one.
Phase-0 usage (the Arm B result side of the primary McNemar):
    CUDA_VISIBLE_DEVICES=0 python scripts/eval_checkpoint.py \
        --checkpoint runs/dpo_arm_b_response/checkpoints/micro_XXXXXXXX \
        --name arm-b-final --gsm8k-n 1319 --math500-n 0
`--quality-gate` additionally runs the SFT-stage quality gate + ledger
records; leave it off for non-SFT checkpoints (the gate's anchors are
SFT-specific).
"""

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.utils.guard import apply_resource_guards


def quality_gate_size_errors(cfg, math500_n: int, gsm8k_n: int, per_set: dict | None = None) -> list[str]:
    """Why a run of these sizes may NOT write the SFT quality-gate receipt (empty = fine).
    The receipt unlocks DPO. `assert_meets_golden` accepts any point estimate inside the target's own
    Wilson CI, which is enormous at tiny n (1/5 passes the 0.09 anchor), so the gate certifies only a
    FULL-size evaluation: both sets, at the configured n the anchors were recorded at. `per_set` (the
    measured report) is checked too, in case a set held fewer problems than requested."""
    want = {"math500": cfg.eval.math500_n, "gsm8k": cfg.eval.gsm8k_n}
    asked = {"math500": math500_n, "gsm8k": gsm8k_n}
    errs = [f"{k}: requested n={asked[k]}, the quality gate needs the full configured n={want[k]}"
            for k in want if asked[k] != want[k]]
    for k, w in want.items():
        got = (per_set or {}).get(k, {}).get("n") if per_set is not None else w
        if got != w:
            errs.append(f"{k}: measured n={got}, the quality gate needs n={w}")
    return errs


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--checkpoint", required=True, help="checkpoint dir (HF-layout safetensors dir)")
    parser.add_argument("--config", default="configs/sft.toml")
    parser.add_argument("--name", required=True, help="report model_name and output stem")
    parser.add_argument("--out", default=None, help="default: outputs/eval/<name>.json")
    parser.add_argument("--math500-n", type=int, default=None, help="0 skips; default: config")
    parser.add_argument("--gsm8k-n", type=int, default=None, help="0 skips; default: config")
    parser.add_argument(
        "--quality-gate", action="store_true", help="SFT-stage gate + ledger records"
    )
    args = parser.parse_args()

    # Guards BEFORE torch import (env vars are load-bearing); heavy=True takes
    # the exclusive job lock -- a 7B generation eval over 1,319 problems is a
    # multi-hour GPU job and must serialise behind training like any other.
    apply_resource_guards(cpu_threads=8, heavy=True, label=f"eval-{args.name}")


    from align.config import load_config  # noqa: E402
    from align.eval.golden import golden_generate_fn, load_checkpoint_model  # noqa: E402
    from align.eval.harness import EvalConfig, load_eval_set, run_eval  # noqa: E402
    from align.eval.verifier import OutcomeVerifier  # noqa: E402
    from align.tokenizer import load_tokenizer  # noqa: E402
    from align.train.stage_smoke import _manifest_hashes, _mde  # noqa: E402

    cfg = load_config(args.config)
    checkpoint = args.checkpoint
    assert os.path.isdir(checkpoint), checkpoint
    out_path = args.out or os.path.join("outputs", "eval", f"{args.name}.json")

    math500_n = cfg.eval.math500_n if args.math500_n is None else args.math500_n
    gsm8k_n = cfg.eval.gsm8k_n if args.gsm8k_n is None else args.gsm8k_n
    if args.quality_gate:
        errs = quality_gate_size_errors(cfg, math500_n, gsm8k_n)
        if errs:
            print("refusing --quality-gate on a subset eval (the receipt unlocks DPO):\n  " + "\n  ".join(errs),
                  file=sys.stderr)
            return 2
    problems = []
    # 0 in this CLI means SKIP the set (checked here); upstream load_eval_set
    # uses limit=0 for "no limit", which this branch structure never passes
    # ambiguously -- a negative value falls through to upstream validation.
    if math500_n != 0:
        problems += load_eval_set(cfg.paths.math500_path, limit=math500_n)
    if gsm8k_n != 0:
        problems += load_eval_set(cfg.paths.gsm8k_path, limit=gsm8k_n)
    if not problems:
        print("nothing to evaluate (both sets skipped)", file=sys.stderr)
        return 2

    # After every refusal above: a 7B load is the expensive step.
    tokenizer = load_tokenizer(checkpoint)
    model = load_checkpoint_model(checkpoint, device="cuda")

    verifier = OutcomeVerifier()
    gen_fn = golden_generate_fn(model, tokenizer, "cuda")
    print(f"evaluating {args.name}: {len(problems)} problems from {checkpoint}", flush=True)
    rep = run_eval(problems, gen_fn, verifier, EvalConfig(), model_name=args.name)
    per_set = rep.per_set
    for name, scores in sorted(per_set.items()):
        print(f"  {name}: n={scores.get('n')} pass@1={scores.get('pass@1')}", flush=True)

    report = {
        "per_set": per_set,
        "config": {
            "mde": _mde(cfg),
            "n_problems": sum(r.get("n", 0) for r in per_set.values()),
            "eval_set_hash": _manifest_hashes(cfg),
            "generator": "align Engine greedy (golden_generate_fn)",
            "checkpoint": checkpoint,
        },
        "pooled": rep.pooled,
        "per_problem": rep.per_problem,
    }
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    tmp = out_path + ".partial"
    with open(tmp, "w") as f:
        json.dump(report, f, indent=1, default=str)
    os.replace(tmp, out_path)
    print(f"report written: {out_path}", flush=True)

    if not args.quality_gate:
        return 0

    from align.train.stage_gates import (  # noqa: E402
        GateLedger,
        eval_source_hash,
        quality_gate,
        write_quality_gate_receipt,
    )

    errs = quality_gate_size_errors(cfg, math500_n, gsm8k_n, per_set=per_set)
    if errs:  # report is already written above; only the receipt is withheld
        print("refusing to write the quality-gate receipt:\n  " + "\n  ".join(errs), file=sys.stderr)
        return 2

    run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft")
    ledger = GateLedger(run_dir, eval_source_hash=eval_source_hash())
    ledger.record("eval", passed=True, detail=json.dumps(per_set, default=str))
    decision = quality_gate(report, cfg)
    ledger.record("quality_gate", passed=decision.passed, detail=decision.halt_message[:400])
    write_quality_gate_receipt(cfg, decision, checkpoint=checkpoint)
    print(f"QUALITY GATE: passed={decision.passed}", flush=True)
    if not decision.passed:
        print(decision.halt_message, file=sys.stderr, flush=True)
    return 0 if decision.passed else 1


if __name__ == "__main__":
    # See module docstring: spawn workers re-import __main__; nothing above
    # this guard may build a model or touch the GPU.
    sys.exit(main())

	"""Standalone checkpoint eval: fresh process, single GPU, golden-proven generator.

	Promotion of the .tmp/eval_gate.py driver (task #26) into a parameterized
	entry point. The design facts it carries forward are all incident-derived:

	* FRESH PROCESS, never in-trainer: the train->eval handoff inside one
	process produced the rank race (peers built a second 7B while rank 0
	saved) and the orphan-GPU-process incident. A checkpoint on disk plus a
	fresh process has neither failure mode.
	* EVERYTHING under the __main__ guard: the verifier's scoring pool uses
	the 'spawn' context, which RE-IMPORTS __main__ in every worker -- the
	unguarded first driver had all 8 workers load their own 7B onto the
	eval GPU.
	* greedy generation via `golden_generate_fn` (in-repo Engine, same encode
	as the golden regression: BOS default, truncation 768). The anchors and
	banked SFT rows were measured with HF generate; the generator swap is
	certified by the golden regression + checkpoint roundtrip gate, and the
	generator is recorded in the report config.
	* Atomic report write (.partial -> os.replace): a half-written report
	masquerades as a complete one.

	Phase-0 usage (the Arm B result side of the primary McNemar):

	CUDA_VISIBLE_DEVICES=0 python scripts/eval_checkpoint.py \
	--checkpoint runs/dpo_arm_b_response/checkpoints/micro_XXXXXXXX \
	--name arm-b-final --gsm8k-n 1319 --math500-n 0

	`--quality-gate` additionally runs the SFT-stage quality gate + ledger
	records; leave it off for non-SFT checkpoints (the gate's anchors are
	SFT-specific).
	"""

	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.utils.guard import apply_resource_guards


	def quality_gate_size_errors(cfg, math500_n: int, gsm8k_n: int, per_set: dict \| None = None) -> list[str]:
	"""Why a run of these sizes may NOT write the SFT quality-gate receipt (empty = fine).

	The receipt unlocks DPO. `assert_meets_golden` accepts any point estimate inside the target's own
	Wilson CI, which is enormous at tiny n (1/5 passes the 0.09 anchor), so the gate certifies only a
	FULL-size evaluation: both sets, at the configured n the anchors were recorded at. `per_set` (the
	measured report) is checked too, in case a set held fewer problems than requested."""
	want = {"math500": cfg.eval.math500_n, "gsm8k": cfg.eval.gsm8k_n}
	asked = {"math500": math500_n, "gsm8k": gsm8k_n}
	errs = [f"{k}: requested n={asked[k]}, the quality gate needs the full configured n={want[k]}"
	for k in want if asked[k] != want[k]]
	for k, w in want.items():
	got = (per_set or {}).get(k, {}).get("n") if per_set is not None else w
	if got != w:
	errs.append(f"{k}: measured n={got}, the quality gate needs n={w}")
	return errs


	def main() -> int:
	parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
	parser.add_argument("--checkpoint", required=True, help="checkpoint dir (HF-layout safetensors dir)")
	parser.add_argument("--config", default="configs/sft.toml")
	parser.add_argument("--name", required=True, help="report model_name and output stem")
	parser.add_argument("--out", default=None, help="default: outputs/eval/<name>.json")
	parser.add_argument("--math500-n", type=int, default=None, help="0 skips; default: config")
	parser.add_argument("--gsm8k-n", type=int, default=None, help="0 skips; default: config")
	parser.add_argument(
	"--quality-gate", action="store_true", help="SFT-stage gate + ledger records"
	)
	args = parser.parse_args()

	# Guards BEFORE torch import (env vars are load-bearing); heavy=True takes
	# the exclusive job lock -- a 7B generation eval over 1,319 problems is a
	# multi-hour GPU job and must serialise behind training like any other.
	apply_resource_guards(cpu_threads=8, heavy=True, label=f"eval-{args.name}")


	from align.config import load_config # noqa: E402
	from align.eval.golden import golden_generate_fn, load_checkpoint_model # noqa: E402
	from align.eval.harness import EvalConfig, load_eval_set, run_eval # noqa: E402
	from align.eval.verifier import OutcomeVerifier # noqa: E402
	from align.tokenizer import load_tokenizer # noqa: E402
	from align.train.stage_smoke import _manifest_hashes, _mde # noqa: E402

	cfg = load_config(args.config)
	checkpoint = args.checkpoint
	assert os.path.isdir(checkpoint), checkpoint
	out_path = args.out or os.path.join("outputs", "eval", f"{args.name}.json")

	math500_n = cfg.eval.math500_n if args.math500_n is None else args.math500_n
	gsm8k_n = cfg.eval.gsm8k_n if args.gsm8k_n is None else args.gsm8k_n
	if args.quality_gate:
	errs = quality_gate_size_errors(cfg, math500_n, gsm8k_n)
	if errs:
	print("refusing --quality-gate on a subset eval (the receipt unlocks DPO):\n " + "\n ".join(errs),
	file=sys.stderr)
	return 2
	problems = []
	# 0 in this CLI means SKIP the set (checked here); upstream load_eval_set
	# uses limit=0 for "no limit", which this branch structure never passes
	# ambiguously -- a negative value falls through to upstream validation.
	if math500_n != 0:
	problems += load_eval_set(cfg.paths.math500_path, limit=math500_n)
	if gsm8k_n != 0:
	problems += load_eval_set(cfg.paths.gsm8k_path, limit=gsm8k_n)
	if not problems:
	print("nothing to evaluate (both sets skipped)", file=sys.stderr)
	return 2

	# After every refusal above: a 7B load is the expensive step.
	tokenizer = load_tokenizer(checkpoint)
	model = load_checkpoint_model(checkpoint, device="cuda")

	verifier = OutcomeVerifier()
	gen_fn = golden_generate_fn(model, tokenizer, "cuda")
	print(f"evaluating {args.name}: {len(problems)} problems from {checkpoint}", flush=True)
	rep = run_eval(problems, gen_fn, verifier, EvalConfig(), model_name=args.name)
	per_set = rep.per_set
	for name, scores in sorted(per_set.items()):
	print(f" {name}: n={scores.get('n')} pass@1={scores.get('pass@1')}", flush=True)

	report = {
	"per_set": per_set,
	"config": {
	"mde": _mde(cfg),
	"n_problems": sum(r.get("n", 0) for r in per_set.values()),
	"eval_set_hash": _manifest_hashes(cfg),
	"generator": "align Engine greedy (golden_generate_fn)",
	"checkpoint": checkpoint,
	},
	"pooled": rep.pooled,
	"per_problem": rep.per_problem,
	}
	os.makedirs(os.path.dirname(out_path), exist_ok=True)
	tmp = out_path + ".partial"
	with open(tmp, "w") as f:
	json.dump(report, f, indent=1, default=str)
	os.replace(tmp, out_path)
	print(f"report written: {out_path}", flush=True)

	if not args.quality_gate:
	return 0

	from align.train.stage_gates import ( # noqa: E402
	GateLedger,
	eval_source_hash,
	quality_gate,
	write_quality_gate_receipt,
	)

	errs = quality_gate_size_errors(cfg, math500_n, gsm8k_n, per_set=per_set)
	if errs: # report is already written above; only the receipt is withheld
	print("refusing to write the quality-gate receipt:\n " + "\n ".join(errs), file=sys.stderr)
	return 2

	run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft")
	ledger = GateLedger(run_dir, eval_source_hash=eval_source_hash())
	ledger.record("eval", passed=True, detail=json.dumps(per_set, default=str))
	decision = quality_gate(report, cfg)
	ledger.record("quality_gate", passed=decision.passed, detail=decision.halt_message[:400])
	write_quality_gate_receipt(cfg, decision, checkpoint=checkpoint)
	print(f"QUALITY GATE: passed={decision.passed}", flush=True)
	if not decision.passed:
	print(decision.halt_message, file=sys.stderr, flush=True)
	return 0 if decision.passed else 1


	if __name__ == "__main__":
	# See module docstring: spawn workers re-import __main__; nothing above
	# this guard may build a model or touch the GPU.
	sys.exit(main())