Upload podflow/calib_train3_v9.py with huggingface_hub

podflow/calib_train3_v9.py

@@ -0,0 +1,409 @@
+"""Exact-objective joint calibration: per-(row,128-block) PQ2_0 scale factors + inner norms.
+
+Objective mirrors eval_service/metrics.py compare() exactly:
+  KL   = sum over answer rows of KL(teacher || model) where each side's distribution is
+         FinalHead/GgufHead applied to the LAST captured layer's state,
+  dist = sum over answer rows, mean over the 8 captured layers, of ||b-a||/||a||.
+Teacher targets come from gen_teacher.py states files (hidden states per row x 8 layers).
+Model states come from the surrogate forward (hidden_states[L] for L in 8..64 step 8,
+pre-final-norm, matching llama-states' captured rows).
+"""
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+REPO = "/root/ingot"
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--surrogate", default="/dev/shm/king_hf.pt")
+    ap.add_argument("--gguf", default="/root/models/king-xii")
+    ap.add_argument("--teacher", default="/root/models/teacher")
+    ap.add_argument("--samples", default="/root/train_samples.json")
+    ap.add_argument("--out", default="/root/models/chal-calib")
+    ap.add_argument("--steps", type=int, default=400)
+    ap.add_argument("--lr", type=float, default=5e-3)
+    ap.add_argument("--dist-weight", type=float, default=0.5)
+    ap.add_argument("--max-states", type=int, default=64)
+    args = ap.parse_args()
+
+    import faulthandler, traceback
+    faulthandler.enable()
+    import torch
+    import torch.nn.functional as F
+    sys.path.insert(0, f"{REPO}/src")
+    from transformers import AutoModelForCausalLM
+    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, _untile
+    from eval_service import config
+    from eval_service.metrics import FinalHead, GgufHead
+
+    dev = "cuda"
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+    sd = torch.load(args.surrogate, map_location="cpu", weights_only=True)
+    model = AutoModelForCausalLM.from_pretrained(args.teacher, dtype=torch.bfloat16,
+                                                 device_map={"": dev})  # no 54GB CPU copy: 125GB-RAM pods
+    model.load_state_dict({k: v.to(torch.bfloat16) for k, v in sd.items()}, strict=False)
+    model.to(dev)
+    model.gradient_checkpointing_enable()
+    model.config.use_cache = False
+
+    ck = GgufCheckpoint(args.gguf, "cpu")
+    g = ck.g
+    hd = ck._heads
+    kv = g.kv
+    bs = int(kv["prism.hadamard.block_size"].value)
+    H = torch.from_numpy(hadamard(bs)).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):
+        listed = [kv["prism.hadamard." + k].value
+                  for k in ("weight_names", "inverse_weight_names")
+                  if "prism.hadamard." + k in kv]
+        if not any(name in ns for ns 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_
+        return None
+
+    def hf_to_stored(name, w):
+        if name.startswith("blk.") and hd.rep > 1:
+            kind = name.split(".")[2]
+            if kind == "attn_qkv":
+                qk = 2 * hd.n_k * hd.k_dim
+                w = torch.cat([w[:qk], tile(w[qk:], 0, hd.v_dim)])
+            elif kind == "attn_gate":
+                w = tile(w, 0, hd.v_dim)
+            elif kind in ("ssm_alpha", "ssm_beta"):
+                w = tile(w, 0, 1)
+            elif kind == "ssm_out" 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, None
+        rows, width = w.shape
+        nblk = width // bs
+        h_l = H.to(w.device)
+        s_l = signs.to(w.device)
+        w = ((w.reshape(rows, nblk, bs) * s_l.reshape(nblk, bs)) @ h_l).reshape(rows, width)
+        return w, signs
+
+    # ---- wrap PQ2_0 linears (skip output.weight: GgufHead path) ----
+    for p in model.parameters():
+        p.requires_grad_(False)
+    gguf_to_hf = {}
+    sd_keys = list(sd)
+    for gg_name in g.tensors:
+        if gg_name == "output.weight" or gg_name.endswith("norm.weight"):
+            continue
+        canon = canonical(gg_name)
+        if canon is None or g.tensors[gg_name].ggml_type != PQ2_0:
+            continue
+        suffix = canon[len("model."):] if canon.startswith("model.") else canon
+        cands = [k for k in sd_keys if k == canon or k.endswith(suffix)]
+        if cands:
+            gguf_to_hf[gg_name] = cands[0]
+
+    scale_params = {}
+
+    # Effective-weight linear with per-(row,128-block) scales, dequantized ON THE FLY from the
+    # GGUF's native 2-bit blocks (6.8GB total vs 51GB for bf16 weights). Custom autograd.Function:
+    # nothing weight-sized is retained for backward (weights are recomputed), and training sees
+    # exactly the bytes the validator's eval loads.
+    def unpack(c):
+        # 4 two-bit codes per byte, LSB-first: code j (j=0..127) from byte j//4, bits (j%4)*2
+        return torch.stack([(c >> sh) & 3 for sh in (0, 2, 4, 6)], dim=-1).view(
+            c.shape[0], c.shape[1], 128)
+
+    class _ScaledWFn(torch.autograd.Function):
+        @staticmethod
+        def build_w(s, codes, d0, hb, signs, rows, nblk_w):
+            width = d0.shape[1] * 128
+            d_eff = d0.float() * s.squeeze(-1).float()  # [rows, nblk] fp32
+            q = unpack(codes)
+            m1 = ((q.float() - 1.0) * d_eff.unsqueeze(-1)).to(torch.bfloat16).view(rows, width)
+            if hb is not None:
+                gsz = hb.shape[0]
+                m1 = (m1.reshape(rows, width // gsz, gsz) @ hb).reshape(rows, width) * signs
+            return m1
+
+        @staticmethod
+        def forward(ctx, x, s, codes, d0, hb, signs, bias, rows, nblk_w, pr, pc):
+            # stored order -> HF order: the fork stores linear-attention value heads tiled
+            # [rep, n_k, hd]; the HF model expects them grouped [n_k, rep, hd] (GgufCheckpoint
+            # _hf_order). Missing this fed head-permuted weights to the model in v7.
+            w = _ScaledWFn.build_w(s, codes, d0, hb, signs, rows, nblk_w)
+            if pr is not None:
+                w = w[pr]
+            if pc is not None:
+                w = w[:, pc]
+            if x.dtype in (torch.long, torch.int, torch.int32, torch.int64):
+                y = F.embedding(x, w)  # token_embd: indices in, rows gathered from effective weight
+            else:
+                y = F.linear(x, w.to(x.dtype), bias)
+            ctx.save_for_backward(x, s, codes, d0, hb, signs, bias, pr, pc)
+            ctx.rows = rows
+            ctx.nblk_w = nblk_w
+            return y
+
+        @staticmethod
+        def backward(ctx, gy):
+            x, s, codes, d0, hb, signs, bias, pr, pc = ctx.saved_tensors
+            rows, nblk_w = ctx.rows, ctx.nblk_w
+            width = d0.shape[1] * 128
+            w = _ScaledWFn.build_w(s, codes, d0, hb, signs, rows, nblk_w)
+            if pr is not None:
+                w = w[pr]
+            if pc is not None:
+                w = w[:, pc]
+            if x.dtype in (torch.long, torch.int, torch.int32, torch.int64):
+                gw = torch.zeros(rows, width, device=gy.device, dtype=gy.dtype)
+                gw.index_add_(0, x.reshape(-1), gy.reshape(-1, width))
+                gx = None
+            else:
+                gx = (gy @ w.to(gy.dtype)) if ctx.needs_input_grad[0] else None
+                gw = gy.reshape(-1, gy.shape[-1]).T @ x.reshape(-1, x.shape[-1])
+            # un-permute the gradient back to stored order (inverse of the forward gathers)
+            if pc is not None:
+                gw_t = torch.empty_like(gw)
+                gw_t[:, pc] = gw
+                gw = gw_t
+            if pr is not None:
+                gw_t = torch.empty_like(gw)
+                gw_t[pr] = gw
+                gw = gw_t
+            if hb is not None:
+                gsz = hb.shape[0]
+                g1 = ((gw * signs).reshape(rows, width // gsz, gsz) @ hb.T).reshape(rows, width)
+            else:
+                g1 = gw
+            # stored[r,k] = (q-1)[r,k] * (d0*s)[r,b]  ->  dL/ds[r,b] = sum_k g1[r,k] (q-1)[r,k] d0[r,b]
+            qs = unpack(codes).float() - 1.0
+            gs = (qs * g1.view(rows, nblk_w, 128).float()).sum(-1) * d0.float()
+            gs = gs.unsqueeze(-1).to(s.dtype)
+            return gx, gs, None, None, None, None, None, None, None, None, None
+
+    def untile_perms(name):
+        """Row/col permutations reproducing GgufCheckpoint._hf_order for this tensor."""
+        if ck._heads.rep == 1 or not name.startswith("blk."):
+            return None, None
+        kind = name.split(".")[2]
+        rows, width = g.tensors[name].shape[1], g.tensors[name].shape[0]
+        pr = pc = None
+        if kind == "attn_qkv":
+            qk = 2 * ck._heads.n_k * ck._heads.k_dim
+            idx = (torch.arange(rows - qk) + qk).reshape(rows - qk, 1)
+            pr = torch.cat([torch.arange(qk), _untile(idx, 0, ck._heads, ck._heads.v_dim).reshape(-1)])
+        elif kind == "attn_gate":
+            pr = _untile(torch.arange(rows).reshape(rows, 1), 0, ck._heads, ck._heads.v_dim).reshape(-1)
+        elif kind in ("ssm_alpha", "ssm_beta"):
+            pr = _untile(torch.arange(rows).reshape(rows, 1), 0, ck._heads, 1).reshape(-1)
+        elif kind == "ssm_out":
+            folded = rotation(g, name) is not None
+            if not (folded and ck._grouped):
+                pc = _untile(torch.arange(width).reshape(1, width), 1, ck._heads, ck._heads.v_dim).reshape(-1)
+        if pr is not None:
+            pr = pr.to(dev)
+        if pc is not None:
+            pc = pc.to(dev)
+        return pr, pc
+
+    king_path = next(Path(args.gguf).glob("*.gguf"))
+    fk = open(king_path, "rb")
+    dequant_checked = set()
+    for gg_name, hf_key in gguf_to_hf.items():
+        mod_name = hf_key.rsplit(".weight", 1)[0]
+        try:
+            mod = model.get_submodule(mod_name)
+        except Exception:
+            continue
+        rows, width = sd[hf_key].shape  # tile() preserves shape: stored-space layout == HF shape
+        nblk_w = width // 128
+        if nblk_w == 0 or rows == 0:
+            continue
+        signs = rot_signs(gg_name)
+        t = g.tensors[gg_name]
+        fk.seek(g.data_start + t.offset)
+        raw = np.frombuffer(fk.read(rows * nblk_w * 34), np.uint8).reshape(rows * nblk_w, 34)
+        d0 = torch.from_numpy(raw[:, :2].copy().view("<f2")).view(rows, nblk_w).to(dev)
+        codes = torch.from_numpy(raw[:, 2:].copy()).view(rows, nblk_w, 32).to(dev)
+        # one-time cross-check: direct dequant must match the surrogate transform path
+        kind = "rot" if signs is not None else "plain"
+        if kind not in dequant_checked:
+            dequant_checked.add(kind)
+            w_mine = _ScaledWFn.build_w(torch.ones(rows, nblk_w, 1, device=dev), codes, d0,
+                                        None, None, rows, nblk_w)
+            w_ref, _ = hf_to_stored(gg_name, sd[hf_key].float().to(dev))
+            rel = (w_mine.float() - w_ref.float()).abs().mean().item() / w_ref.float().abs().mean().item()
+            print(f"DEQUANT_CHECK[{kind}] {gg_name}: rel_mean_diff={rel:.3e}", flush=True)
+            assert rel < 0.02, f"direct dequant mismatch on {gg_name} ({kind}): {rel}"
+        s = torch.nn.Parameter(torch.ones(rows, nblk_w, 1, device=dev))
+        # hadamard entries and rotation signs are +-1: exact in bf16
+        hb = H.to(torch.bfloat16) if signs is not None else None
+        if signs is not None:
+            signs = signs.to(torch.bfloat16)
+        pr, pc = untile_perms(gg_name)
+
+        def make_fwd(mod=mod, codes=codes, signs=signs, hb=hb, rows=rows, nblk_w=nblk_w,
+                     s=s, d0=d0, bias=getattr(mod, "bias", None), pr=pr, pc=pc):
+            def fwd(x):
+                return _ScaledWFn.apply(x, s, codes, d0, hb, signs, bias, rows, nblk_w, pr, pc)
+            return fwd
+
+        mod.forward = make_fwd()
+        with torch.no_grad():
+            mod.weight = torch.nn.Parameter(torch.empty(0, device=dev), requires_grad=False)
+        scale_params[gg_name] = (s, rows, nblk_w)
+
+    # inner norms trainable (exclude output/final norm: handled by GgufHead from file)
+    norm_params = []
+    for n_, p_ in model.named_parameters():
+        if n_.endswith("norm.weight") and "model.norm" not in n_ and not n_.endswith(".model.norm.weight"):
+            p_.requires_grad_(True)
+            norm_params.append((n_, p_))
+    trained_norm_names = {n_ for n_, _ in norm_params}
+    all_p = [v[0] for v in scale_params.values()] + [p_ for _, p_ in norm_params]
+    print(f"wrapped {len(scale_params)} linears "
+          f"({sum(v[0].numel() for v in scale_params.values()):,} scale params), "
+          f"{len(norm_params)} inner norms", flush=True)
+
+    # ---- heads ----
+    print("WRAP_DONE", flush=True)
+    teacher_head = FinalHead(Path(args.teacher), dev)
+    own_head = GgufHead(Path(args.gguf), dev)
+    LAYERS = list(config.LAYERS)  # 8..64 step 8
+
+    # ---- data: teacher states + teacher logprobs cached ----
+    records = json.loads(Path(args.samples).read_text())[: args.max_states]
+    data = []
+    from safetensors.torch import load_file
+    for r in records:
+        st = load_file(r["states"])["hidden_states"].float()      # [rows, 8, hidden]
+        data.append({
+            "ids": torch.tensor(r["prompt"] + r["answer"], device=dev).unsqueeze(0),
+            "n_prompt": len(r["prompt"]),
+            "n_answer": len(r["answer"]),
+            "teacher_states": st.cpu(),
+        })
+        if len(data) >= args.max_states:
+            break
+    print(f"{len(data)} trajectory samples", flush=True)
+    import gc
+    try:
+        del sd
+    except NameError:
+        pass
+    gc.collect()
+    torch.cuda.empty_cache()
+    print("MEM after setup:", torch.cuda.memory_allocated() // (1 << 20), "MB", flush=True)
+    print("DATA_DONE mem:", torch.cuda.memory_allocated() // (1 << 20), "MB", flush=True)
+
+    opt = torch.optim.Adam(all_p, lr=args.lr)
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.steps)
+    model.train()  # gradient checkpointing (enabled above) only activates in training mode
+    step = 0
+    while step < args.steps:
+        for d in data:
+            ids = d["ids"]
+            out = model.model(ids, output_hidden_states=True, use_cache=False)  # skip lm_head + cache: loss uses hidden states only
+            hs = out.hidden_states                       # tuple: [emb, L1..L64, post-norm]
+            n_ans = d["n_answer"]
+            start = d["n_prompt"] - 1                    # row0 = last prompt token
+            rows = slice(start, start + n_ans)
+
+            own_last = torch.stack([hs[L][0, start:start + n_ans] for L in LAYERS], 1)  # [n,8,h]
+            t_st = d["teacher_states"].to(dev)
+            # states files hold rows 0..start+n_ans-1 of the sequence; the eval window starts at
+            # the last prompt row (metrics.py compare: "rows that read or predict answer tokens")
+            t_win = t_st[start:start + n_ans]
+
+            # KL from last captured layer via own head
+            ref_last = t_win[:, -1].float()
+            kl = torch.zeros((), device=dev)
+            CH = 256
+            for c0 in range(0, ref_last.shape[0], CH):
+                own_lp_c = own_head.logprobs(own_last[c0:c0 + CH, -1].float())
+                with torch.no_grad():
+                    ref_lp_c = teacher_head.logprobs(ref_last[c0:c0 + CH])
+                    p_c = ref_lp_c.exp()
+                kl = kl + (p_c * (ref_lp_c - own_lp_c)).sum(-1).sum()
+                del ref_lp_c, own_lp_c, p_c
+            kl = kl / ref_last.shape[0]
+
+            # dist over captured layers, rows 1.. (answer token rows; matches metrics.py)
+            a = t_win[1:].float()
+            b = own_last[1:].float()
+            dist = ((b - a).norm(dim=-1) / (a.norm(dim=-1) + 1e-9)).mean()
+
+            loss = kl + args.dist_weight * dist
+            opt.zero_grad(set_to_none=True)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(all_p, 1.0)
+            opt.step()
+            step += 1
+            sched.step()
+            if step % 20 == 0:
+                print(f"step {step}: KL {kl.item():.5f}  dist {dist.item():.5f}", flush=True)
+            if step >= args.steps:
+                break
+
+    # ---- export ----
+    import shutil
+    king = next(Path(args.gguf).glob("*.gguf"))
+    out_dir = Path(args.out)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    dst = out_dir / king.name
+    shutil.copy(king, dst)
+    g_out = read_header(dst)
+    with open(dst, "r+b") as fh, open(king, "rb") as fk:
+        for gg_name, (s, rows, nblk_w) in scale_params.items():
+            t = g_out.tensors[gg_name]
+            nblk_flat = int(np.prod(t.shape)) // 128
+            fk.seek(g_out.data_start + t.offset)
+            blocks = np.frombuffer(fk.read(nblk_flat * 34), np.uint8).reshape(-1, 34).copy()
+            d0 = blocks[:, :2].copy().view("<f2").astype(np.float32).reshape(rows, nblk_w)
+            s_np = s.detach().cpu().numpy().reshape(rows, nblk_w)
+            blocks[:, :2] = (d0 * s_np).astype(np.float16).astype("<f2").view(np.uint8).reshape(-1, 2)
+            fh.seek(g_out.data_start + t.offset)
+            fh.write(blocks.tobytes())
+        named = dict(model.named_parameters())
+        n_norm = 0
+        for gg_name, t in g_out.tensors.items():
+            if not gg_name.endswith("norm.weight") or gg_name == "output_norm.weight":
+                continue
+            parts = gg_name.split(".")
+            suffix = (f"layers.{parts[1]}.{parts[2]}.weight" if gg_name.startswith("blk.")
+                      else None)
+            if suffix is None:
+                continue
+            cand = [n_ for n_ in named if n_.endswith(suffix) and n_ in trained_norm_names]
+            if not cand:
+                continue
+            v = named[cand[0]].detach().float().cpu().numpy() + 1.0
+            if v.size != int(np.prod(t.shape)):
+                continue
+            fh.seek(g_out.data_start + t.offset)
+            fh.write(v.astype("<f4").tobytes())
+            n_norm += 1
+    print(f"exported scales + {n_norm} norms -> {dst}")
+
+
+if __name__ == "__main__":
+    main()