push artifacts/scripts/scale_refit.py

artifacts/scripts/scale_refit.py

@@ -0,0 +1,158 @@
+"""Challenger v1: the king with every PQ2_0 block's fp16 scale refit by least squares against the
+teacher's rotated weights (codes untouched). Works entirely in the stored (rotated) layout, where
+stored = W_hf · signs · H per block for both folded weights and latent embeddings.
+
+    .venv/bin/python scale_refit.py [--out /root/models/chal-v1]
+"""
+import argparse
+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.dequant import hadamard, rotation
+from common.gguf.types import PQ2_0, TYPES
+from model_dedup.gguf_checkpoint import GgufCheckpoint, Heads
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--king", default="/root/models/king-ix")
+ap.add_argument("--teacher", default="/root/models/teacher")
+ap.add_argument("--out", default="/root/models/chal-v1")
+ap.add_argument("--max-tensors", type=int, default=0, help="limit for testing")
+args = ap.parse_args()
+
+ck = GgufCheckpoint(args.king, "cpu")
+g = ck.g
+heads = ck._heads
+block_size = int(g.kv["prism.hadamard.block_size"].value)
+H = hadamard(block_size)
+print("block_size", block_size, "heads", heads)
+
+teacher_shards = {}
+import json
+
+idx = json.loads((Path(args.teacher) / "model.safetensors.index.json").read_text())
+weight_map = idx["weight_map"]
+
+def teacher_tensor(hf_name: str) -> np.ndarray:
+    raw_name = (
+        hf_name
+        if hf_name == "lm_head.weight"
+        else "model.language_model." + hf_name[len("model.") :]
+    )
+    shard = weight_map[raw_name]
+    if shard not in teacher_shards:
+        import safetensors
+
+        teacher_shards[shard] = safetensors.safe_open(str(Path(args.teacher) / shard), "pt")
+    t = teacher_shards[shard].get_tensor(raw_name)
+    return t.to(dtype=torch.float32).numpy()
+
+import torch  # noqa: E402
+
+
+def tile(w: np.ndarray, dim: int, head: int) -> np.ndarray:
+    """Inverse of GgufCheckpoint._untile: HF grouped [n_k, rep, head] -> stored tiled [rep, n_k, head]."""
+    shape = list(w.shape)
+    x = w.reshape(*shape[:dim], heads.n_k, heads.rep, head, *shape[dim + 1 :])
+    return x.transpose(dim, dim + 1).reshape(shape)
+
+
+def hf_to_stored(name: str, w: np.ndarray) -> np.ndarray:
+    """Teacher tensor, HF layout -> the king's stored layout (row order, then rotation)."""
+    if name.startswith("layers.") and heads.rep > 1:
+        kind = name.split(".")[2]
+        if kind == "in_proj_qkv" and name.endswith("in_proj_qkv.weight"):
+            qk = 2 * heads.n_k * heads.k_dim
+            w = np.concatenate([w[:qk], tile(w[qk:], 0, heads.v_dim)])
+        elif kind == "in_proj_z":
+            w = tile(w, 0, heads.v_dim)
+        elif kind in ("in_proj_a", "in_proj_b"):
+            w = tile(w, 0, 1)
+        elif kind == "out_proj" and not (rot_signs(name) is not None and ck._grouped):
+            w = tile(w.T, 0, heads.v_dim).T
+    signs = rot_signs(name)
+    if signs is None:
+        return w
+    rows, width = w.shape
+    nblk = width // block_size
+    # stored = (W · signs) · H per col block, as get()'s inverse: W = (stored · H) · signs
+    w = (w.reshape(rows, nblk, block_size) * signs.reshape(nblk, block_size)) @ H
+    return w.reshape(rows, width)
+
+
+def rot_signs(name: str) -> np.ndarray | None:
+    kv = g.kv
+    listed = [
+        kv["prism.hadamard." + key].value
+        for key in ("weight_names", "inverse_weight_names")
+        if "prism.hadamard." + key in kv
+    ]
+    if not any(name in names for names in listed):
+        return None
+    widths = kv["prism.hadamard.sign_widths"].value
+    values = kv["prism.hadamard.sign_values"].value
+    width = g.tensors[name].shape[0]
+    at = 0
+    for w_ in widths:
+        if w_ == width:
+            return np.asarray(values[at : at + w_], np.float32)
+        at += w_
+    raise ValueError(f"no sign vector of width {width} for {name}")
+
+
+# the gguf tensor name each HF name maps to (dedup canonical, with the .weight suffix)
+from model_dedup.gguf_checkpoint import canonical as gguf_canonical  # noqa: E402
+
+out_dir = Path(args.out)
+out_dir.mkdir(parents=True, exist_ok=True)
+src = Path(args.king) / "Ingot-IX-27B-Q2.gguf"
+dst = out_dir / "chal-v1.gguf"
+shutil.copy(src, dst)
+
+done = 0
+with open(dst, "r+b") as fh:
+    for name, t in g.tensors.items():
+        if t.ggml_type != PQ2_0:
+            continue
+        if args.max_tensors and done >= args.max_tensors:
+            break
+        canon = gguf_canonical(name)
+        if canon is None:
+            print("skip (no hf name):", name)
+            continue
+        hf_name = "model." + canon if not canon.startswith(("lm_head", "norm.")) else canon
+        # canonical gives e.g. layers.0.self_attn.q_proj.weight; HF: model.layers.0...
+        w_t = teacher_tensor(hf_name)
+        w_t = hf_to_stored(name, w_t)
+        rows, width = w_t.shape
+        nblk = width // 128  # the PQ2_0 block is 128 wide; the Hadamard block is not
+        # read raw blocks
+        size = rows * nblk * 34
+        raw = np.memmap(dst, np.uint8, "r+", g.data_start + t.offset, (size,))
+        blocks = raw.reshape(-1, 34)
+        qs = blocks[:, 2:]
+        codes = np.stack([(qs >> (2 * k)) & 3 for k in range(4)], axis=-1).reshape(rows, nblk, 128)
+        if (codes == 3).any():
+            raise SystemExit(f"{name}: a code 3 in the king")
+        c_vals = codes.astype(np.float32) - 1.0
+        w_t_b = w_t.reshape(rows, nblk, 128)
+        d_orig = blocks[:, :2].copy().view("<f2").astype(np.float32).reshape(rows, nblk)
+        support = (c_vals * c_vals).sum(-1)
+        num = (c_vals * w_t_b).sum(-1)
+        # a block of mostly-zero codes has no stable least-squares scale: ridge + a support floor,
+        # and never more than double the king's own scale (its codes were fit with it)
+        d_new = num / (support + 4.0)
+        d_new = np.where(support >= 8, d_new, d_orig)
+        d_new = np.clip(d_new, -2.0 * np.abs(d_orig), 2.0 * np.abs(d_orig))
+        d_new = np.clip(d_new, -5000.0, 5000.0)
+        d_new = d_new.astype(np.float16)
+        blocks[:, :2] = d_new.astype("<f2").view(np.uint8).reshape(-1, 2)
+        done += 1
+        if done % 50 == 0:
+            print(f"refit {done} tensors ({name})", flush=True)
+
+print(f"refit {done} tensors -> {dst}")