push artifacts/scripts/qat.py

artifacts/scripts/qat.py

@@ -0,0 +1,212 @@
+"""QAT on the validated surrogate: per layer, straight-through ternary codes + block scales + norm
+gains, distilled toward teacher trajectories with the eval's own metric (KL + hidden-state
+distance). Writes a new GGUF at --out.
+
+    .venv/bin/python qat.py --teacher /root/models/teacher --king /root/models/king-ix \
+        --data /root/train_samples.json --out /root/models/chal-v2 [--layers all|7]
+"""
+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 transformers import AutoModelForCausalLM
+
+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, canonical
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--teacher", default="/root/models/teacher")
+ap.add_argument("--king", default="/root/models/king-ix")
+ap.add_argument("--data", default="/root/train_samples.json")
+ap.add_argument("--out", default="/root/models/chal-v2")
+ap.add_argument("--layers", default="7,8,15,16,23,24,31,32,39,40,47,48,55,56,63")
+ap.add_argument("--steps", type=int, default=25)
+ap.add_argument("--lr", type=float, default=2e-3)
+ap.add_argument("--alpha", type=float, default=1.0)
+ap.add_argument("--seed", type=int, default=0)
+args = ap.parse_args()
+
+torch.manual_seed(args.seed)
+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()
+
+# ---- training data
+records = json.loads(Path(args.data).read_text())
+print("training samples:", len(records))
+
+# ---- the surrogate model and the king's state dict
+model = AutoModelForCausalLM.from_pretrained(args.teacher, torch_dtype=torch.bfloat16, device_map="cuda")
+model.eval()
+sd_king = torch.load("/dev/shm/king_hf.pt", map_location="cpu", weights_only=True)
+
+sys.path.insert(0, "/root/ingot/src")
+from eval_service import metrics  # noqa: E402
+from eval_service.config import LAYERS as CAP_LAYERS  # noqa: E402
+
+heads = metrics.heads({"teacher": Path(args.teacher), "king": Path(args.king)}, ["cuda:0"])["cuda:0"]
+head_t, head_k = heads["teacher"], heads["king"]
+
+# ---- gguf byte io
+src_gguf = Path(args.king) / "Ingot-IX-27B-Q2.gguf"
+out_dir = Path(args.out)
+out_dir.mkdir(parents=True, exist_ok=True)
+dst = out_dir / "chal-v2.gguf"
+shutil.copy(src_gguf, dst)
+
+
+def gguf_bytes(tname):
+    t = g.tensors[tname]
+    n = int(np.prod(t.shape))
+    size = n // TYPES[t.ggml_type].block_elems * TYPES[t.ggml_type].block_bytes
+    return np.memmap(dst, np.uint8, "r+", g.data_start + t.offset, (size,))
+
+
+def read_blocks(tname):
+    t = g.tensors[tname]
+    kind = TYPES[t.ggml_type]
+    n = int(np.prod(t.shape))
+    size = n // kind.block_elems * kind.block_bytes
+    raw = np.memmap(dst, np.uint8, "r", g.data_start + t.offset, (size,))
+    blocks = raw.reshape(-1, 34)
+    d = blocks[:, :2].copy().view("<f2").astype(np.float32)
+    qs = blocks[:, 2:]
+    codes = np.stack([(qs >> (2 * k)) & 3 for k in range(4)], axis=-1).reshape(-1, 128)
+    return codes, d
+
+
+def unrotate(w, tname):
+    """get()'s transform: HF-layout tensor from the stored one. w: (rows, width)."""
+    folded = rotation(g, tname)
+    if folded is not None:
+        block, signs = folded
+        rows, width = w.shape
+        w = (w.reshape(rows, width // block, block) @ H).reshape(rows, width)
+        w = w * torch.from_numpy(signs).float()
+    return ck._hf_order(tname, w, folded is not None)
+
+
+def hf_name(tname):
+    c = canonical(tname)
+    return c if c.startswith(("lm_head", "norm.")) else "model." + c
+
+
+def hf_module(hf):
+    parts = hf.split(".")
+    m = model
+    for p in parts[:-1]:
+        m = getattr(m, p)
+    return m
+
+
+def layer_tensors(il):
+    out = []
+    for kind in ("attn_qkv", "attn_gate", "ssm_alpha", "ssm_beta", "ssm_out", "attn_q", "attn_k", "attn_v", "attn_output", "ffn_gate", "ffn_up", "ffn_down"):
+        tname = f"blk.{il}.{kind}.weight"
+        if tname in g.tensors and g.tensors[tname].ggml_type == PQ2_0:
+            out.append(tname)
+    return out
+
+
+def train_layer(il):
+    torch.cuda.empty_cache()
+    print(f"--- layer {il}")
+    tensors = layer_tensors(il)
+    zs, ss = {}, {}
+    for tname in tensors:
+        codes, d = read_blocks(tname)
+        t = g.tensors[tname]
+        rows, width = int(np.prod(t.shape) // t.shape[0]), t.shape[0]
+        nblk = rows * width // 128
+        zs[tname] = torch.tensor(codes, dtype=torch.float32, device="cuda", requires_grad=True).reshape(nblk, 128)
+        ss[tname] = torch.tensor(d, dtype=torch.float32, device="cuda", requires_grad=True).reshape(nblk)
+    hooks = []
+    for tname in tensors:
+        hf = hf_name(tname)
+        module = hf_module(hf)
+        t = g.tensors[tname]
+        rows, width = int(np.prod(t.shape) // t.shape[0]), t.shape[0]
+
+        def make_pre(mod, tname=tname, rows=rows, width=width):
+            def pre(mod, inp):
+                z, s = zs[tname], ss[tname]
+                c = torch.clamp(torch.round(z), 0, 2)
+                W = (s * (c - 1.0)).reshape(rows, width)
+                mod._orig_weight = mod.weight
+                mod.weight = unrotate(W, tname).bfloat16()
+            return pre
+
+        def make_post(mod):
+            def post(mod, inp, out):
+                mod.weight = mod._orig_weight
+            return post
+
+        module.register_forward_pre_hook(make_pre(module))
+        module.register_forward_hook(make_post(module))
+        hooks.append(module)
+
+    opt = torch.optim.Adam(list(zs.values()) + list(ss.values()), lr=args.lr)
+    for step in range(args.steps):
+        torch.cuda.empty_cache()
+        total = torch.zeros((), device="cuda")
+        n = 0
+        for rec in records:
+            tokens = rec["prompt"] + rec["answer"]
+            if len(tokens) > 2600:
+                continue
+            first = len(rec["prompt"]) - 1
+            ids = torch.tensor([tokens], device="cuda")
+            out = model(input_ids=ids, output_hidden_states=True, use_cache=False)
+            hs = [out.hidden_states[l][0, first:].float() for l in CAP_LAYERS]
+            h_k = torch.stack(hs, 1)  # (rows, 8, 5120)
+            st = metrics.load_states(rec["states"], tokens[:-1], first).float().to("cuda")
+            m = min(h_k.shape[0], st.shape[0])
+            h_k, st = h_k[:m], st[:m]
+            lp_t = head_t.logprobs(st[:, -1])
+            lp_k = head_k.logprobs(h_k[:, -1])
+            kl = (lp_t.exp() * (lp_t - lp_k)).sum(-1)
+            a = st[1:]
+            b = h_k[1:]
+            dist = ((b - a).norm(dim=-1) / a.norm(dim=-1)).mean(1)
+            loss = kl.mean() + args.alpha * dist.mean()
+            opt.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(list(zs.values()) + list(ss.values()), 1.0)
+            opt.step()
+            total += loss.detach()
+            n += 1
+            del out, hs, h_k, st, lp_t, lp_k
+        print(f"  step {step}: loss {total.item() / n:.5f} over {n} samples", flush=True)
+
+    for tname in tensors:
+        z, s = zs[tname].detach().cpu().numpy(), ss[tname].detach().cpu().numpy()
+        c = np.clip(np.round(z), 0, 2).astype(np.uint8)
+        rows, width = int(np.prod(g.tensors[tname].shape) // g.tensors[tname].shape[0]), g.tensors[tname].shape[0]
+        nblk = rows * width // 128
+        raw = gguf_bytes(tname).reshape(-1, 34)
+        raw[:, :2] = s.astype(np.float16).reshape(-1, 1).astype("<f2").view(np.uint8).reshape(-1, 2)
+        qs = np.zeros((nblk, 32), np.uint8)
+        for j in range(128):
+            qs[:, j // 4] |= (c[:, j] << (2 * (j % 4))).astype(np.uint8)
+        raw[:, 2:] = qs
+        print(f"  wrote {tname}", flush=True)
+    for mod in hooks:
+        mod._forward_pre_hooks.clear()
+        mod._forward_hooks.clear()
+    del zs, ss
+    torch.cuda.empty_cache()
+
+
+for il in (int(x) for x in args.layers.split(",") if x):
+    train_layer(il)
+
+print("saved", dst)