| """V1 same-input replay; private dumps, aggregate parity receipts only.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import importlib |
| import json |
| import os |
| from pathlib import Path |
| import sys |
| import time |
| from typing import Any |
|
|
| import torch |
| from torch import Tensor |
|
|
| from .model import Drafter, load |
| from .ops import frozen_linear |
| from .data import D2Blocks, sha256 |
|
|
|
|
| def v1_import(root: str) -> None: |
| if not (Path(root)/"verify/run_ref.py").is_file(): |
| raise FileNotFoundError(f"V1 verify modules missing under {root}") |
| sys.path.insert(0, root) |
|
|
|
|
| def g1_drafts(path: Path, data: D2Blocks) -> dict[str, dict[int, list[int]]]: |
| """Read D2-DATA g1.py's local seq_id/P output; require exact block coverage. |
| |
| G1's JSONL has no embedded manifest hash. Its file digest and the current |
| manifest identity must therefore be kept together in the evaluation receipt. |
| Use only the coordinator's G1 output for that exact dataset, never a guess. |
| """ |
| result: dict[str, dict[int, list[int]]] = {} |
| for line in path.read_text().splitlines(): |
| row = json.loads(line) |
| req = row["seq_id"] |
| if req in result: |
| raise ValueError("duplicate sequence in G1 K2 output") |
| result[req] = {int(p): [int(t) for t in ids] for p, ids in row["drafts"].items()} |
| expected_requests = {req for _, _, req in data.sequences} |
| if set(result) != expected_requests: |
| raise ValueError("G1 K2 request set differs from the selected dataset split") |
| for _, index, req in data.sequences: |
| length, gen_start = index[1:3] |
| if set(result[req]) != set(range(gen_start-1, length-2)): |
| raise ValueError("G1 K2 position coverage is incomplete or mismatched") |
| if any(len(ids) != 5 or any(t < 0 or t >= 129280 for t in ids) for ids in result[req].values()): |
| raise ValueError("G1 K2 draft shape or vocabulary invalid") |
| return result |
|
|
|
|
| def load_capture_set(path: str, reader: str | None = None, split: str | None = None, |
| k2_mapping: str | None = None) -> Any: |
| if reader is None: |
| return importlib.import_module("verify.captures").CaptureSet(path) |
| if split is None: |
| raise ValueError("D2 shard replay requires --split") |
| manifest = Path(path) |
| data = D2Blocks(manifest, split) |
| spec = importlib.util.spec_from_file_location("d2_data_shards", reader) |
| if spec is None or spec.loader is None: |
| raise ImportError(f"cannot load D2-DATA reader {reader}") |
| module = importlib.util.module_from_spec(spec) |
| sys.modules[spec.name] = module |
| spec.loader.exec_module(module) |
| drafts = None |
| by_seq = None |
| if k2_mapping: |
| if Path(k2_mapping).suffix == ".jsonl": |
| by_seq = g1_drafts(Path(k2_mapping), data) |
| else: |
| mapping = json.loads(Path(k2_mapping).read_text()) |
| if mapping["arm"] != "k2" or mapping["split"] != split or mapping["manifest_sha256"] != sha256(manifest): |
| raise ValueError("K2 draft mapping identity/split mismatch") |
| drafts = mapping["drafts"] |
| paths = [str((manifest.parent/entry["file"]).resolve()) for entry in data.manifest["splits"][split]["files"]] |
| cs = module.ShardCaptureSet(paths, drafts=drafts, drafts_by_seq=by_seq, verify_sha256=True) |
| if k2_mapping and any(s.engine_drafts_placeholder for s in cs.samples): |
| raise ValueError("K2 replay mapping has missing samples") |
| return cs |
|
|
|
|
| def capture_inputs(cs: Any, s: Any, device: str) -> tuple[dict[str, Tensor], dict[str, object]]: |
| positions = cs.window(s) |
| status = cs.window_status(s) |
| if status["missing"]: |
| raise ValueError(f"incomplete V1 window for {s.key}") |
| rows = [cs.rows[s.req][p] for p in positions] |
| has = [r.taps is not None for r in rows] |
| taps = torch.stack([r.taps.reshape(-1) if r.taps is not None else torch.zeros(15360, dtype=torch.bfloat16) for r in rows]) |
| mx = torch.stack([r.mx.reshape(-1) if r.mx is not None else torch.zeros(5120) for r in rows]) |
| inp = {"taps": taps[None], "engine_mx": mx[None], "has_taps": torch.tensor([has]), |
| "positions": torch.tensor([positions]), "mask": torch.ones(1, len(rows), dtype=torch.bool), |
| "anchor": torch.tensor([s.anchor]), "P": torch.tensor([s.P])} |
| meta = {"n": len(rows), "taps_rows": sum(has), "engine_mx_rows": len(rows)-sum(has), |
| "P_source": "taps" if has[-1] else "engine_mx", "first": positions[0], "last": positions[-1]} |
| return {k: v.to(device) for k, v in inp.items()}, meta |
|
|
|
|
| def read_checkpoint(path: Path) -> tuple[dict[str, Any], str]: |
| digest = sha256(path) |
| state = torch.load(path, map_location="cpu", weights_only=False) |
| if sha256(path) != digest: |
| raise ValueError("checkpoint changed while being read; use an immutable archived step") |
| return state, digest |
|
|
|
|
| def apply_trainables(model: Drafter, values: dict[str, Tensor]) -> None: |
| named = dict(model.named_parameters()) |
| if set(values) != set(named): |
| raise ValueError("checkpoint trainable names differ from Phase 2a model") |
| with torch.no_grad(): |
| for name, p in named.items(): |
| if p.shape != values[name].shape or not torch.isfinite(values[name]).all(): |
| raise ValueError(f"bad adapter tensor {name}") |
| p.copy_(values[name]) |
|
|
|
|
| def apply_adapter(model: Drafter, path: str) -> str: |
| state, digest = read_checkpoint(Path(path)) |
| if state.get("schema") != "dsv41.d2.train-checkpoint.v1" or state.get("frozen_identity") != model.frozen_identity: |
| raise ValueError("adapter checkpoint frozen source differs from loaded model") |
| apply_trainables(model, state["trainable"]) |
| return digest |
|
|
|
|
| @torch.no_grad() |
| def dump_sample(model: Drafter, inp: dict[str, Tensor], window: dict[str, object], meta: dict[str, object]) -> dict[str, Any]: |
| tr: dict[str, Any] = {} |
| hidden, hx = model.hidden(**inp, trace=tr) |
| ids, conf = model.greedy(hidden, hx, inp["anchor"]) |
| base = torch.cat([frozen_linear(hidden.float(), model.t["head.weight"][lo:lo+4096]) |
| for lo in range(0, model.c.vocab, 4096)], -1) |
| provenance = {"git_sha": os.environ.get("D2_GIT_SHA", "uncommitted-snapshot"), "torch": torch.__version__, |
| "source_sha256": {p.name: sha256(p) for p in sorted(Path(__file__).parent.glob("*.py"))}, |
| "frozen_identity": model.frozen_identity, "reference_math": model.reference_math, |
| "dtype_policy": "BF16 frozen masters; packed FP4 forward FP64 sum; FP32 activations/router/Markov; STE QDQ", **meta} |
| out: dict[str, Any] = {"drafts": ids[0].tolist(), "conf": conf[0].cpu(), "base_logits": base[0].cpu(), |
| "P": int(inp["P"][0]), "anchor": int(inp["anchor"][0]), "window": window, "meta": provenance, |
| "head_hidden": hidden[0].cpu(), "head_x": hx[0].cpu(), |
| "main_x_P": tr["main_x"][0, -1].cpu()} |
| kvs = [] |
| for i in range(model.c.stages): |
| st = tr[f"s{i}"] |
| kvs.append(st["kv"][0].cpu()) |
| out[f"s{i}"] = {n: st[n][0].cpu() for n in ("attn_out", "ffn_in", "ffn_out", "h_out")} |
| out[f"s{i}"]["route"] = {"idx": st["route"]["idx"][0].cpu(), "biased": st["route"]["biased"].cpu()} |
| out["kv_window"] = torch.stack(kvs, 1) |
| return out |
|
|
|
|
| def relative(a: Tensor, b: Tensor) -> float: |
| return float((a.double()-b.double()).norm()/b.double().norm().clamp_min(1e-30)) |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(__doc__) |
| ap.add_argument("--v1-root", required=True) |
| ap.add_argument("--captures", required=True) |
| ap.add_argument("--ckpt", required=True) |
| ap.add_argument("--references", required=True, help="V1 dumps root with ref and mirror directories") |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--adapter") |
| ap.add_argument("--limit", type=int, default=50) |
| ap.add_argument("--device", default="cpu") |
| ap.add_argument("--oracle", action="store_true") |
| ap.add_argument("--reference-math", action="store_true", help="CPU unquantized reference math; GPU quantized GEMMs") |
| ap.add_argument("--data-reader", help="D2-DATA shards.py; --captures then names MANIFEST.json") |
| ap.add_argument("--split", choices=("dev", "heldout_legal", "heldout_general", "acap")) |
| ap.add_argument("--k2-mapping") |
| ap.add_argument("--native-manifest") |
| ap.add_argument("--package-root", type=Path, help="verified merged native arm; owns native manifest and adapter provenance") |
| ap.add_argument("--parity-gate", action="store_true", help="also require exact agreement for an adapted arm") |
| a = ap.parse_args() |
| if a.device != "cpu": |
| from .safety import assert_gpu_lease |
| assert_gpu_lease() |
| v1_import(a.v1_root) |
| cs = load_capture_set(a.captures, a.data_reader, a.split, a.k2_mapping) |
| if cs.plumbing: |
| raise ValueError(f"V1 plumbing errors: {len(cs.plumbing)}") |
| samples = sorted(cs.samples, key=lambda s: (s.req, s.P, s.round)) |
| samples = [s for s in samples if (Path(a.references)/"ref"/f"{s.key}.pt").is_file()] |
| |
| if a.limit and len(samples) > a.limit: |
| samples = [samples[j*len(samples)//a.limit] for j in range(a.limit)] |
| if not samples: |
| raise ValueError("no reference-covered samples") |
| torch.set_num_threads(8) |
| package = None |
| if a.package_root: |
| from .reference import verify_package |
| if a.adapter or a.native_manifest or Path(a.ckpt).resolve() != a.package_root.resolve(): |
| raise ValueError("--package-root must equal --ckpt and cannot combine with adapter/native-manifest") |
| package = verify_package(a.package_root) |
| a.native_manifest = str(a.package_root/"NATIVE_MANIFEST.json") |
| model = load(a.ckpt, a.device, oracle=a.oracle, reference_math=a.reference_math, native_manifest=a.native_manifest).eval() |
| adapter_identity = apply_adapter(model, a.adapter) if a.adapter else (package["adapter_sha256"] if package else None) |
| out = Path(a.out) |
| out.mkdir(mode=0o700, parents=True, exist_ok=True) |
| receipts: list[dict[str, object]] = [] |
| begin = time.monotonic() |
| for s in samples: |
| inp, window = capture_inputs(cs, s, model.activation_device) |
| arm = dump_sample(model, inp, window, {"arm": "D2-TRAIN", "adapter": a.adapter, "native_root": a.ckpt, |
| "adapter_sha256": adapter_identity, "package_identity": package["drafter_identity"] if package else None, |
| "accumulation": str(model.acc), "uuid": os.environ.get("CUDA_VISIBLE_DEVICES")}) |
| arm["sample"] = {"key": s.key, "req": s.req, "P": s.P, "realized": s.realized, |
| "engine_drafts": s.engine_drafts, "engine_drafts_placeholder": getattr(s, "engine_drafts_placeholder", False)} |
| if a.k2_mapping: |
| arm["meta"]["k2_mapping_sha256"] = sha256(Path(a.k2_mapping)) |
| arm["meta"]["target_manifest_sha256"] = sha256(Path(a.captures)) |
| torch.save(arm, out/f"{s.key}.pt") |
| ref = torch.load(Path(a.references)/"ref"/f"{s.key}.pt", weights_only=False, map_location="cpu") |
| mir = torch.load(Path(a.references)/"mirror"/f"{s.key}.pt", weights_only=False, map_location="cpu") |
| err = relative(arm["base_logits"], ref["base_logits"]) |
| band = relative(mir["base_logits"], ref["base_logits"]) |
| row = {"key": s.key, "drafts_equal": arm["drafts"] == ref["drafts"], "logit_relerr": err, |
| "storage_band": band, "within_v1_3x_band": err <= 3*band, |
| "seconds": time.monotonic()-begin} |
| receipts.append(row) |
| (out/"parity.json").write_text(json.dumps({"samples": len(receipts), "rows": receipts, |
| "reference_math": model.reference_math, "frozen_identity": model.frozen_identity, |
| "acc": str(model.acc), "pass": all(r["drafts_equal"] and r["within_v1_3x_band"] for r in receipts)}, indent=2)+"\n") |
| print(json.dumps(row), flush=True) |
| if ((not a.adapter and not package) or a.parity_gate) and not all(r["drafts_equal"] and r["within_v1_3x_band"] for r in receipts): |
| raise SystemExit("D2 unadapted parity gate failed; see parity.json") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
