Nadir v2: element head, inner-page data, runtime fixes

864f427 verified 21 days ago

16.1 kB

#!/usr/bin/env python3
"""Calibrated weight-only quantization of NadirNet into ONNX MatMulNBits (8-bit and 4-bit).
Layer-sequential GPTQ: each layer's Hessian comes from real crawled pages flowing through the already
quantized earlier layers. Each group of 32 weights gets asymmetric uint codes whose clipping range is
chosen by an MSE search. Codes are packed in ORT's MatMulNBits layout, and every variant is checked
three ways: against a PyTorch fake-quant copy (packing parity), against fp32 (fidelity) and against
the Haiku targets (quality)."""
import argparse, copy, json, os, pickle, time
import numpy as np
import torch
import torch.nn as nn

from model import NadirNet
from train import prep
from export import feeds_for, elem_feeds, fidelity, bench, NAMES_IN

GS = 32


class Stop(Exception):
    pass


def group_params(w, bits, maxshrink, grid=100):
    """Per-row scale/zero-point for one group (rows, GS), minimizing |q(w)-w|^2.4 over clipping ranges."""
    maxq = 2 ** bits - 1
    xmin = w.min(1).values.clamp(max=0)
    xmax = w.max(1).values.clamp(min=0)
    best = torch.full_like(xmin, float("inf"))
    bs, bz = torch.ones_like(xmin), torch.zeros_like(xmin)
    for i in range(int(maxshrink * grid) + 1):
        p = 1 - i / grid
        lo, hi = p * xmin, p * xmax
        s = ((hi - lo) / maxq).clamp(min=1e-10)
        z = torch.round(-lo / s).clamp(0, maxq)
        q = (torch.clamp(torch.round(w / s[:, None]) + z[:, None], 0, maxq) - z[:, None]) * s[:, None]
        err = (q - w).abs().pow(2.4).sum(1)
        m = err < best
        best[m], bs[m], bz[m] = err[m], s[m], z[m]
    return bs, bz


def gptq(W, H, bits, percdamp=0.01, blocksize=128, gs=GS):
    """W (out, Kpad) with Kpad % gs == 0, H (Kpad, Kpad); gs=None quantizes per output channel (one group per row).
    Returns dequantized W, codes, scales, zero points."""
    maxq = 2 ** bits - 1
    maxshrink = 0.5 if bits == 4 else 0.2
    W = W.clone().double()
    H = H.clone().double()
    rows, cols = W.shape
    dead = torch.diag(H) == 0
    H[dead, dead] = 1
    W[:, dead] = 0
    H += percdamp * torch.mean(torch.diag(H)) * torch.eye(cols, dtype=H.dtype)
    Hinv = torch.linalg.cholesky(torch.cholesky_inverse(torch.linalg.cholesky(H)), upper=True)
    Q = torch.zeros_like(W)
    codes = torch.zeros(rows, cols, dtype=torch.uint8)
    if gs is None:  # per channel: one range per output row, chosen on the whole row up front
        s, z = group_params(W, bits, maxshrink)
        S, Z = s[:, None].clone(), z[:, None].clone()
    else:
        S = torch.zeros(rows, cols // gs, dtype=torch.float64)
        Z = torch.zeros(rows, cols // gs, dtype=torch.float64)
    for i1 in range(0, cols, blocksize):
        i2 = min(i1 + blocksize, cols)
        W1 = W[:, i1:i2].clone()
        E1 = torch.zeros_like(W1)
        H1 = Hinv[i1:i2, i1:i2]
        for i in range(i2 - i1):
            c = i1 + i
            if gs is not None and c % gs == 0:  # groups never straddle blocks (blocksize % gs == 0)
                s, z = group_params(W1[:, i:i + gs], bits, maxshrink)
                S[:, c // gs], Z[:, c // gs] = s, z
            w = W1[:, i]
            qc = torch.clamp(torch.round(w / s) + z, 0, maxq)
            q = (qc - z) * s
            codes[:, c] = qc.to(torch.uint8)
            Q[:, c] = q
            err = (w - q) / H1[i, i]
            W1[:, i:] -= err[:, None] * H1[i, i:][None, :]
            E1[:, i] = err
        W[:, i2:] -= E1 @ Hinv[i1:i2, i2:]
    return Q.float(), codes, S.float(), Z.to(torch.uint8)


def pack(codes, S, Z, bits):
    """MatMulNBits layout: B (N, k_blocks, blob), scales (N, k_blocks), zero points (N, ceil(k_blocks*bits/8))."""
    c = codes.numpy()
    N, Kp = c.shape
    kb = Kp // GS
    z = Z.numpy().astype(np.uint8)
    if bits == 8:
        B = c.reshape(N, kb, GS).astype(np.uint8)
        zp = z
    else:
        c = c.reshape(N, kb, GS)
        B = (c[..., 0::2] | (c[..., 1::2] << 4)).astype(np.uint8)
        if kb % 2:
            z = np.concatenate([z, np.full((N, 1), 8, np.uint8)], 1)
        zp = (z[:, 0::2] | (z[:, 1::2] << 4)).astype(np.uint8)
    return B, S.numpy().astype(np.float32), zp


def linear_order(model, sample):
    """Linear modules in the order the forward pass calls them."""
    order, hooks = [], []
    for name, mod in model.named_modules():
        if isinstance(mod, nn.Linear):
            hooks.append(mod.register_forward_pre_hook(lambda m, a, n=name: order.append(n) if n not in order else None))
    with torch.no_grad():
        model(*[torch.from_numpy(sample[k]) for k in NAMES_IN])
    for h in hooks:
        h.remove()
    return order


def onnx_matmuls(model, onnx_model):
    """Map each Linear to the MatMul node whose initializer is its transposed weight (Gemm exports stay fp32)."""
    from onnx import numpy_helper
    inits = {i.name: i for i in onnx_model.graph.initializer}
    by_shape = {}
    for node in onnx_model.graph.node:
        if node.op_type == "MatMul" and node.input[1] in inits:
            arr = numpy_helper.to_array(inits[node.input[1]])
            by_shape.setdefault(arr.shape, []).append((node, arr))
    out = {}
    for name, mod in model.named_modules():
        if isinstance(mod, nn.Linear):
            wt = mod.weight.detach().numpy().T
            for node, arr in by_shape.get(wt.shape, []):
                if np.array_equal(arr, wt):
                    out[name] = node
                    break
    return out


def hessian(model, name, calib):
    """2/n * X^T X of the inputs to one Linear, over all real token rows of the calibration pages."""
    mod = dict(model.named_modules())[name]
    K = mod.in_features
    H = torch.zeros(K, K, dtype=torch.float64)
    n = [0]

    def hook(m, a):
        x = a[0].detach().reshape(-1, K).double()
        H.add_(x.T @ x)
        n[0] += x.shape[0]
        raise Stop

    h = mod.register_forward_pre_hook(hook)
    with torch.no_grad():
        for f in calib:
            try:
                model(*[torch.from_numpy(f[k]) for k in NAMES_IN])
            except Stop:
                pass
    h.remove()
    return H * (2.0 / max(1, n[0]))


def quantize_model(model, calib, bits_for, order):
    """Sequential GPTQ over `order`; returns the fake-quant model and packed tensors per layer."""
    qm = copy.deepcopy(model).eval()
    mods = dict(qm.named_modules())
    packed = {}
    for name in order:
        bits = bits_for.get(name)
        if not bits:
            continue
        mod = mods[name]
        K = mod.in_features
        H = hessian(qm, name, calib)
        t = time.time()
        # blocks of 32 along K (zero-padded to a multiple of 32)
        Kp = -(-K // GS) * GS
        Hp = torch.zeros(Kp, Kp, dtype=torch.float64); Hp[:K, :K] = H
        Wp = torch.zeros(mod.out_features, Kp); Wp[:, :K] = mod.weight.detach()
        Q, codes, S, Z = gptq(Wp, Hp, bits)
        if bits == 8:
            # 8-bit: uint8 codes as a (K/32, 32, N) tensor with per-block scale and zero point, dequantized in-graph
            N, kb = mod.out_features, Kp // GS
            packed[name] = ("dq8", K, codes.numpy().T.reshape(kb, GS, N).copy(), S.numpy().T[:, None, :].astype(np.float32).copy(),
                            Z.numpy().T[:, None, :].astype(np.uint8).copy())
        else:
            # 4-bit: packed as MatMulNBits
            packed[name] = ("nbits", K) + pack(codes, S, Z, bits)
        with torch.no_grad():
            mod.weight.copy_(Q[:, :K])
        print(f"  {name} {tuple(mod.weight.shape)} {bits}-bit  {time.time() - t:.1f}s", flush=True)
    return qm, packed


def build_onnx(fp32_path, mapping, packed, out_path, accuracy_level=0):
    import onnx
    from onnx import helper, numpy_helper
    m = onnx.load(fp32_path)
    g = m.graph
    drop_inits = set()
    for name, (kind, K, B, S, zp) in packed.items():
        node = mapping[name]
        wname = node.input[1]
        idx = list(g.node).index(node)
        if kind == "dq8":
            # uint8 codes dequantized with plain Cast/Sub/Mul: every onnxruntime constant-folds this into exact fp32
            # weights at load. (DequantizeLinear->MatMul gets rewritten by onnxruntime into an activation-quantized
            # kernel, which costs accuracy.)
            kb, gs, N = B.shape
            inits = [numpy_helper.from_array(B, wname + "_q8"), numpy_helper.from_array(zp, wname + "_zero_point"),
                     numpy_helper.from_array(S, wname + "_scale"), numpy_helper.from_array(np.array([kb * gs, N], np.int64), wname + "_kn")]
            g.initializer.extend(inits)
            base = (node.name or wname) + "_deq8"
            F = onnx.TensorProto.FLOAT
            deq = [helper.make_node("Cast", [inits[0].name], [wname + "_f"], name=base + "_cast", to=F),
                   helper.make_node("Cast", [inits[1].name], [wname + "_zf"], name=base + "_castz", to=F),
                   helper.make_node("Sub", [wname + "_f", wname + "_zf"], [wname + "_c"], name=base + "_sub"),
                   helper.make_node("Mul", [wname + "_c", inits[2].name], [wname + "_d3"], name=base + "_mul"),
                   helper.make_node("Reshape", [wname + "_d3", inits[3].name], [wname + ("_dq" if kb * gs == K else "_dp")], name=base + "_reshape")]
            if kb * gs != K:  # drop the zero padding rows
                lims = [numpy_helper.from_array(np.array(v, np.int64), wname + n) for v, n in (([0], "_s0"), ([K], "_sK"), ([0], "_ax"))]
                g.initializer.extend(lims)
                deq.append(helper.make_node("Slice", [wname + "_dp"] + [i.name for i in lims], [wname + "_dq"], name=base + "_slice"))
            g.node[idx].input[1] = wname + "_dq"  # `node` belongs to another copy of the graph: edit this one
            for k, dn in enumerate(deq):
                g.node.insert(idx + k, dn)
        else:
            N = B.shape[0]
            inits = [numpy_helper.from_array(B, wname + "_Q4"), numpy_helper.from_array(S, wname + "_scales"),
                     numpy_helper.from_array(zp, wname + "_zero_points")]
            g.initializer.extend(inits)
            kw = dict(K=K, N=N, bits=4, block_size=GS)
            if accuracy_level:
                kw["accuracy_level"] = accuracy_level
            new = helper.make_node("MatMulNBits", [node.input[0]] + [i.name for i in inits], list(node.output),
                                   name=(node.name or wname) + "_Q4", domain="com.microsoft", **kw)
            g.node.remove(node)
            g.node.insert(idx, new)
        drop_inits.add(wname)
    keep = [i for i in g.initializer if i.name not in drop_inits]
    del g.initializer[:]
    g.initializer.extend(keep)
    if any(n.domain == "com.microsoft" for n in g.node) and not any(o.domain == "com.microsoft" for o in m.opset_import):
        m.opset_import.append(helper.make_opsetid("com.microsoft", 1))
    onnx.checker.check_model(m)
    onnx.save(m, out_path)


def quality(sess, samples, esess=None, ref=None):
    """dE vs the training targets (Haiku for src 0, native dark for src 1); with the element graph, how often the
    element decisions computed from this page graph's outputs match those from the fp32 page graph."""
    agg = {0: [0.0, 0.0], 1: [0.0, 0.0]}
    ok = n = 0
    for s in samples:
        f = feeds_for(s)
        lab, _, ctx = sess.run(None, f)
        k = len(s["role"])
        m = s["tmask"][:k]
        if m.any():
            d = np.linalg.norm(lab[0, :k] - s["lab_t"][:k], axis=-1)
            w = s["w"][:k] * m
            agg[int(s["src"])][0] += float((d * w).sum() / max(w.sum(), 1e-9)); agg[int(s["src"])][1] += 1
        if esess is not None and ref is not None and len(s["E_ut"]):
            rl, _, rc = ref.run(None, f)
            a = esess.run(None, elem_feeds(s, lab, ctx)[0])[0][0].argmax(-1)
            b = esess.run(None, elem_feeds(s, rl, rc)[0])[0][0].argmax(-1)
            ok += int((a == b).sum()); n += len(a)
    out = {"haiku_dE": agg[0][0] / max(1, agg[0][1]), "native_dE": agg[1][0] / max(1, agg[1][1])}
    if n:
        out["elem_decisions_same_as_fp32"] = ok / n
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--data", required=True)
    ap.add_argument("--fp32", required=True, help="fp32 ONNX from export.py")
    ap.add_argument("--out", required=True)
    ap.add_argument("--calib", type=int, default=256)
    ap.add_argument("--eval", type=int, default=400)
    ap.add_argument("--threads", type=int, default=8)
    ap.add_argument("--elem", default="", help="element graph (fp32) to measure decision agreement")
    args = ap.parse_args()
    torch.set_num_threads(args.threads)
    os.makedirs(args.out, exist_ok=True)
    import onnx, onnxruntime as ort
    ck = torch.load(args.ckpt, map_location="cpu")
    model = NadirNet(**ck["cfg"]); model.load_state_dict(ck["model"]); model.eval()
    data = pickle.load(open(args.data, "rb"))
    hv = [s for s in data["val"] if s["src"] == 0]; nv = [s for s in data["val"] if s["src"] == 1]
    val = [prep(s) for s in hv[: args.eval * 3 // 4] + nv[: args.eval - min(len(hv), args.eval * 3 // 4)]]  # both sources
    rng = np.random.default_rng(0)
    pick = rng.choice(len(data["train"]), size=min(args.calib, len(data["train"])), replace=False)
    calib = [feeds_for(prep(data["train"][i])) for i in pick]
    order = linear_order(model, calib[0])
    mapping = onnx_matmuls(model, onnx.load(args.fp32))
    quantizable = [n for n in order if n in mapping and not n.endswith(".rel")]
    print(f"{len(quantizable)}/{len(order)} linears map to MatMul; fp32 kept: {[n for n in order if n not in quantizable]}", flush=True)
    edges = [n for n in quantizable if n.startswith(("tok_in.", "img_in.", "color_head.", "ctx_head."))]
    plans = {"int8": {n: 8 for n in quantizable},
             "int4": {n: (8 if n in edges else 4) for n in quantizable}}
    so = ort.SessionOptions(); so.intra_op_num_threads = 4
    ref = ort.InferenceSession(args.fp32, so, providers=["CPUExecutionProvider"])
    report = {"format": {"8-bit layers": f"uint8 blocks of {GS} with fp32 scale + uint8 zero point, dequantized in-graph (Cast/Sub/Mul, folded at load)", "4-bit layers": f"uint4 blocks of {GS}, MatMulNBits"},
              "group_size": GS, "calib_pages": len(calib), "eval_pages": len(val), "fp32_layers": [n for n in order if n not in quantizable],
              "int4_8bit_edge_layers": edges, "variants": {}}
    esess = ort.InferenceSession(args.elem, so, providers=["CPUExecutionProvider"]) if args.elem else None
    bfeeds = [feeds_for(s) for s in val[:20]]
    r = quality(ref, val); r["size_mb"] = round(os.path.getsize(args.fp32) / 1e6, 2); r["cpu_ms_median"] = round(bench(ref, bfeeds), 2)
    report["variants"]["fp32"] = r
    print("fp32", r, flush=True)
    for vname, bits_for in plans.items():
        print(f"== {vname}", flush=True)
        qm, packed = quantize_model(model, calib, bits_for, quantizable)
        path = os.path.join(args.out, f"nadir-{vname}.onnx")
        build_onnx(args.fp32, mapping, packed, path)
        sess = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
        # packing parity: ORT MatMulNBits graph vs the PyTorch fake-quant model
        par = 0.0
        with torch.no_grad():
            for s in val[:8]:
                f = feeds_for(s)
                tq = qm(*[torch.from_numpy(f[k]) for k in NAMES_IN])[0].numpy()
                par = max(par, float(np.abs(tq - sess.run(None, f)[0]).max()))
        r = fidelity(sess, ref, val)
        r.update(quality(sess, val, esess, ref))
        r["pack_parity_max_abs"] = par
        r["size_mb"] = round(os.path.getsize(path) / 1e6, 2)
        r["cpu_ms_median"] = round(bench(sess, bfeeds), 2)
        report["variants"][vname] = r
        print(vname, r, flush=True)
        torch.save({"cfg": ck["cfg"], "model": qm.state_dict()}, os.path.join(args.out, f"fakequant-{vname}.pt"))
    json.dump(report, open(os.path.join(args.out, "quant_report.json"), "w"), indent=1)


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""Calibrated weight-only quantization of NadirNet into ONNX MatMulNBits (8-bit and 4-bit).

	Layer-sequential GPTQ: each layer's Hessian comes from real crawled pages flowing through the already
	quantized earlier layers. Each group of 32 weights gets asymmetric uint codes whose clipping range is
	chosen by an MSE search. Codes are packed in ORT's MatMulNBits layout, and every variant is checked
	three ways: against a PyTorch fake-quant copy (packing parity), against fp32 (fidelity) and against
	the Haiku targets (quality)."""
	import argparse, copy, json, os, pickle, time
	import numpy as np
	import torch
	import torch.nn as nn

	from model import NadirNet
	from train import prep
	from export import feeds_for, elem_feeds, fidelity, bench, NAMES_IN

	GS = 32


	class Stop(Exception):
	pass


	def group_params(w, bits, maxshrink, grid=100):
	"""Per-row scale/zero-point for one group (rows, GS), minimizing \|q(w)-w\|^2.4 over clipping ranges."""
	maxq = 2 ** bits - 1
	xmin = w.min(1).values.clamp(max=0)
	xmax = w.max(1).values.clamp(min=0)
	best = torch.full_like(xmin, float("inf"))
	bs, bz = torch.ones_like(xmin), torch.zeros_like(xmin)
	for i in range(int(maxshrink * grid) + 1):
	p = 1 - i / grid
	lo, hi = p * xmin, p * xmax
	s = ((hi - lo) / maxq).clamp(min=1e-10)
	z = torch.round(-lo / s).clamp(0, maxq)
	q = (torch.clamp(torch.round(w / s[:, None]) + z[:, None], 0, maxq) - z[:, None]) * s[:, None]
	err = (q - w).abs().pow(2.4).sum(1)
	m = err < best
	best[m], bs[m], bz[m] = err[m], s[m], z[m]
	return bs, bz


	def gptq(W, H, bits, percdamp=0.01, blocksize=128, gs=GS):
	"""W (out, Kpad) with Kpad % gs == 0, H (Kpad, Kpad); gs=None quantizes per output channel (one group per row).
	Returns dequantized W, codes, scales, zero points."""
	maxq = 2 ** bits - 1
	maxshrink = 0.5 if bits == 4 else 0.2
	W = W.clone().double()
	H = H.clone().double()
	rows, cols = W.shape
	dead = torch.diag(H) == 0
	H[dead, dead] = 1
	W[:, dead] = 0
	H += percdamp * torch.mean(torch.diag(H)) * torch.eye(cols, dtype=H.dtype)
	Hinv = torch.linalg.cholesky(torch.cholesky_inverse(torch.linalg.cholesky(H)), upper=True)
	Q = torch.zeros_like(W)
	codes = torch.zeros(rows, cols, dtype=torch.uint8)
	if gs is None: # per channel: one range per output row, chosen on the whole row up front
	s, z = group_params(W, bits, maxshrink)
	S, Z = s[:, None].clone(), z[:, None].clone()
	else:
	S = torch.zeros(rows, cols // gs, dtype=torch.float64)
	Z = torch.zeros(rows, cols // gs, dtype=torch.float64)
	for i1 in range(0, cols, blocksize):
	i2 = min(i1 + blocksize, cols)
	W1 = W[:, i1:i2].clone()
	E1 = torch.zeros_like(W1)
	H1 = Hinv[i1:i2, i1:i2]
	for i in range(i2 - i1):
	c = i1 + i
	if gs is not None and c % gs == 0: # groups never straddle blocks (blocksize % gs == 0)
	s, z = group_params(W1[:, i:i + gs], bits, maxshrink)
	S[:, c // gs], Z[:, c // gs] = s, z
	w = W1[:, i]
	qc = torch.clamp(torch.round(w / s) + z, 0, maxq)
	q = (qc - z) * s
	codes[:, c] = qc.to(torch.uint8)
	Q[:, c] = q
	err = (w - q) / H1[i, i]
	W1[:, i:] -= err[:, None] * H1[i, i:][None, :]
	E1[:, i] = err
	W[:, i2:] -= E1 @ Hinv[i1:i2, i2:]
	return Q.float(), codes, S.float(), Z.to(torch.uint8)


	def pack(codes, S, Z, bits):
	"""MatMulNBits layout: B (N, k_blocks, blob), scales (N, k_blocks), zero points (N, ceil(k_blocks*bits/8))."""
	c = codes.numpy()
	N, Kp = c.shape
	kb = Kp // GS
	z = Z.numpy().astype(np.uint8)
	if bits == 8:
	B = c.reshape(N, kb, GS).astype(np.uint8)
	zp = z
	else:
	c = c.reshape(N, kb, GS)
	B = (c[..., 0::2] \| (c[..., 1::2] << 4)).astype(np.uint8)
	if kb % 2:
	z = np.concatenate([z, np.full((N, 1), 8, np.uint8)], 1)
	zp = (z[:, 0::2] \| (z[:, 1::2] << 4)).astype(np.uint8)
	return B, S.numpy().astype(np.float32), zp


	def linear_order(model, sample):
	"""Linear modules in the order the forward pass calls them."""
	order, hooks = [], []
	for name, mod in model.named_modules():
	if isinstance(mod, nn.Linear):
	hooks.append(mod.register_forward_pre_hook(lambda m, a, n=name: order.append(n) if n not in order else None))
	with torch.no_grad():
	model(*[torch.from_numpy(sample[k]) for k in NAMES_IN])
	for h in hooks:
	h.remove()
	return order


	def onnx_matmuls(model, onnx_model):
	"""Map each Linear to the MatMul node whose initializer is its transposed weight (Gemm exports stay fp32)."""
	from onnx import numpy_helper
	inits = {i.name: i for i in onnx_model.graph.initializer}
	by_shape = {}
	for node in onnx_model.graph.node:
	if node.op_type == "MatMul" and node.input[1] in inits:
	arr = numpy_helper.to_array(inits[node.input[1]])
	by_shape.setdefault(arr.shape, []).append((node, arr))
	out = {}
	for name, mod in model.named_modules():
	if isinstance(mod, nn.Linear):
	wt = mod.weight.detach().numpy().T
	for node, arr in by_shape.get(wt.shape, []):
	if np.array_equal(arr, wt):
	out[name] = node
	break
	return out


	def hessian(model, name, calib):
	"""2/n * X^T X of the inputs to one Linear, over all real token rows of the calibration pages."""
	mod = dict(model.named_modules())[name]
	K = mod.in_features
	H = torch.zeros(K, K, dtype=torch.float64)
	n = [0]

	def hook(m, a):
	x = a[0].detach().reshape(-1, K).double()
	H.add_(x.T @ x)
	n[0] += x.shape[0]
	raise Stop

	h = mod.register_forward_pre_hook(hook)
	with torch.no_grad():
	for f in calib:
	try:
	model(*[torch.from_numpy(f[k]) for k in NAMES_IN])
	except Stop:
	pass
	h.remove()
	return H * (2.0 / max(1, n[0]))


	def quantize_model(model, calib, bits_for, order):
	"""Sequential GPTQ over `order`; returns the fake-quant model and packed tensors per layer."""
	qm = copy.deepcopy(model).eval()
	mods = dict(qm.named_modules())
	packed = {}
	for name in order:
	bits = bits_for.get(name)
	if not bits:
	continue
	mod = mods[name]
	K = mod.in_features
	H = hessian(qm, name, calib)
	t = time.time()
	# blocks of 32 along K (zero-padded to a multiple of 32)
	Kp = -(-K // GS) * GS
	Hp = torch.zeros(Kp, Kp, dtype=torch.float64); Hp[:K, :K] = H
	Wp = torch.zeros(mod.out_features, Kp); Wp[:, :K] = mod.weight.detach()
	Q, codes, S, Z = gptq(Wp, Hp, bits)
	if bits == 8:
	# 8-bit: uint8 codes as a (K/32, 32, N) tensor with per-block scale and zero point, dequantized in-graph
	N, kb = mod.out_features, Kp // GS
	packed[name] = ("dq8", K, codes.numpy().T.reshape(kb, GS, N).copy(), S.numpy().T[:, None, :].astype(np.float32).copy(),
	Z.numpy().T[:, None, :].astype(np.uint8).copy())
	else:
	# 4-bit: packed as MatMulNBits
	packed[name] = ("nbits", K) + pack(codes, S, Z, bits)
	with torch.no_grad():
	mod.weight.copy_(Q[:, :K])
	print(f" {name} {tuple(mod.weight.shape)} {bits}-bit {time.time() - t:.1f}s", flush=True)
	return qm, packed


	def build_onnx(fp32_path, mapping, packed, out_path, accuracy_level=0):
	import onnx
	from onnx import helper, numpy_helper
	m = onnx.load(fp32_path)
	g = m.graph
	drop_inits = set()
	for name, (kind, K, B, S, zp) in packed.items():
	node = mapping[name]
	wname = node.input[1]
	idx = list(g.node).index(node)
	if kind == "dq8":
	# uint8 codes dequantized with plain Cast/Sub/Mul: every onnxruntime constant-folds this into exact fp32
	# weights at load. (DequantizeLinear->MatMul gets rewritten by onnxruntime into an activation-quantized
	# kernel, which costs accuracy.)
	kb, gs, N = B.shape
	inits = [numpy_helper.from_array(B, wname + "_q8"), numpy_helper.from_array(zp, wname + "_zero_point"),
	numpy_helper.from_array(S, wname + "_scale"), numpy_helper.from_array(np.array([kb * gs, N], np.int64), wname + "_kn")]
	g.initializer.extend(inits)
	base = (node.name or wname) + "_deq8"
	F = onnx.TensorProto.FLOAT
	deq = [helper.make_node("Cast", [inits[0].name], [wname + "_f"], name=base + "_cast", to=F),
	helper.make_node("Cast", [inits[1].name], [wname + "_zf"], name=base + "_castz", to=F),
	helper.make_node("Sub", [wname + "_f", wname + "_zf"], [wname + "_c"], name=base + "_sub"),
	helper.make_node("Mul", [wname + "_c", inits[2].name], [wname + "_d3"], name=base + "_mul"),
	helper.make_node("Reshape", [wname + "_d3", inits[3].name], [wname + ("_dq" if kb * gs == K else "_dp")], name=base + "_reshape")]
	if kb * gs != K: # drop the zero padding rows
	lims = [numpy_helper.from_array(np.array(v, np.int64), wname + n) for v, n in (([0], "_s0"), ([K], "_sK"), ([0], "_ax"))]
	g.initializer.extend(lims)
	deq.append(helper.make_node("Slice", [wname + "_dp"] + [i.name for i in lims], [wname + "_dq"], name=base + "_slice"))
	g.node[idx].input[1] = wname + "_dq" # `node` belongs to another copy of the graph: edit this one
	for k, dn in enumerate(deq):
	g.node.insert(idx + k, dn)
	else:
	N = B.shape[0]
	inits = [numpy_helper.from_array(B, wname + "_Q4"), numpy_helper.from_array(S, wname + "_scales"),
	numpy_helper.from_array(zp, wname + "_zero_points")]
	g.initializer.extend(inits)
	kw = dict(K=K, N=N, bits=4, block_size=GS)
	if accuracy_level:
	kw["accuracy_level"] = accuracy_level
	new = helper.make_node("MatMulNBits", [node.input[0]] + [i.name for i in inits], list(node.output),
	name=(node.name or wname) + "_Q4", domain="com.microsoft", **kw)
	g.node.remove(node)
	g.node.insert(idx, new)
	drop_inits.add(wname)
	keep = [i for i in g.initializer if i.name not in drop_inits]
	del g.initializer[:]
	g.initializer.extend(keep)
	if any(n.domain == "com.microsoft" for n in g.node) and not any(o.domain == "com.microsoft" for o in m.opset_import):
	m.opset_import.append(helper.make_opsetid("com.microsoft", 1))
	onnx.checker.check_model(m)
	onnx.save(m, out_path)


	def quality(sess, samples, esess=None, ref=None):
	"""dE vs the training targets (Haiku for src 0, native dark for src 1); with the element graph, how often the
	element decisions computed from this page graph's outputs match those from the fp32 page graph."""
	agg = {0: [0.0, 0.0], 1: [0.0, 0.0]}
	ok = n = 0
	for s in samples:
	f = feeds_for(s)
	lab, _, ctx = sess.run(None, f)
	k = len(s["role"])
	m = s["tmask"][:k]
	if m.any():
	d = np.linalg.norm(lab[0, :k] - s["lab_t"][:k], axis=-1)
	w = s["w"][:k] * m
	agg[int(s["src"])][0] += float((d * w).sum() / max(w.sum(), 1e-9)); agg[int(s["src"])][1] += 1
	if esess is not None and ref is not None and len(s["E_ut"]):
	rl, _, rc = ref.run(None, f)
	a = esess.run(None, elem_feeds(s, lab, ctx)[0])[0][0].argmax(-1)
	b = esess.run(None, elem_feeds(s, rl, rc)[0])[0][0].argmax(-1)
	ok += int((a == b).sum()); n += len(a)
	out = {"haiku_dE": agg[0][0] / max(1, agg[0][1]), "native_dE": agg[1][0] / max(1, agg[1][1])}
	if n:
	out["elem_decisions_same_as_fp32"] = ok / n
	return out


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--ckpt", required=True)
	ap.add_argument("--data", required=True)
	ap.add_argument("--fp32", required=True, help="fp32 ONNX from export.py")
	ap.add_argument("--out", required=True)
	ap.add_argument("--calib", type=int, default=256)
	ap.add_argument("--eval", type=int, default=400)
	ap.add_argument("--threads", type=int, default=8)
	ap.add_argument("--elem", default="", help="element graph (fp32) to measure decision agreement")
	args = ap.parse_args()
	torch.set_num_threads(args.threads)
	os.makedirs(args.out, exist_ok=True)
	import onnx, onnxruntime as ort
	ck = torch.load(args.ckpt, map_location="cpu")
	model = NadirNet(**ck["cfg"]); model.load_state_dict(ck["model"]); model.eval()
	data = pickle.load(open(args.data, "rb"))
	hv = [s for s in data["val"] if s["src"] == 0]; nv = [s for s in data["val"] if s["src"] == 1]
	val = [prep(s) for s in hv[: args.eval * 3 // 4] + nv[: args.eval - min(len(hv), args.eval * 3 // 4)]] # both sources
	rng = np.random.default_rng(0)
	pick = rng.choice(len(data["train"]), size=min(args.calib, len(data["train"])), replace=False)
	calib = [feeds_for(prep(data["train"][i])) for i in pick]
	order = linear_order(model, calib[0])
	mapping = onnx_matmuls(model, onnx.load(args.fp32))
	quantizable = [n for n in order if n in mapping and not n.endswith(".rel")]
	print(f"{len(quantizable)}/{len(order)} linears map to MatMul; fp32 kept: {[n for n in order if n not in quantizable]}", flush=True)
	edges = [n for n in quantizable if n.startswith(("tok_in.", "img_in.", "color_head.", "ctx_head."))]
	plans = {"int8": {n: 8 for n in quantizable},
	"int4": {n: (8 if n in edges else 4) for n in quantizable}}
	so = ort.SessionOptions(); so.intra_op_num_threads = 4
	ref = ort.InferenceSession(args.fp32, so, providers=["CPUExecutionProvider"])
	report = {"format": {"8-bit layers": f"uint8 blocks of {GS} with fp32 scale + uint8 zero point, dequantized in-graph (Cast/Sub/Mul, folded at load)", "4-bit layers": f"uint4 blocks of {GS}, MatMulNBits"},
	"group_size": GS, "calib_pages": len(calib), "eval_pages": len(val), "fp32_layers": [n for n in order if n not in quantizable],
	"int4_8bit_edge_layers": edges, "variants": {}}
	esess = ort.InferenceSession(args.elem, so, providers=["CPUExecutionProvider"]) if args.elem else None
	bfeeds = [feeds_for(s) for s in val[:20]]
	r = quality(ref, val); r["size_mb"] = round(os.path.getsize(args.fp32) / 1e6, 2); r["cpu_ms_median"] = round(bench(ref, bfeeds), 2)
	report["variants"]["fp32"] = r
	print("fp32", r, flush=True)
	for vname, bits_for in plans.items():
	print(f"== {vname}", flush=True)
	qm, packed = quantize_model(model, calib, bits_for, quantizable)
	path = os.path.join(args.out, f"nadir-{vname}.onnx")
	build_onnx(args.fp32, mapping, packed, path)
	sess = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
	# packing parity: ORT MatMulNBits graph vs the PyTorch fake-quant model
	par = 0.0
	with torch.no_grad():
	for s in val[:8]:
	f = feeds_for(s)
	tq = qm(*[torch.from_numpy(f[k]) for k in NAMES_IN])[0].numpy()
	par = max(par, float(np.abs(tq - sess.run(None, f)[0]).max()))
	r = fidelity(sess, ref, val)
	r.update(quality(sess, val, esess, ref))
	r["pack_parity_max_abs"] = par
	r["size_mb"] = round(os.path.getsize(path) / 1e6, 2)
	r["cpu_ms_median"] = round(bench(sess, bfeeds), 2)
	report["variants"][vname] = r
	print(vname, r, flush=True)
	torch.save({"cfg": ck["cfg"], "model": qm.state_dict()}, os.path.join(args.out, f"fakequant-{vname}.pt"))
	json.dump(report, open(os.path.join(args.out, "quant_report.json"), "w"), indent=1)


	if __name__ == "__main__":
	main()