tools: q4nx_diff names the layout that matches, instead of calling a good container BAD (#1110)

tools/q4nx_diff.py

@@ -172,12 +172,72 @@ def compare_cmd(args):
     return 0
 
 
-def decode_cmd(args):
+# The worst per-element error a CORRECT container can show: a q4_1 container packed from a Q4_K_M GGUF
+# is a requantization, so exact equality is impossible and "no error anywhere" would call every honest
+# container BAD. Measured on a Qwen3-0.6B pack (Q4_K_M -> q4_1): worst max|err| 4.1e-2. A container whose
+# tiles belong to different weights lands an order of magnitude above it (measured 1.1-2.0).
+LAYOUT_MATCH_MAX_ERR = 0.1
+LAYOUT_MATCH_RATIO = 3.0     # ... and it must beat the other layout by this much
+
+
+def verdict(results):
+    """(rc, line) for {layout: (exact, inexact, worst_max_err)}.
+
+    A container is laid out ONE way, and a caller that guesses wrong gets a confident "BAD" for a good
+    container. So name the layout that matched, and when none did, name the closest and say so instead
+    of implying the container is broken.
+    """
+    if not results:
+        return 2, "nothing was compared"
+    order = sorted(results, key=lambda k: results[k][2])
+    best = order[0]
+    worst = results[best][2]
+    rival = results[order[1]][2] if len(order) > 1 else None
+    exact = results[best][0]
+    took = (exact > 0 and worst == 0.0) or (
+        worst <= LAYOUT_MATCH_MAX_ERR and (rival is None or rival >= worst * LAYOUT_MATCH_RATIO))
+    if took:
+        other = "" if rival is None else ", next layout %s worst max|err| %.3e" % (order[1], rival)
+        how = "%d tensor(s) exactly" % exact if worst == 0.0 and exact else "worst max|err| %.3e" % worst
+        return 0, "matches layout %s (%s%s)" % (best, how, other)
+    return 1, ("matches NO layout: closest is %s with worst max|err| %.3e (> %.3e, or not %.0fx better "
+               "than %s) - the container's tiles are not these weights"
+               % (best, worst, LAYOUT_MATCH_MAX_ERR, LAYOUT_MATCH_RATIO, order[1] if rival is not None else "-"))
+
+
+def _decode_pass(args, layout, quiet=False):
+    """One layout over every comparable tensor. Returns (exact, inexact, worst_max_err)."""
     import gguf
     from gguf.quants import dequantize
 
     hdr, base, mm = load(os.path.join(args.q4nx, "model.q4nx"))
-    dec = DECODERS[args.layout]
+    dec = DECODERS[layout]
+    r = gguf.GGUFReader(args.gguf)
+    want = [t for t in r.tensors if t.name.endswith(".weight") or t.name in ("output.weight",)]
+    if args.tensor:
+        want = [t for t in want if args.tensor in t.name]
+    if not want:
+        print("no tensors matched", file=sys.stderr)
+        return 0, 0, 0.0
+    exact = inexact = skipped = 0
+    worst = (0.0, "")
+    # our converter's names, so a GGUF tensor can be looked up in the container
+    lay = {"attn_q": "self_attn.q_proj", "attn_k": "self_attn.k_proj", "attn_v": "self_attn.v_proj",
+           "attn_output": "self_attn.o_proj", "attn_q_norm": "self_attn.q_norm",
+           "attn_k_norm": "self_attn.k_norm", "ffn_gate": "mlp.gate_proj",
+           "ffn_up": "mlp.up_proj", "ffn_down": "mlp.down_proj", "attn_norm": "input_layernorm",
+           "ffn_norm": "post_attention_layernorm", "token_embd": "model.embed_tokens",
+           "output_norm": "model.norm", "output": "lm_head"}
+
+    def q4nx_name(n):
+        head, _, field = n.rpartition(".")
+        if head.startswith("blk."):
+            idx = head.split(".")[1]
+            return "model.layers.%s.%s.weight" % (idx, lay[head.split(".", 2)[2]])
+        return lay[head] + "." + field
+
+    hdr, base, mm = load(os.path.join(args.q4nx, "model.q4nx"))
+    dec = DECODERS[layout]
     r = gguf.GGUFReader(args.gguf)
     want = [t for t in r.tensors if t.name.endswith(".weight") or t.name in ("output.weight",)]
     if args.tensor:
@@ -219,7 +279,8 @@ def q4nx_name(n):
         try:
             got = dec(raw, rows, cols).astype(np.float64)
         except ValueError as e:
-            print("skip %-46s %s" % (name, e))
+            if not quiet:
+                print("skip %-46s %s" % (name, e))
             skipped += 1
             continue
         ref = dequantize(t.data, t.tensor_type).astype(np.float64).reshape(rows, cols)
@@ -229,18 +290,59 @@ def q4nx_name(n):
         inexact += (not ok)
         if err.max() > worst[0]:
             worst = (float(err.max()), t.name)
-        if args.verbose or not ok:
+        if not quiet and (args.verbose or not ok):
             print("%s %-46s %-5s max|err| %.3e  rms %.3e"
                   % ("ok  " if ok else "BAD ", t.name, t.tensor_type.name, err.max(),
                      float(np.sqrt((err ** 2).mean()))))
-    print("\nlayout=%s: exact=%d inexact=%d skipped=%d   worst max|err| %.3e (%s)"
-          % (args.layout, exact, inexact, skipped, worst[0], worst[1] or "-"))
-    return 0 if inexact == 0 else 1
+    if not quiet:
+        print("layout=%s: exact=%d inexact=%d skipped=%d   worst max|err| %.3e (%s)"
+              % (layout, exact, inexact, skipped, worst[0], worst[1] or "-"))
+    return exact, inexact, worst[0]
+
+
+def decode_cmd(args):
+    layouts = sorted(DECODERS) if args.layout == "both" else [args.layout]
+    results = {}
+    for layout in layouts:
+        results[layout] = _decode_pass(args, layout, quiet=(len(layouts) > 1 and layout != layouts[-1]))
+    rc, line = verdict(results)
+    print(line)
+    return rc
+
+
+def self_test():
+    """The verdict logic, with the four cases this lane measured on real containers."""
+    bad = []
+    cases = [
+        # what a correct container looks like: one layout exact (a round trip), the other noise
+        ({"ours": (12, 0, 0.0), "flm": (0, 12, 1.6)}, 0, "ours", "an exact round trip picks its own layout"),
+        # a correct requantized container: 4.1e-2 is the honest Q4_K_M -> q4_1 error, 1.6 is the other
+        ({"ours": (0, 12, 4.1e-2), "flm": (0, 12, 1.6)}, 0, "ours", "a lossy but correct pack still matches"),
+        ({"ours": (0, 12, 4.1e-2)}, 0, "ours", "a single layout under the tolerance matches"),
+        # a wrong converter: both layouts far above the tolerance, so the container is named as wrong
+        ({"ours": (0, 12, 2.0), "flm": (0, 12, 1.1)}, 1, "NO", "tiles from other weights match nothing"),
+        # ambiguous: close to the tolerance and not clearly better than its rival
+        ({"ours": (0, 12, 0.09), "flm": (0, 12, 0.12)}, 1, "NO", "not clearly better than the rival"),
+        ({}, 2, "nothing", "nothing compared"),
+    ]
+    for results, want_rc, want_word, what in cases:
+        rc, line = verdict(results)
+        if rc != want_rc or want_word not in line:
+            bad.append("%s: got rc=%d %r" % (what, rc, line))
+    if bad:
+        print("self-test FAILED:")
+        for b in bad:
+            print("  " + b)
+        return 1
+    print("self-test passed: the verdict names the matching layout, accepts a lossy-but-correct pack, "
+          "and refuses tiles from other weights")
+    return 0
 
 
 def main():
     ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
-    sub = ap.add_subparsers(dest="cmd", required=True)
+    ap.add_argument("--self-test", action="store_true", help="the layout verdict, on synthetic results")
+    sub = ap.add_subparsers(dest="cmd", required=False)
     c = sub.add_parser("compare", help="compare two model.q4nx containers")
     c.add_argument("a")
     c.add_argument("b")
@@ -249,11 +351,16 @@ def main():
     d = sub.add_parser("decode", help="decode a container and check it against its GGUF")
     d.add_argument("--gguf", required=True)
     d.add_argument("--q4nx", required=True, help="the model directory")
-    d.add_argument("--layout", required=True, choices=sorted(DECODERS))
+    d.add_argument("--layout", required=True, choices=sorted(DECODERS) + ["both"],
+                   help="which tile layout the container uses; 'both' tries each and names the one that "
+                        "matches (a wrong guess reports a good container as BAD, which is how the "
+                        "converter mismatch in engine#1110 hid)")
     d.add_argument("--tensor", help="only tensors whose GGUF name contains this")
     d.add_argument("-v", "--verbose", action="store_true")
     d.set_defaults(func=decode_cmd)
     args = ap.parse_args()
+    if args.self_test:
+        return self_test()
     return args.func(args)