| |
| """P2 integrator: assert the seams between the independently built pieces, on real files. |
| |
| .venv/bin/python scripts/p2/check_seams.py keys --generations G.jsonl --verifications V.jsonl |
| .venv/bin/python scripts/p2/check_seams.py sft --generations G0 G1 --verifications V0 V1 --shards D/shards |
| |
| Each check exists because the pieces were written by different hands against a prose contract, |
| and each failure below is silent when only one side is looked at: |
| |
| keys generation -> verification. A verification file whose key set differs from the |
| generation file's (a sample never verified, a verdict for a sample that is not on disk, a |
| verdict at another cap or round) scores a different experiment than the one generated. |
| Also checks `code_sha256` against the generation on disk (a stale verdict). |
| sft verification -> trainer input. The tokens the trainer takes loss on must be EXACTLY the |
| `completion_ids` vLLM sampled for a sample the verifier ACCEPTED, and the context must be |
| exactly its `prompt_ids`. This re-derives the labels through the same `ShardDataset` + |
| `collate_sft` the trainer uses and does not trust `build_sft_data.py`'s own read-back |
| (an independent implementation of the same property is the only thing that catches a |
| shared misunderstanding). |
| |
| Exit status 0 only if every check passes; a report JSON is printed either way. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib.util |
| import json |
| import sys |
| from pathlib import Path |
|
|
| REPO = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(REPO)) |
|
|
|
|
| def _common(): |
| if "p2_common" in sys.modules: |
| return sys.modules["p2_common"] |
| spec = importlib.util.spec_from_file_location("p2_common", Path(__file__).with_name("p2_common.py")) |
| mod = importlib.util.module_from_spec(spec) |
| sys.modules["p2_common"] = mod |
| spec.loader.exec_module(mod) |
| return mod |
|
|
|
|
| class SeamError(AssertionError): |
| pass |
|
|
|
|
| def _load(paths, key_fields): |
| out: dict[tuple, dict] = {} |
| for p in paths: |
| for rec in _common().iter_jsonl(p): |
| k = tuple(rec[f] for f in key_fields) |
| if k in out: |
| raise SeamError(f"duplicate key {k} in {p}") |
| out[k] = rec |
| return out |
|
|
|
|
| def check_keys(gen_paths, ver_paths) -> dict: |
| """Generation key set == verification key set, and every verdict is for the code on disk.""" |
| C = _common() |
| gens = _load(gen_paths, C.KEY_FIELDS) |
| vers = _load(ver_paths, C.KEY_FIELDS) |
| missing = sorted(set(gens) - set(vers)) |
| extra = sorted(set(vers) - set(gens)) |
| if missing or extra: |
| raise SeamError(f"key sets differ: {len(missing)} generated but unverified (e.g. {missing[:2]}), " |
| f"{len(extra)} verified but not generated (e.g. {extra[:2]})") |
| stale = [k for k, g in gens.items() if vers[k]["code_sha256"] != C.code_sha(g["lean_code"])] |
| if stale: |
| raise SeamError(f"{len(stale)} verdicts were computed for different code than the generation holds, e.g. {stale[:2]}") |
| |
| bad = [k for k, g in gens.items() if g["finish_reason"] == "length" and vers[k]["accepted"]] |
| if bad: |
| raise SeamError(f"{len(bad)} length-truncated samples are marked accepted, e.g. {bad[:2]}") |
| rounds: dict = {} |
| for k, v in vers.items(): |
| r = rounds.setdefault(k[2], [0, 0]) |
| r[0] += 1 |
| r[1] += int(bool(v["accepted"])) |
| return {"keys": len(gens), "by_round": {str(r): {"samples": a, "accepted": b} for r, (a, b) in sorted(rounds.items())}} |
|
|
|
|
| def check_sft(gen_paths, ver_paths, shards_dir) -> dict: |
| """Every stored row's loss tokens == completion_ids of an ACCEPTED generation, context == prompt_ids.""" |
| import numpy as np |
|
|
| from align.data.loader import collate_sft |
| from align.data.store import ShardDataset |
|
|
| C = _common() |
| gens = _load(gen_paths, C.KEY_FIELDS) |
| vers = _load(ver_paths, C.KEY_FIELDS) |
| rows = 0 |
| kinds: dict = {} |
| seen: set = set() |
| for split in ("train", "held"): |
| d = Path(shards_dir) / "sft" / split |
| if not d.is_dir(): |
| continue |
| ds = ShardDataset(d) |
| for i in range(len(ds)): |
| row = ds[i] |
| ex = row.extras |
| key = (row.problem_id, int(ex["sample_idx"]), int(ex["round"]), int(ex["cap_tokens"])) |
| if key not in gens: |
| raise SeamError(f"{split}[{i}] {key}: no such generation") |
| if not (key in vers and vers[key]["accepted"]): |
| raise SeamError(f"{split}[{i}] {key}: trained on a sample the verifier did not accept") |
| if key in seen: |
| raise SeamError(f"{key}: stored twice") |
| seen.add(key) |
| g = gens[key] |
| p = int(ex["prompt_len"]) |
| ids = np.asarray(row.input_ids) |
| if p != len(g["prompt_ids"]) or not np.array_equal(ids[:p], np.asarray(g["prompt_ids"])): |
| raise SeamError(f"{key}: stored context is not the sampled prompt_ids") |
| lab = collate_sft([row], pad_token_id=0).labels[0].numpy() |
| trained = lab[lab != -100] |
| if not np.array_equal(trained, np.asarray(g["completion_ids"], dtype=trained.dtype)): |
| raise SeamError(f"{key}: tokens with loss ({len(trained)}) != completion_ids ({len(g['completion_ids'])})") |
| if not np.all(lab[:p] == -100): |
| raise SeamError(f"{key}: loss leaks into the prompt") |
| kinds[str(ex["kind"])] = kinds.get(str(ex["kind"]), 0) + 1 |
| rows += 1 |
| if rows == 0: |
| raise SeamError("no SFT rows were found; a vacuous pass") |
| accepted = sum(1 for v in vers.values() if v["accepted"]) |
| return {"rows": rows, "by_kind": kinds, "accepted_samples": accepted, "accepted_not_used": accepted - rows} |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| for name in ("keys", "sft"): |
| s = sub.add_parser(name) |
| s.add_argument("--generations", nargs="+", required=True) |
| s.add_argument("--verifications", nargs="+", required=True) |
| if name == "sft": |
| s.add_argument("--shards", required=True) |
| a = ap.parse_args(argv) |
| try: |
| rep = check_keys(a.generations, a.verifications) if a.cmd == "keys" else check_sft(a.generations, a.verifications, a.shards) |
| rep["passed"] = True |
| except SeamError as e: |
| rep = {"passed": False, "error": str(e)} |
| print(json.dumps(rep, indent=2)) |
| return 0 if rep["passed"] else 1 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
