| """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 |
|
|
|
|
| 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 |
| |
| 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}") |
|
|
|
|
| |
| from model_dedup.gguf_checkpoint import canonical as gguf_canonical |
|
|
| 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 |
| |
| w_t = teacher_tensor(hf_name) |
| w_t = hf_to_stored(name, w_t) |
| rows, width = w_t.shape |
| nblk = width // 128 |
| |
| 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) |
| |
| |
| 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}") |
|
|
