push artifacts/scripts/identity_check.py

artifacts/scripts/identity_check.py

@@ -0,0 +1,97 @@
+"""Identity round-trip sanity v2 (chunked): rebuild PQ2_0 scales from the king's own
+weights (codes unchanged) along the correct block axis; write back as a challenger."""
+import hashlib
+import shutil
+import sys
+from pathlib import Path
+
+sys.path.insert(0, "/root/ingot/src")
+import numpy as np
+
+from common.gguf import read_header
+from common.gguf.types import PQ2_0
+
+KING = Path("/root/models/king-ix/Ingot-IX-27B-Q2.gguf")
+OUT_DIR = Path("/root/models/chal-identity")
+OUT = OUT_DIR / KING.name
+CHUNK = 1 << 20  # blocks per chunk
+
+g = read_header(KING)
+pq = [(name, t) for name, t in g.tensors.items() if t.ggml_type == PQ2_0]
+print(f"PQ2_0 tensors: {len(pq)}")
+
+
+def codes_d(blocks):
+    d = blocks[:, :2].copy().view("<f2").astype(np.float32).reshape(-1)
+    qs = blocks[:, 2:]
+    codes = np.stack([(qs >> (2 * k)) & 3 for k in range(4)], axis=-1).reshape(-1, 128).astype(np.float32) - 1.0
+    return codes, d
+
+
+def refit_scales(blocks):
+    codes, d = codes_d(blocks)
+    den = (codes * codes).sum(1)
+    num = (codes * codes * d[:, None]).sum(1)
+    return np.where(den > 0, num / np.maximum(den, 1e-9), d)
+
+
+# ---- Part A: axis diagnostic on small PQ2_0 tensors ----
+print("=== AXIS DIAGNOSTIC ===")
+shown = 0
+for name, t in pq:
+    n_blocks = int(np.prod(t.shape)) // 128
+    if n_blocks > 200_000:
+        continue
+    b = np.asarray(np.memmap(KING, np.uint8, "r", g.data_start + t.offset, (n_blocks * 34,)).reshape(-1, 34))
+    d_orig = b[:, :2].copy().view("<f2").astype(np.float32).reshape(-1)
+    d_ok = refit_scales(b).astype(np.float16).astype(np.float32)
+    err = np.abs(d_ok - d_orig)
+    scale = max(np.abs(d_orig).max(), 1e-9)
+    # v1's wrong axis on a square grid, for the failure record
+    side = int(round(np.sqrt(n_blocks)))
+    bad = None
+    if side * side == n_blocks:
+        codes, _ = codes_d(b)
+        dT = d_orig.reshape(side, side).T.reshape(-1)
+        den = (codes * codes).sum(1)
+        num = (codes * codes * dT[:, None]).sum(1)
+        d_bad = np.where(den > 0, num / np.maximum(den, 1e-9), dT).astype(np.float16).astype(np.float32)
+        bad = (np.abs(d_bad - d_orig).max() / scale)
+    print(f"  {name[:46]:<46} blocks={n_blocks:>8} maxerr={err.max()/scale:.2e} bits_flipped={(d_ok.astype('<f2').view(np.uint16) != d_orig.astype('<f2').view(np.uint16)).sum()}"
+          + (f"   v1-axis maxerr={bad:.2e}" if bad is not None else ""))
+    shown += 1
+    if shown >= 8:
+        break
+
+# ---- Part B: write identity challenger ----
+OUT_DIR.mkdir(parents=True, exist_ok=True)
+shutil.copy(KING, OUT)
+patched = changed = 0
+with open(OUT, "r+b") as fh:
+    for name, t in pq:
+        n_blocks = int(np.prod(t.shape)) // 128
+        raw = np.memmap(OUT, np.uint8, "r+", g.data_start + t.offset, (n_blocks * 34,))
+        for s in range(0, n_blocks, CHUNK):
+            e = min(s + CHUNK, n_blocks)
+            b = np.asarray(raw[s * 34:e * 34]).reshape(-1, 34)
+            d_new = refit_scales(b).astype(np.float16).astype("<f2").view(np.uint8).reshape(-1, 2)
+            changed += int((d_new != b[:, :2]).sum() // 2)
+            b[:, :2] = d_new
+        patched += 1
+        if patched % 100 == 0:
+            print(f"  patched {patched}/{len(pq)}", flush=True)
+
+print(f"\npatched {patched} tensors; scales differing from original: {changed}")
+
+
+def h(p):
+    m = hashlib.sha256()
+    with open(p, "rb") as f:
+        for c in iter(lambda: f.read(1 << 24), b""):
+            m.update(c)
+    return m.hexdigest()[:16]
+
+
+print(f"king sha256     {h(KING)}")
+print(f"identity sha256 {h(OUT)}")
+print("WROTE", OUT)