File size: 15,073 Bytes
e1acc15 | 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 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 | """V1 paired replay scoring and dev-only selection/G2/packaging gates.
Scores are aggregate-only. Candidate logits use that candidate's Markov head;
using the stock Markov head after adapting it would invalidate both KL and
V1's own-draft self-consistency check. Never score NOISE placeholder baselines.
"""
from __future__ import annotations
import argparse
from collections import defaultdict
from decimal import Decimal
import json
from pathlib import Path
import random
from statistics import mean
from typing import Any, TypedDict, cast
import torch
from torch import Tensor
from ...native.official import OfficialShards
from .data import D2Blocks, sha256
from .loss import bucket_divergence
from .replay import read_checkpoint
class Score(TypedDict):
request: str
candidate: float
k2: float
kl_candidate: list[float | None]
kl_baseline: list[float | None]
def match_length(drafts: list[int], truth: list[int | None]) -> int | None:
for i, (d, t) in enumerate(zip(drafts, truth)):
if t is None:
return None
if d != t:
return i
return len(drafts)
def clustered(rows: list[Score], iterations: int = 10000, seed: int = 23) -> dict[str, Any]:
if not rows or iterations < 100:
raise ValueError("empty scores or insufficient bootstrap iterations")
by: dict[str, list[float]] = defaultdict(list)
for r in rows:
by[r["request"]].append(r["candidate"]-r["k2"])
rng = random.Random(seed)
keys = sorted(by)
boots = []
for _ in range(iterations):
samples = [by[rng.choice(keys)] for _ in keys]
boots.append(sum(map(sum, samples))/sum(map(len, samples)))
boots.sort()
return {"rounds": len(rows), "requests": len(by), "candidate": mean(r["candidate"] for r in rows),
"k2": mean(r["k2"] for r in rows), "gain": mean(r["candidate"]-r["k2"] for r in rows),
"ci95_request_cluster": [boots[int(iterations*.025)], boots[int(iterations*.975)]],
"worst_request_gain": min(mean(v) for v in by.values()),
"kl_candidate": [mean(xs) if (xs := [cast(float, r["kl_candidate"][i]) for r in rows if r["kl_candidate"][i] is not None]) else None for i in range(5)],
"kl_baseline": [mean(xs) if (xs := [cast(float, r["kl_baseline"][i]) for r in rows if r["kl_baseline"][i] is not None]) else None for i in range(5)]}
def g2(legal: dict[str, Any], general: dict[str, Any], acap: dict[str, Any]) -> dict[str, Any]:
checks = {"legal_requests_ge20": legal["requests"] >= 20, "general_requests_ge10": general["requests"] >= 10,
"acap_requests_eq4": acap["requests"] == 4, "legal_gain_ge015": legal["gain"] >= .15,
"legal_ci_positive": legal["ci95_request_cluster"][0] > 0, "general_gain_ge_minus003": general["gain"] >= -.03,
"every_request_ge_minus02": all(s["worst_request_gain"] >= -.2 for s in (legal, general, acap)),
"kl_lower_every_position": all(a is not None and b is not None and a < b
for s in (legal, general, acap) for a, b in zip(s["kl_candidate"], s["kl_baseline"]))}
return {"verdict": "PASS" if all(checks.values()) else "FAIL", "checks": checks,
"legal": legal, "general": general, "acap": acap}
def select_dev(summaries: list[dict[str, Any]]) -> dict[str, Any]:
if not summaries or any(s["split"] != "dev" for s in summaries):
raise ValueError("checkpoint selection is dev-only")
if len({s["dataset"] for s in summaries}) != 1:
raise ValueError("dev dataset changed between candidates")
if any(any(x is None for x in s["metrics"]["kl_candidate"]) for s in summaries):
raise ValueError("missing dev KL position")
return min(summaries, key=lambda s: sum(w*k for w, k in zip((1., .8, .65, .5, .4), s["metrics"]["kl_candidate"])))
def packaging_gate(before: float, after: float, classification: dict[str, Any], expected_identity: str | None = None) -> dict[str, Any]:
classes = classification["classes"]
# Evaluate the specified decimal gate without binary .20*.75 rounding
# above the exact .15 boundary. No statistical tolerance is added.
checks = {"gain_retained_ge75pct": before > 0 and Decimal(str(after)) >= Decimal(".75")*Decimal(str(before)),
"zero_state_defects": classes.get("state_provenance_defect", 0) == 0 and classes.get("agree_but_state_flag", 0) == 0,
"zero_unexplained": classes.get("unexplained", 0) == 0,
"layer_excess_lt3x": bool(classification.get("relerr_p99"))
and set(classification["relerr_p99"]) == set(classification.get("arith_band_p99", {}))
and all(error < 3*classification["arith_band_p99"][name]
or error == classification["arith_band_p99"][name] == 0
for name, error in classification["relerr_p99"].items()),
"native_representation": classification.get("representation") == "native",
"nonempty_classification": classification["samples"] >= 50,
"own_reference_identity": expected_identity is None or classification.get("candidate_reference_identity") == expected_identity}
return {"pass": all(checks.values()), "checks": checks, "gain_before": before, "gain_after": after}
def bucket_kl(logits: Tensor, ids: Tensor, lp: Tensor, tail: Tensor) -> Tensor:
z = logits.double()
logz = z.logsumexp(-1)
top = z.gather(-1, ids.long()) - logz[:, None]
complement = z.scatter(-1, ids.long(), -torch.inf).logsumexp(-1)-logz
return bucket_divergence(top, complement, ids, lp.double(), tail.double(), "kl")
def markov(adapter: dict[str, Any] | None, native: Path) -> tuple[Tensor, Tensor]:
if adapter is not None:
state = adapter["trainable"]
prefix = "t.mtp__2__markov_head__"
return state[prefix+"embed__weight"].float(), state[prefix+"head__weight"].float()
return native_markov(native)
def native_markov(native: Path, manifest: Path | None = None) -> tuple[Tensor, Tensor]:
source = OfficialShards(str(native)) if manifest is None else OfficialShards(str(native), str(manifest))
names = ["mtp.2.markov_head.embed.weight", "mtp.2.markov_head.head.weight"]
for n in names:
source.verify(n)
return source.get(names[0]).float(), source.get(names[1]).float()
def package_markov(root: Path, adapter_digest: str) -> tuple[Tensor, Tensor, dict[str, Any]]:
from .reference import verify_package
receipt = verify_package(root)
if receipt["adapter_sha256"] != adapter_digest:
raise ValueError("package source adapter differs from selected checkpoint")
emb, head = native_markov(root, root/"NATIVE_MANIFEST.json")
return emb, head, receipt
def score_dumps(manifest: Path, split: str, candidate: Path, baseline: Path, k2: Path,
adapter: Path, native: Path, package_root: Path | None = None) -> dict[str, Any]:
if split == "train":
raise ValueError("evaluation must be held out")
data = D2Blocks(manifest, split)
state, adapter_digest = read_checkpoint(adapter)
package = None
if package_root is None:
emb, head = markov(state, native)
else:
emb, head, package = package_markov(package_root, adapter_digest)
be, bh = markov(None, native)
by_seq = None
if k2.is_file():
from .replay import g1_drafts
by_seq = g1_drafts(k2, data)
rows: list[Score] = []
# V1 dump names are keyed by absolute P and sequence index. ShardCaptureSet
# enumerates selected split sequences in manifest order (same as D2Blocks).
for seq_index, (path, seq, req) in enumerate(data.sequences):
L, gen = seq[1:3]
first_index = data.ends[seq_index-1] if seq_index else 0
next_round = gen-1
for offset, P in enumerate(range(gen-1, L-2)):
if P != next_round:
continue # fixed K2 trajectory; candidate predictions never select the scored rounds
block = data[first_index+offset]
t = block["targets"]
mask = t["position_mask"][0]
if not bool(mask.all()):
break
key = f"r{P:06d}_b{data.sequence_b[seq_index]}.pt"
a, b = [torch.load(d/key, map_location="cpu", weights_only=False) for d in (candidate, baseline)]
if a.get("meta", {}).get("adapter_sha256") != adapter_digest:
raise ValueError("candidate dump checkpoint identity differs from scored adapter")
if a.get("meta", {}).get("package_identity") != (package["drafter_identity"] if package else None):
raise ValueError("candidate dump package identity differs from scored arm")
if b.get("meta", {}).get("adapter_sha256") is not None:
raise ValueError("native baseline must be unadapted")
k = None if by_seq is not None else torch.load(k2/key, map_location="cpu", weights_only=False)
if k is not None and "k2" not in str(k.get("meta", {}).get("arm", "")).lower():
raise ValueError("explicit K2 baseline provenance required")
if k is not None and k.get("sample", {}).get("engine_drafts_placeholder", False):
raise ValueError("placeholder K2 baseline")
for dump in ([a, b, k] if k is not None else [a, b]):
if dump["P"] != P or dump["anchor"] != int(block["inputs"]["anchor"][0]) or dump["sample"]["req"] != req:
raise ValueError("mismatched trajectory or placeholder K2 baseline")
truth = [int(x) if bool(v) else None for x, v in zip(t["labels"][0], mask)]
if by_seq is not None:
k2_ids = by_seq[req][P]
else:
assert k is not None
k2_ids = k["drafts"]
cand_n, k2_n = match_length(a["drafts"], truth), match_length(k2_ids, truth)
if cand_n is None or k2_n is None:
raise ValueError("complete target block has an indeterminate draft score")
if len(a["drafts"]) != 5 or len(k2_ids) != 5:
raise ValueError("candidate or K2 does not contain five drafts")
next_round = P+k2_n+1
previous = t["previous"][0]
aq = a["base_logits"] + emb[previous] @ head.T
bq = b["base_logits"] + be[previous] @ bh.T
kl_a = bucket_kl(aq[mask], t["top_ids"][0, mask], t["top_logp"][0, mask], t["tail_mass"][0, mask]).tolist()
kl_b = bucket_kl(bq[mask], t["top_ids"][0, mask], t["top_logp"][0, mask], t["tail_mass"][0, mask]).tolist()
rows.append({"request": req, "candidate": cand_n, "k2": k2_n,
"kl_candidate": kl_a+[None]*(5-len(kl_a)), "kl_baseline": kl_b+[None]*(5-len(kl_b))})
return {"split": split, "dataset": data.fingerprint, "checkpoint": str(adapter), "checkpoint_sha256": adapter_digest,
"metrics": clustered(rows), "proxy": "V1 committed-trajectory paired vs explicit K2",
"k2_source": str(k2), "k2_jsonl_sha256": sha256(k2) if by_seq is not None else None,
"target_manifest_sha256": sha256(manifest), "package_identity": package["drafter_identity"] if package else None}
def main() -> None:
ap = argparse.ArgumentParser(__doc__)
commands = ap.add_subparsers(dest="command", required=True)
sel = commands.add_parser("select-dev")
sel.add_argument("summaries", nargs="+", type=Path)
sel.add_argument("--out", required=True, type=Path)
verdict = commands.add_parser("g2")
for n in ("legal", "general", "acap", "selection", "out"):
verdict.add_argument("--"+n, required=True, type=Path)
score = commands.add_parser("score")
for n in ("manifest", "candidate", "baseline", "k2", "adapter", "native", "out"):
score.add_argument("--"+n, required=True, type=Path)
score.add_argument("--split", choices=("dev", "heldout_legal", "heldout_general", "acap"), required=True)
score.add_argument("--package-root", type=Path, help="score the merged native arm with its actual stored Markov head")
score.add_argument("--selection", type=Path, help="immutable dev selection; required before test scoring")
package = commands.add_parser("packaging-gate")
package.add_argument("--gain-before", type=float, required=True)
package.add_argument("--gain-after", type=float, required=True)
for name in ("classification", "package-root", "out"):
package.add_argument("--"+name, type=Path, required=True)
a = ap.parse_args()
if a.command == "select-dev":
result = select_dev([json.loads(p.read_text()) for p in a.summaries])
elif a.command == "g2":
source = [json.loads(p.read_text()) for p in (a.legal, a.general, a.acap)]
if [s["split"] for s in source] != ["heldout_legal", "heldout_general", "acap"]:
raise ValueError("wrong held-out split labels")
chosen = json.loads(a.selection.read_text())["checkpoint_sha256"]
if any(s["checkpoint_sha256"] != chosen for s in source):
raise ValueError("test checkpoint differs from dev selection")
result = g2(*(s["metrics"] for s in source))
result["checkpoint_sha256"] = chosen
elif a.command == "packaging-gate":
from .reference import verify_package
receipt = verify_package(a.package_root)
result = packaging_gate(a.gain_before, a.gain_after, json.loads(a.classification.read_text()), receipt["drafter_identity"])
else:
if a.split != "dev":
if a.selection is None:
raise ValueError("test scoring requires a prior dev-only selection")
chosen = json.loads(a.selection.read_text())
if chosen["split"] != "dev" or chosen["checkpoint_sha256"] != sha256(a.adapter):
raise ValueError("test checkpoint differs from dev-only selection")
if a.package_root:
from .reference import verify_package
receipt = verify_package(a.package_root)
if receipt["adapter_sha256"] != chosen["checkpoint_sha256"]:
raise ValueError("package source differs from dev-selected checkpoint")
ledger_name = f"D2-PACKAGE-ONCE-{receipt['drafter_identity']}-{a.split}.json"
else:
ledger_name = f"D2-TEST-ONCE-{a.split}.json"
ledger = a.manifest.parent/ledger_name
with ledger.open("x") as f:
json.dump({"checkpoint_sha256": chosen["checkpoint_sha256"], "output": str(a.out),
"state": "test-scoring-started; do not retune or rerun"}, f)
result = score_dumps(a.manifest, a.split, a.candidate, a.baseline, a.k2, a.adapter, a.native, a.package_root)
# Exclusive creation enforces one reported test result and immutable selection.
with a.out.open("x") as f:
json.dump(result, f, indent=2)
f.write("\n")
if __name__ == "__main__":
main()
|