File size: 17,717 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 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 | """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()
|