| |
| """Verify the public Q25 checkpoint, static plan, and evaluation bundle.""" |
|
|
| from __future__ import annotations |
|
|
| from collections import Counter |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| ROOT = Path(__file__).resolve().parent |
| EXPECTED_DENSE = "10e8559713ef1d951c604605f8f3666a027a25a341363d0c17006f628cc38c1f" |
| EXPECTED_Q25 = "51120c7ecca5234a4e7cee424c199582e444eb04e434a2d7f08263e3ccc74a90" |
|
|
|
|
| def sha256(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as handle: |
| for chunk in iter(lambda: handle.read(16 << 20), b""): |
| digest.update(chunk) |
| return digest.hexdigest() |
|
|
|
|
| def verify_manifest() -> int: |
| rows = json.loads((ROOT / "inventory.json").read_text(encoding="utf-8")) |
| for relative, expected in rows.items(): |
| path = ROOT / relative |
| if not path.is_file(): |
| raise AssertionError(f"missing release file: {relative}") |
| if path.stat().st_size != expected["bytes"] or sha256(path) != expected["sha256"]: |
| raise AssertionError(f"release file changed: {relative}") |
| return len(rows) |
|
|
|
|
| def main() -> None: |
| files = verify_manifest() |
| dense_path = ROOT / "checkpoints/dense_model.pt" |
| q25_path = ROOT / "checkpoints/q25_export.pt" |
| if sha256(dense_path) != EXPECTED_DENSE or sha256(q25_path) != EXPECTED_Q25: |
| raise AssertionError("checkpoint digest differs from the article manifest") |
|
|
| plan = json.loads((ROOT / "selection/Q25/plan.json").read_text(encoding="utf-8")) |
| assignments = { |
| (int(row["layer"]), int(row["head"])): row["mode"] for row in plan["groups"] |
| } |
| if len(assignments) != 96: |
| raise AssertionError("Q25 plan does not contain 96 unique heads") |
| plan_modes = Counter(assignments.values()) |
| if plan_modes != Counter({"LOCAL": 81, "LOCAL_GRAPH": 15}): |
| raise AssertionError(f"unexpected plan modes: {plan_modes}") |
|
|
| state = torch.load(q25_path, map_location="cpu", weights_only=True) |
| state_modes: dict[tuple[int, int], int] = {} |
| for layer in range(24): |
| key = f"base_model.blocks.{layer}.attention.modes" |
| values = state[key].tolist() |
| if len(values) != 16: |
| raise AssertionError(f"layer {layer} does not contain 16 head modes") |
| for head, mode in enumerate(values): |
| state_modes[(layer, head)] = int(mode) |
| counts = Counter(state_modes.values()) |
| if counts != Counter({0: 288, 2: 81, 1: 15}): |
| raise AssertionError(f"checkpoint mode counts differ: {counts}") |
| for identity, mode in assignments.items(): |
| expected = 1 if mode == "LOCAL_GRAPH" else 2 |
| if state_modes[identity] != expected: |
| raise AssertionError(f"checkpoint/plan mismatch at {identity}") |
|
|
| exported_layers = { |
| int(key.split(".")[2]) |
| for key in state |
| if key.endswith("attention.global_head_indices") |
| } |
| if exported_layers != {2, 4, 5, 17, 18, 19, 20, 21, 22, 23}: |
| raise AssertionError(f"unexpected physically packed layers: {exported_layers}") |
|
|
| with np.load(ROOT / "evaluation/q25_470_documents.npz") as bundle: |
| if bundle["tokens"].shape != (470, 8192) or bundle["tokens"].dtype != np.uint16: |
| raise AssertionError("evaluation token bundle has the wrong shape or dtype") |
| languages = Counter(str(item) for item in bundle["languages"]) |
| if languages != Counter({"zh": 300, "en": 50, "de": 50, "es": 50, "ar": 20}): |
| raise AssertionError(f"evaluation language support differs: {languages}") |
| if len(set(str(item) for item in bundle["document_sha256"])) != 470: |
| raise AssertionError("evaluation documents are duplicated") |
|
|
| parameter_elements = sum(value.numel() for value in state.values()) |
| print(json.dumps({ |
| "verified_files": files, |
| "dense_checkpoint_sha256": EXPECTED_DENSE, |
| "q25_checkpoint_sha256": EXPECTED_Q25, |
| "checkpoint_state_elements_including_buffers": parameter_elements, |
| "global_heads": counts[0], |
| "local_heads": counts[2], |
| "local_graph_heads": counts[1], |
| "physically_packed_layers": sorted(exported_layers), |
| "evaluation_documents": 470, |
| "status": "verified", |
| }, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
