| """GPU scale refit: per PQ2_0 block, d* = <c, teacher_rotated>/(support + 4), capped at 2x the |
| king's own scale, written back to a copy of the GGUF. All block math on CUDA.""" |
| import argparse |
| import json |
| import shutil |
| import sys |
| from pathlib import Path |
|
|
| sys.path.insert(0, "/root/ingot/src") |
| import numpy as np |
| import torch |
|
|
| from common.gguf import read_header |
| from common.gguf.dequant import hadamard, rotation |
| from common.gguf.types import PQ2_0 |
| from model_dedup.gguf_checkpoint import GgufCheckpoint, canonical |
|
|
| 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("--device", default="cuda") |
| args = ap.parse_args() |
|
|
| dev = args.device |
| ck = GgufCheckpoint(args.king, "cpu") |
| g = ck.g |
| hd = ck._heads |
| block_size = int(g.kv["prism.hadamard.block_size"].value) |
| H = torch.from_numpy(hadamard(block_size)).float().to(dev) |
|
|
| idx = json.loads((Path(args.teacher) / "model.safetensors.index.json").read_text()) |
| weight_map = idx["weight_map"] |
|
|
| import safetensors |
|
|
| _shards = {} |
|
|
|
|
| def teacher_tensor(hf_name): |
| raw_name = "lm_head.weight" if hf_name == "lm_head.weight" else "model.language_model." + hf_name[len("model.") :] |
| shard = weight_map[raw_name] |
| if shard not in _shards: |
| _shards[shard] = safetensors.safe_open(str(Path(args.teacher) / shard), "pt") |
| t = _shards[shard].get_tensor(raw_name) |
| return t.float().to(dev) |
|
|
|
|
| def tile(w, dim, head): |
| shape = list(w.shape) |
| x = w.reshape(*shape[:dim], hd.n_k, hd.rep, head, *shape[dim + 1 :]) |
| return x.transpose(dim, dim + 1).reshape(shape) |
|
|
|
|
| def rot_signs(name): |
| 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 torch.from_numpy(np.asarray(values[at : at + w_], np.float32)).to(dev) |
| at += w_ |
| raise ValueError(f"no sign vector of width {width} for {name}") |
|
|
|
|
| def hf_to_stored(name, w): |
| if name.startswith("layers.") and hd.rep > 1: |
| kind = name.split(".")[2] |
| if kind == "in_proj_qkv" and name.endswith("in_proj_qkv.weight"): |
| qk = 2 * hd.n_k * hd.k_dim |
| w = torch.cat([w[:qk], tile(w[qk:], 0, hd.v_dim)]) |
| elif kind == "in_proj_z": |
| w = tile(w, 0, hd.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, hd.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) |
|
|
|
|
| out_dir = Path(args.out) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| dst = out_dir / "chal-v1.gguf" |
| shutil.copy(Path(args.king) / "Ingot-IX-27B-Q2.gguf", dst) |
|
|
| done = 0 |
| with open(dst, "r+b") as fh: |
| for name, t in g.tensors.items(): |
| if t.ggml_type != PQ2_0: |
| continue |
| canon = canonical(name) |
| if canon is None: |
| continue |
| hf_name = canon if canon.startswith(("lm_head", "norm.")) else "model." + canon |
| w_t = hf_to_stored(name, teacher_tensor(hf_name)) |
| rows, width = w_t.shape |
| nblk = rows // 128 |
| size = rows * width // 128 * 34 |
| raw = np.memmap(dst, np.uint8, "r+", g.data_start + t.offset, (size,)) |
| blocks = raw.reshape(-1, 34) |
| qs = torch.from_numpy(blocks[:, 2:].copy()).to(dev) |
| codes = torch.stack([(qs >> (2 * k)) & 3 for k in range(4)], -1).reshape(-1, 128) |
| c_vals = (codes.float() - 1.0).reshape(width, nblk, 128).permute(1, 0, 2) |
| d_orig = torch.from_numpy(blocks[:, :2].copy().view("<f2").astype(np.float32)).to(dev) |
| d_orig = d_orig.reshape(width, nblk).T |
| w_b = w_t.reshape(nblk, 128, width).permute(0, 2, 1) |
| support = (c_vals * c_vals).sum(-1) |
| num = (c_vals * w_b).sum(-1) |
| d_new = num / (support + 4.0) |
| d_new = torch.where(support >= 8, d_new, d_orig) |
| cap = 2.0 * d_orig.abs() |
| d_new = torch.clamp(d_new, -cap, cap).clamp(-5000.0, 5000.0) |
| blocks[:, :2] = d_new.T.reshape(-1).cpu().numpy().astype(np.float16).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}") |
|
|
