| """Reference-only pilot evaluation using G2's top20 KL and fixed K2 trajectory. |
| |
| This is the pre-registered matched pilot, not dev checkpoint selection or the |
| final G2 test attempt. Outputs contain aggregates and identities, never rows. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| from collections import Counter |
| import json |
| import statistics |
| from pathlib import Path |
| from typing import NotRequired, TypedDict, cast |
|
|
| import torch |
|
|
| from .data import D2Blocks, sha256 |
| from .evaluate import bucket_kl, match_length |
| from .loss import distillation_loss |
| from .model import Drafter, load |
| from .pilot import ARMS, SEED, Recipe, Spec, digest, initialization, read_spec, schedule, write_json, verify_units |
| from .perf import Targets |
| from .replay import apply_trainables, read_checkpoint |
| from .pilot_statistics import DEV_SPLITS, RequestTotal, split_roles |
| from .pilot_sample import SPLITS as SPLITS, RoundKey, Sample, complete_rounds, draw_sample, key_hash, population_splits |
|
|
| EVALUATOR = "G2-reference-top20-tail-fixed-K2-amend6ab-v2" |
| PART_SCHEMA = "dsv41.d2.pilot-evaluation-part.v1" |
|
|
|
|
| class Metrics(TypedDict): |
| requests: int |
| rounds: int |
| population_sha256: str |
| kl_by_position: list[float] |
| loss: float |
| accepted_length: float |
| kl_median_by_position: list[float] |
| request_totals: list[RequestTotal] |
|
|
|
|
| class Checkpoint(TypedDict): |
| path: str |
| sha256: str |
| step: int |
|
|
|
|
| class Arm(TypedDict): |
| schema: str |
| arm: str |
| plan_sha256: str |
| batch_plan_identity: str |
| checkpoints: dict[str, Checkpoint] |
| step_count: int |
| peak_reserved_bytes: int |
| complete: bool |
| seconds: float |
| synthetic: bool |
| recipe: NotRequired[Recipe] |
|
|
|
|
| class Evaluation(TypedDict): |
| schema: str |
| plan_sha256: str |
| synthetic: bool |
| evaluator: str |
| k2_sha256: dict[str, str] |
| baseline: dict[str, Metrics] |
| arms: dict[str, dict[str, dict[str, Metrics]]] |
| checkpoints: dict[str, dict[str, Checkpoint]] |
| step_agreement: dict[str, object] |
| sample_sha256: str |
| split: NotRequired[str] |
| split_roles: NotRequired[dict[str, str]] |
|
|
|
|
| def read_arm(root: Path, name: str, spec: Spec, plan: Path | str, length: int) -> Arm: |
| arm = cast(Arm, json.loads((root/name/"arm.json").read_text())) |
| batch, _, _ = schedule(length, spec["checkpoints"], name, spec.get("recipe", "kl")) |
| identity = digest({"plan": spec, "arm": name, "schedule": batch.identity}) |
| if (arm["schema"] != "dsv41.d2.pilot-arm.v1" or arm["arm"] != name or not arm["complete"] |
| or arm["plan_sha256"] != (sha256(plan) if isinstance(plan, Path) else plan) or arm["synthetic"] != spec["synthetic"] |
| or arm["batch_plan_identity"] != identity or arm["step_count"] != len(batch.batches) |
| or arm.get("recipe", None if "recipe" in spec else "kl") != spec.get("recipe", "kl") |
| or arm.get("units") != spec["matched_units"][name] |
| or set(arm["checkpoints"]) != {str(n) for n in spec["checkpoints"]}): |
| raise ValueError("pilot arm receipt mismatch or missing checkpoint") |
| for n, step in zip(spec["checkpoints"], batch.checkpoints): |
| record = arm["checkpoints"][str(n)] |
| state, actual = read_checkpoint(Path(record["path"])) |
| if (actual != record["sha256"] or record["step"] != step or state["step"] != step or state["blocks"] != n |
| or state.get("valid_supervised_tokens") != spec["matched_units"][name][str(n)]["valid_supervised_tokens"] |
| or state.get("optimizer_updates") != step |
| or [row.get("applied_lr") for row in state["curve"]] != spec["matched_units"][name][str(n)]["learning_rates"] |
| or state.get("batch_plan_identity") != identity or state["run"] != spec["run"] |
| or state["frozen_identity"] != spec["frozen_identity"] or state["reference_math"] != (name == "REF") |
| or state.get("fast_training", False) != (name != "REF") |
| or any(not torch.isfinite(p).all() for p in state["trainable"].values()) |
| or any(not all(torch.isfinite(torch.tensor(v)) for v in row.values()) for row in state["curve"])): |
| raise ValueError("pilot checkpoint identity/settings/nonfinite mismatch") |
| return arm |
|
|
|
|
| def validate_sample_metrics(sample: Sample, split: str, metrics: Metrics) -> None: |
| keys = sample["keys"][split] |
| if (metrics["rounds"] != len(keys) or metrics["requests"] != len({k["request"] for k in keys}) |
| or metrics["population_sha256"] != key_hash(keys) |
| or {r["request_sha256"]: r["rounds"] for r in metrics["request_totals"]} |
| != {digest(request): count for request, count in Counter(k["request"] for k in keys).items()}): |
| raise ValueError("pilot scored keys differ from the frozen sample") |
|
|
|
|
| def validate_sample_evaluation(spec: Spec, evaluation: Evaluation, splits: tuple[str, ...]) -> None: |
| sample = spec["sample"] |
| if (evaluation.get("sample_sha256") != sample["sha256"] or evaluation["k2_sha256"] != sample["k2_sha256"] |
| or set(evaluation["baseline"]) != set(splits) or set(evaluation["arms"]) != set(ARMS) |
| or set(evaluation["checkpoints"]) != set(ARMS)): |
| raise ValueError("pilot evaluation sample hash/K2/split mismatch") |
| if evaluation.get("split_roles") != split_roles(sample["keys"]): |
| raise ValueError("pilot selection splits must be labelled dev; reserved test is frozen-candidate only") |
| for split in splits: |
| validate_sample_metrics(sample, split, evaluation["baseline"][split]) |
| for arm in ARMS: |
| if (set(evaluation["arms"][arm]) != {str(n) for n in spec["checkpoints"]} |
| or set(evaluation["checkpoints"][arm]) != {str(n) for n in spec["checkpoints"]}): |
| raise ValueError("missing sampled pilot checkpoint") |
| for metrics in evaluation["arms"][arm].values(): |
| if set(metrics) != set(splits): |
| raise ValueError("sampled pilot split coverage mismatch") |
| for split in splits: |
| validate_sample_metrics(sample, split, metrics[split]) |
|
|
|
|
| @torch.no_grad() |
| def score(model: Drafter, data: D2Blocks, k2: Path, keys: list[RoundKey], chunk: int = 8192, |
| synthetic_windows: dict[str, dict[str, str]] | None = None) -> Metrics: |
| if not model.reference_math or model.fast_training: |
| raise ValueError("pilot evaluation requires qualified reference math") |
| expected = {(k["request"], k["P"]) for k in keys} |
| selected = [(index, key) for index, key in complete_rounds(data, k2) if (key["request"], key["P"]) in expected] |
| if not keys or len(expected) != len(keys) or [key for _, key in selected] != keys: |
| raise ValueError("pilot scored keys differ from the frozen sample") |
| total_kl = torch.zeros(5, dtype=torch.float64) |
| total_loss, total_accepted = 0., 0. |
| actual: list[RoundKey] = [] |
| kl_rows: list[list[float]] = [] |
| request_totals: dict[str, RequestTotal] = {} |
| for index, key in selected: |
| request, pos = key["request"], key["P"] |
| block = data[index] |
| if block["request"] != request or int(block["inputs"]["P"][0]) != pos: |
| raise ValueError("pilot scored keys differ from the frozen sample") |
| targets = cast(Targets, {k: t.to(model.activation_device) for k, t in block["targets"].items()}) |
| if not bool(targets["position_mask"].all()): |
| raise ValueError("sampled pilot round has incomplete targets") |
| inp = {k: t.to(model.activation_device) for k, t in block["inputs"].items()} |
| if synthetic_windows is not None: |
| entry = synthetic_windows[request] |
| exact = data.verified(entry["file"], entry["sha256"]) |
| original = cast(dict[str, torch.Tensor], torch.load(exact, map_location="cpu", weights_only=False)) |
| if int(original["P"][0]) != pos or int(original["anchor"][0]) != int(inp["anchor"][0]): |
| raise ValueError("synthetic capture-window alignment mismatch") |
| inp = {k: t.to(model.activation_device) for k, t in original.items()} |
| h, hx = model.hidden_from(inp) |
| picks, _ = model.greedy(h, hx, inp["anchor"], chunk=chunk) |
| candidate_n = match_length(picks[0].tolist(), targets["labels"][0].tolist()) |
| if candidate_n is None: |
| raise ValueError("indeterminate pilot replay round") |
| logits = torch.cat([model.logit_chunk(h, targets["previous"], lo, min(lo+chunk, model.c.vocab)) |
| for lo in range(0, model.c.vocab, chunk)], -1)[0] |
| kl = bucket_kl(logits, targets["top_ids"][0], targets["top_logp"][0], targets["tail_mass"][0]).cpu() |
| loss, _ = distillation_loss(model, h, chunk=chunk, **targets) |
| if not torch.isfinite(kl).all() or not torch.isfinite(loss): |
| raise FloatingPointError("nonfinite pilot evaluation") |
| if request not in data.request_families: |
| raise ValueError("pilot request missing source family") |
| aggregate = request_totals.setdefault(request, {"request_sha256": digest(request), |
| "family": data.request_families[request], "task_label": data.request_tasks[request], "rounds": 0, "accepted_sum": 0.}) |
| aggregate["rounds"] += 1 |
| aggregate["accepted_sum"] += candidate_n |
| kl_rows.append(kl.tolist()) |
| total_kl += kl |
| total_loss += float(loss) |
| total_accepted += candidate_n |
| actual.append(key) |
| if actual != keys: |
| raise ValueError("pilot scored keys differ from the frozen sample") |
| rounds = len(actual) |
| return {"requests": len({k["request"] for k in actual}), "rounds": rounds, "population_sha256": key_hash(actual), |
| "kl_by_position": (total_kl/rounds).tolist(), "loss": total_loss/rounds, |
| "accepted_length": total_accepted/rounds, |
| "kl_median_by_position": [statistics.median(row[i] for row in kl_rows) for i in range(5)], |
| "request_totals": [request_totals[r] for r in sorted(request_totals)]} |
|
|
|
|
| def step_agreement(root: Path, steps: int) -> dict[str, object]: |
| result: dict[str, object] = {"informational_only": True, "interpretation": "REF/FAST evolving matched trajectories; FAST-B target agreement"} |
| rows = [] |
| for step in range(1, steps+1): |
| arms = [json.loads((root/a/f"steps/step-{step:09d}.json").read_text()) for a in ARMS] |
| a, b = [torch.tensor(r["drafts"]) for r in arms[:2]] |
| if a.shape != b.shape or a.ndim != 2 or a.shape[1] != 5: |
| raise ValueError("matched per-step diagnostic shape mismatch") |
| rows.append({"step": step, "position1": float((a[:, 0] == b[:, 0]).float().mean()), |
| "exact_blocks": float((a == b).all(-1).float().mean()), |
| "loss_FAST_minus_REF": arms[1]["loss"]-arms[0]["loss"], |
| "target_agreement": {arm: {k: v for k, v in record.items() if k != "drafts"} |
| for arm, record in zip(ARMS, arms)}}) |
| result["steps"] = rows |
| return result |
|
|
|
|
| def verify_sample_draw(spec: Spec, manifest: Path, k2: dict[str, Path]) -> dict[str, D2Blocks]: |
| """Recompute the registered draw without modifying the frozen plan. |
| |
| Verify both populations even for a single-split evaluation job, before any |
| checkpoint/model load. A self-consistently rehashed alternative is refused. |
| """ |
| if sha256(manifest) != spec["manifest_sha256"]: |
| raise ValueError("pilot evaluation manifest mismatch") |
| SPLITS = population_splits(spec["sample"]["keys"]) |
| hashes = {s: sha256(k2[s]) for s in SPLITS} |
| if hashes != spec["sample"]["k2_sha256"]: |
| raise ValueError("pilot K2 sources differ from frozen sample") |
| data = {split: D2Blocks(manifest, split) for split in SPLITS} |
| populations = {s: [key for _, key in complete_rounds(data[s], k2[s])] for s in SPLITS} |
| expected = draw_sample(spec["manifest_sha256"], populations, hashes) |
| if expected != spec["sample"]: |
| raise ValueError("pilot frozen sample differs from seed-20261024 population draw") |
| return data |
|
|
|
|
| def evaluate(a: argparse.Namespace) -> None: |
| from .safety import assert_gpu_lease |
| assert_gpu_lease() |
| spec = read_spec(a.plan, sampled=True) |
| sample = spec["sample"] |
| SPLITS = population_splits(sample["keys"]) |
| splits = SPLITS if a.split == "all" else (a.split,) |
| k2 = {SPLITS[0]: a.k2_dev if SPLITS[0] == "dev" else a.k2_acap, "heldout_legal": a.k2_legal} |
| data = verify_sample_draw(spec, a.manifest, k2) |
| train_data = D2Blocks(a.manifest, "train") |
| verify_units(spec, train_data) |
| arms = {name: read_arm(a.root, name, spec, a.plan, len(train_data)) for name in ARMS} |
| torch.manual_seed(SEED) |
| model = load(a.ckpt, "cuda", rank=32, reference_math=True).eval() |
| if model.frozen_identity != spec["frozen_identity"] or initialization(model) != spec["initialization_sha256"]: |
| raise ValueError("pilot evaluator source/init mismatch") |
| extras = json.loads(a.manifest.read_text()).get("synthetic_eval_windows") |
| if extras is not None and not spec["synthetic"]: |
| raise ValueError("original capture-window sidecars are synthetic rehearsal only") |
| result: Evaluation = {"schema": "dsv41.d2.pilot-evaluation.v1" if a.split == "all" else PART_SCHEMA, "plan_sha256": sha256(a.plan), |
| "synthetic": spec["synthetic"], "evaluator": EVALUATOR, "split_roles": split_roles(spec["sample"]["keys"]), "sample_sha256": sample["sha256"], |
| "k2_sha256": {s: sha256(p) for s, p in k2.items()}, |
| "baseline": {s: score(model, data[s], k2[s], sample["keys"][s], synthetic_windows=extras) for s in splits}, "arms": {}, |
| "checkpoints": {name: arm["checkpoints"] for name, arm in arms.items()}, |
| "step_agreement": step_agreement(a.root, arms["REF"]["step_count"])} |
| if a.split != "all": |
| result["split"] = a.split |
| for name in ARMS: |
| result["arms"][name] = {} |
| for n in spec["checkpoints"]: |
| record = arms[name]["checkpoints"][str(n)] |
| state, actual = read_checkpoint(Path(record["path"])) |
| if actual != record["sha256"]: |
| raise ValueError("checkpoint changed during evaluation") |
| apply_trainables(model, state["trainable"]) |
| result["arms"][name][str(n)] = {s: score(model, data[s], k2[s], sample["keys"][s], synthetic_windows=extras) for s in splits} |
| print(json.dumps({"arm": name, "blocks": n, "metrics": result["arms"][name][str(n)]}), flush=True) |
| validate_sample_evaluation(spec, result, splits) |
| write_json(a.out, result) |
|
|
|
|
| def merge_parts(plan: Path, paths: list[Path], out: Path) -> None: |
| """CPU-only join of the two independently leased reference split jobs.""" |
| from .pilot_verdict import finite |
| spec = read_spec(plan, sampled=True) |
| SPLITS = population_splits(spec["sample"]["keys"]) |
| parts: dict[str, Evaluation] = {} |
| for path in paths: |
| part = cast(Evaluation, json.loads(path.read_text())) |
| split = part.get("split") |
| if (part["schema"] != PART_SCHEMA or split not in SPLITS or split in parts |
| or part["plan_sha256"] != sha256(plan) or part["synthetic"] != spec["synthetic"] |
| or part["evaluator"] != EVALUATOR): |
| raise ValueError("pilot evaluation part identity/split mismatch") |
| assert split is not None |
| finite(part) |
| validate_sample_evaluation(spec, part, (split,)) |
| parts[split] = part |
| if set(parts) != set(SPLITS): |
| raise ValueError("pilot merge requires exactly both held-out split parts") |
| first, second = [parts[s] for s in SPLITS] |
| if first["checkpoints"] != second["checkpoints"] or first["step_agreement"] != second["step_agreement"]: |
| raise ValueError("pilot evaluation parts disagree on checkpoints/step evidence") |
| result: Evaluation = {"schema": "dsv41.d2.pilot-evaluation.v1", "plan_sha256": sha256(plan), |
| "synthetic": spec["synthetic"], "evaluator": EVALUATOR, "split_roles": split_roles(spec["sample"]["keys"]), |
| "sample_sha256": spec["sample"]["sha256"], "k2_sha256": spec["sample"]["k2_sha256"], |
| "baseline": {s: parts[s]["baseline"][s] for s in SPLITS}, |
| "arms": {a: {str(n): {s: parts[s]["arms"][a][str(n)][s] for s in SPLITS} |
| for n in spec["checkpoints"]} for a in ARMS}, |
| "checkpoints": first["checkpoints"], "step_agreement": first["step_agreement"]} |
| validate_sample_evaluation(spec, result, SPLITS) |
| write_json(out, result) |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser(__doc__) |
| for n in ("plan", "out"): |
| p.add_argument("--"+n, required=True, type=Path) |
| for n in ("manifest", "root", "k2-acap", "k2-dev", "k2-legal"): |
| p.add_argument("--"+n, type=Path) |
| p.add_argument("--ckpt") |
| p.add_argument("--split", choices=(*SPLITS, "dev", "all"), default="all") |
| p.add_argument("--merge-parts", nargs=2, type=Path, help="CPU-only merge of dev and legal split outputs") |
| a = p.parse_args() |
| torch.set_num_threads(8) |
| if a.merge_parts: |
| merge_parts(a.plan, a.merge_parts, a.out) |
| else: |
| if any(getattr(a, n) is None for n in ("manifest", "root", "k2_legal", "ckpt")): |
| p.error("evaluation requires --manifest, --root, --k2-dev, --k2-legal and --ckpt") |
| if a.k2_dev is None and a.k2_acap is None: |
| p.error("evaluation requires --k2-dev (historical rehearsal: --k2-acap)") |
| evaluate(a) |
|
|
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
