push artifacts/scripts/qat_norms.py

artifacts/scripts/qat_norms.py

@@ -0,0 +1,119 @@
+"""End-to-end norm tuning on the surrogate: every RMSNorm weight (F32, free in the GGUF) trained
+against teacher trajectories with the eval's metric (KL + hidden-state distance), the same loss
+the eval uses. Writes a new GGUF.
+
+    .venv/bin/python qat_norms.py --data /root/train_samples.json --out /root/models/chal-norms
+"""
+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.types import TYPES
+from model_dedup.gguf_checkpoint import GgufCheckpoint
+
+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-norms")
+ap.add_argument("--steps", type=int, default=15)
+ap.add_argument("--lr", type=float, default=3e-3)
+ap.add_argument("--alpha", type=float, default=1.0)
+ap.add_argument("--max-tokens", type=int, default=2200)
+args = ap.parse_args()
+
+ck = GgufCheckpoint(args.king, "cpu")
+g = ck.g
+records = json.loads(Path(args.data).read_text())
+print("samples:", len(records))
+
+model = AutoModelForCausalLM.from_pretrained(args.teacher, torch_dtype=torch.bfloat16, device_map="cuda")
+model.eval()
+model.gradient_checkpointing_enable()
+sd_king = torch.load("/dev/shm/king_hf.pt", map_location="cpu", weights_only=True)
+model.load_state_dict(sd_king, strict=True)
+
+from eval_service import metrics  # noqa: E402
+from eval_service.config import 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"]
+
+params = []
+for name, p in model.named_parameters():
+    p.requires_grad_(False)
+    if any(k in name for k in ("layernorm.weight", "norm.weight")):
+        p.requires_grad_(True)
+        params.append(p)
+print("trainable norms:", len(params))
+
+opt = torch.optim.Adam(params, lr=args.lr)
+for step in range(args.steps):
+    torch.cuda.empty_cache()
+    total, n = 0.0, 0
+    for rec in records:
+        tokens = rec["prompt"] + rec["answer"]
+        if len(tokens) > args.max_tokens:
+            continue
+        first = len(rec["prompt"]) - 1
+        ids = torch.tensor([tokens], device="cuda")
+        out = model(input_ids=ids, output_hidden_states=True, use_cache=False)
+        h_k = torch.stack([out.hidden_states[l][0, first:].float() for l in LAYERS], 1)
+        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).mean()
+        a, b = st[1:], h_k[1:]
+        dist = ((b - a).norm(dim=-1) / a.norm(dim=-1)).mean()
+        loss = kl + args.alpha * dist
+        opt.zero_grad()
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(params, 1.0)
+        opt.step()
+        total += loss.item()
+        n += 1
+        del out, h_k, st, lp_t, lp_k, ids, loss
+    print(f"step {step}: loss {total / n:.5f} over {n} samples", flush=True)
+
+# ---- write the norms back (attn/post/output norms stored as hf + 1; ssm_norm as-is)
+out_dir = Path(args.out)
+out_dir.mkdir(parents=True, exist_ok=True)
+dst = out_dir / "chal-norms.gguf"
+shutil.copy(Path(args.king) / "Ingot-IX-27B-Q2.gguf", dst)
+
+def write_tensor(tname, values):
+    t = g.tensors[tname]
+    size = int(np.prod(t.shape)) * 4
+    raw = np.memmap(dst, np.uint8, "r+", g.data_start + t.offset, (size,))
+    raw.view("<f4")[:] = values.astype(np.float32).ravel()
+
+n = int(g.kv["qwen35.block_count"].value)
+with torch.no_grad():
+    for il in range(n):
+        for gguf_kind, hf_kind, plus1 in (
+            ("attn_norm", "input_layernorm", True),
+            ("post_attention_norm", "post_attention_layernorm", True),
+            ("ssm_norm", "linear_attn.norm", False),
+        ):
+            if f"blk.{il}.{gguf_kind}.weight" not in g.tensors:
+                continue
+            hf = f"model.layers.{il}.{hf_kind}.weight"
+            mod = model
+            for part in hf.split(".")[:-1]:
+                mod = getattr(mod, part)
+            w = mod.weight.detach().cpu().float()
+            write_tensor(f"blk.{il}.{gguf_kind}.weight", w.numpy() + (1.0 if plus1 else 0.0))
+    mod = model.model
+    w = mod.norm.weight.detach().cpu().float()
+    write_tensor("output_norm.weight", w.numpy() + 1.0)
+print("saved", dst)