File size: 4,334 Bytes
e69b72a | 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 | #!/usr/bin/env python3
"""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()
|