--analog-imatrix: Reduce error in output by 15%

hexstate_requantize.py

@@ -10,6 +10,14 @@ This bypasses the tokenizer parsing problem entirely — the source GGUF
 
 Usage:
     python3 hexstate_requantize.py input.gguf output.gguf
+
+ANALOG mode (minimises LAYER OUTPUT error instead of weight RMSE):
+    --hessian FILE.hxh     full X^T X per tensor (see write_hessian)
+    --analog-imatrix       diagonal Hessian from --imatrix (weaker, no new
+                           calibration needed)
+    --analog-sweeps N      coordinate-descent sweeps (default 8)
+
+  In ANALOG mode weight RMSE RISES by design. Evaluate with perplexity.
 """
 
 import struct
@@ -53,6 +61,22 @@ def _load_hexstate_lib():
             ctypes.c_int,                     # verbose
         ]
 
+        # ANALOG Q2_K (Hessian / output-error objective)
+        if hasattr(lib, 'hexstate_quantize_tensor_q2k_analog'):
+            lib.hexstate_quantize_tensor_q2k_analog.restype = None
+            lib.hexstate_quantize_tensor_q2k_analog.argtypes = [
+                ctypes.POINTER(ctypes.c_float),  # weights
+                ctypes.c_int64,                   # n_elements
+                ctypes.c_int64,                   # n_in (row length)
+                ctypes.POINTER(ctypes.c_float),   # H (n_in*n_in) or NULL
+                ctypes.c_void_p,                  # output
+                ctypes.POINTER(ctypes.c_float),   # out_error
+                ctypes.c_int,                     # n_sweeps
+                ctypes.c_int,                     # verbose
+            ]
+            lib.hexstate_q2k_hessian_bytes.restype = ctypes.c_int64
+            lib.hexstate_q2k_hessian_bytes.argtypes = [ctypes.c_int64]
+
         lib.hexstate_q2k_block_bytes.restype = ctypes.c_int
         lib.hexstate_q2k_block_bytes.argtypes = []
         lib.hexstate_q2k_block_elements.restype = ctypes.c_int
@@ -124,6 +148,121 @@ def _skip_gguf_kv_value(f, vtype):
         f.read(sz)
 
 
+# ─── Hessian (X^T X) support for the ANALOG objective ───────────────────────
+#
+# FILE FORMAT  (.hxh)  — little-endian throughout
+#
+#   magic     4 bytes   b"HXH1"
+#   n_entries uint32
+#   repeated n_entries times:
+#       name_len   uint32
+#       name       name_len bytes, UTF-8 tensor name (GGUF naming)
+#       n_in       uint32           row length; must be a multiple of 256
+#       H          n_in*n_in float32, ROW-MAJOR
+#
+# H must be X^T X / n_samples for that layer's INPUT activations, damped on
+# the diagonal (add ~1% of mean(diag) to keep it invertible). This is NOT a
+# llama.cpp imatrix: an imatrix stores only diag(X^T X), the n_in diagonal
+# entries. Entries may be shared between tensors with the same input space
+# (e.g. attn_q / attn_k / attn_v) by writing the same H under each name.
+#
+# Memory: n_in^2 * 4 bytes. 4 MiB at n_in=1024, 64 MiB at n_in=4096.
+
+HESSIAN_MAGIC = b"HXH1"
+
+
+def read_hessian(path):
+    """Read a .hxh file. Returns dict: tensor_name -> float32 (n_in, n_in)."""
+    import struct as st
+    out = {}
+    with open(path, 'rb') as f:
+        if f.read(4) != HESSIAN_MAGIC:
+            raise ValueError(f"{path}: not a HXH1 Hessian file")
+        n_entries = st.unpack('<I', f.read(4))[0]
+        for _ in range(n_entries):
+            nl = st.unpack('<I', f.read(4))[0]
+            name = f.read(nl).decode('utf-8', errors='replace')
+            n_in = st.unpack('<I', f.read(4))[0]
+            buf = f.read(n_in * n_in * 4)
+            if len(buf) != n_in * n_in * 4:
+                raise ValueError(f"{path}: truncated entry for {name}")
+            out[name] = np.frombuffer(buf, dtype=np.float32).reshape(n_in, n_in)
+    return out
+
+
+def write_hessian(path, entries):
+    """entries: dict name -> (n_in, n_in) array convertible to float32."""
+    import struct as st
+    with open(path, 'wb') as f:
+        f.write(HESSIAN_MAGIC)
+        f.write(st.pack('<I', len(entries)))
+        for name, H in entries.items():
+            H = np.ascontiguousarray(H, dtype=np.float32)
+            assert H.ndim == 2 and H.shape[0] == H.shape[1], name
+            nb = name.encode('utf-8')
+            f.write(st.pack('<I', len(nb)))
+            f.write(nb)
+            f.write(st.pack('<I', H.shape[0]))
+            f.write(H.tobytes())
+
+
+def hessian_from_imatrix(imat_row, n_in, damp=0.01):
+    """Build a DIAGONAL Hessian from an existing imatrix row.
+
+    An imatrix gives diag(X^T X) only, so this discards all cross-column
+    correlation. It is strictly weaker than a real Hessian -- measured on a
+    synthetic layer, diagonal-only recovers roughly half the output-error
+    gain that a full H delivers -- but it needs no new calibration run and
+    lets the analog path be exercised with data you already have.
+    """
+    d = np.asarray(imat_row, dtype=np.float64).ravel()
+    if d.size < n_in:
+        d = np.tile(d, int(np.ceil(n_in / max(d.size, 1))))
+    d = d[:n_in].copy()
+    d[~np.isfinite(d)] = 0.0
+    if d.max() <= 0:
+        d[:] = 1.0
+    d += damp * d.mean()
+    H = np.zeros((n_in, n_in), dtype=np.float32)
+    np.fill_diagonal(H, d.astype(np.float32))
+    return H
+
+
+def quantize_tensor_q2k_analog(f32_data, n_in, H, sweeps=0, verbose=1):
+    """Assign fresh 2-bit codes minimising LAYER OUTPUT error (not weight RMSE).
+
+    Weight RMSE is EXPECTED to rise; judge this path by perplexity.
+    Returns (bytes, n_blocks). Raises if the C library lacks the entry point.
+    """
+    lib = _load_hexstate_lib()
+    if lib is None or not hasattr(lib, 'hexstate_quantize_tensor_q2k_analog'):
+        raise RuntimeError("libhexstate_q2k.so has no analog entry point; rebuild")
+
+    n_elements = len(f32_data)
+    if n_elements % QK_K != 0:
+        pad = QK_K - (n_elements % QK_K)
+        f32_data = np.concatenate([f32_data, np.zeros(pad, dtype=np.float32)])
+        n_elements = len(f32_data)
+
+    n_blocks = n_elements // QK_K
+    output = np.zeros(n_blocks * lib.hexstate_q2k_block_bytes(), dtype=np.uint8)
+    error = ctypes.c_float(0.0)
+
+    f32c = np.ascontiguousarray(f32_data, dtype=np.float32)
+    Hc = np.ascontiguousarray(H, dtype=np.float32)
+    lib.hexstate_quantize_tensor_q2k_analog(
+        f32c.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
+        ctypes.c_int64(n_elements),
+        ctypes.c_int64(int(n_in)),
+        Hc.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
+        output.ctypes.data_as(ctypes.c_void_p),
+        ctypes.byref(error),
+        ctypes.c_int(int(sweeps)),
+        ctypes.c_int(int(verbose)),
+    )
+    return output.tobytes(), n_blocks
+
+
 def read_imatrix(path):
     """Read llama.cpp importance matrix file (GGUF or legacy .dat format).
 
@@ -734,6 +873,27 @@ def main():
                 print(f"  WARNING: imatrix file not found: {imat_path}")
             break
 
+    # ── ANALOG objective options ──
+    hessian_data = None
+    analog_from_imat = '--analog-imatrix' in sys.argv
+    analog_sweeps = 0
+    for i, arg in enumerate(sys.argv):
+        if arg == '--hessian' and i + 1 < len(sys.argv):
+            hpath = sys.argv[i + 1]
+            if os.path.exists(hpath):
+                hessian_data = read_hessian(hpath)
+                print(f"  Loaded Hessian: {len(hessian_data)} tensors from {hpath}")
+            else:
+                print(f"  WARNING: Hessian file not found: {hpath}")
+        elif arg == '--analog-sweeps' and i + 1 < len(sys.argv):
+            analog_sweeps = int(sys.argv[i + 1])
+    use_analog = (hessian_data is not None) or analog_from_imat
+    if use_analog:
+        print("  ANALOG MODE: minimising LAYER OUTPUT error, not weight RMSE.")
+        print("  Weight RMSE will RISE. Evaluate with perplexity, not RMSE.")
+        if analog_from_imat and hessian_data is None:
+            print("  Using DIAGONAL Hessian from imatrix (~half the gain of full H).")
+
     # Check for HPC C library
     use_hpc = _load_hexstate_lib() is not None
 
@@ -1057,7 +1217,14 @@ def main():
             total_quant_bytes = 0
             total_keep_bytes = 0
             total_rmse = 0.0
-            q2k_rmse_sum = 0.0
+            # Element-weighted global accumulators. A plain mean of per-tensor
+            # RMSEs weights a 4K-element norm tensor the same as a 100M-element
+            # FFN, so it is dominated by the small tensors and is the wrong
+            # thing to tune against. Track sum-of-squared-error and element
+            # count instead, and derive the true global RMSE at the end.
+            q2k_se_sum = 0.0          # Σ squared error over all Q2_K elements
+            q2k_elem_count = 0        # Σ elements
+            q2k_rmse_sum = 0.0        # kept for the unweighted mean (reported too)
             q2k_tensor_count = 0
 
             for i, ti in enumerate(tensor_infos):
@@ -1228,7 +1395,29 @@ def main():
                         imat_full = np.tile(iw, n_rows)[:ti['n_elements']]
 
                     n_el = ti['n_elements']
-                    if use_hpc and n_el <= HPC_CHUNK:
+
+                    # ── ANALOG path: fresh 2-bit codes vs output error ──
+                    analog_H = None
+                    analog_n_in = ti['dims'][0] if ti.get('dims') else 0
+                    if use_analog and analog_n_in and analog_n_in % QK_K == 0 \
+                       and n_el % analog_n_in == 0:
+                        if hessian_data and ti['name'] in hessian_data:
+                            Hcand = hessian_data[ti['name']]
+                            if Hcand.shape[0] == analog_n_in:
+                                analog_H = Hcand
+                        if analog_H is None and analog_from_imat and imatrix_data \
+                           and ti['name'] in imatrix_data:
+                            analog_H = hessian_from_imatrix(
+                                imatrix_data[ti['name']], analog_n_in)
+
+                    if analog_H is not None:
+                        q2k_data, n_blocks = quantize_tensor_q2k_analog(
+                            f32, analog_n_in, analog_H, sweeps=analog_sweeps)
+                        src = 'full-H' if (hessian_data and ti['name'] in hessian_data) \
+                              else 'diag-H'
+                        print(f"    [analog:{src}] n_in={analog_n_in} "
+                              f"rows={n_el // analog_n_in}")
+                    elif use_hpc and n_el <= HPC_CHUNK:
                         # Small tensor — single HPC pass
                         q2k_data, n_blocks = quantize_tensor_q2k_hpc(f32, opt_mode=2, importance=imat_full)
                     elif use_hpc:
@@ -1274,8 +1463,8 @@ def main():
                     # ── Compute and report exact per-tensor RMSE ──
                     try:
                         CHUNK_BLK = 100_000  # blocks per chunk to bound memory
-                        total_se = 0.0
-                        total_n = 0
+                        total_se = 0.0   # float64: np.sum on float32 would
+                        total_n = 0      # round away small terms on big tensors
                         for ci in range(0, n_blocks, CHUNK_BLK):
                             ce = min(ci + CHUNK_BLK, n_blocks)
                             chunk_q = q2k_data[ci*84:ce*84]
@@ -1283,9 +1472,11 @@ def main():
                             orig_chunk = f32[ci*QK_K:ce*QK_K]
                             n_valid = min(len(orig_chunk), len(deq_chunk))
                             diff = orig_chunk[:n_valid] - deq_chunk[:n_valid]
-                            total_se += np.sum(diff ** 2)
+                            total_se += float(np.sum(diff.astype(np.float64) ** 2))
                             total_n += n_valid
                         tensor_rmse = np.sqrt(total_se / max(total_n, 1))
+                        q2k_se_sum += total_se
+                        q2k_elem_count += total_n
                         q2k_rmse_sum += tensor_rmse
                         q2k_tensor_count += 1
                         print(f"\n  [Q2_K] {ti['name'][:55]}  RMSE={tensor_rmse:.6e}")
@@ -1326,7 +1517,13 @@ def main():
     print(f"  ║  Compression:       {compression:>42.1f}x ║")
     if q2k_tensor_count > 0:
         mean_rmse = q2k_rmse_sum / q2k_tensor_count
-        print(f"  ║  Mean Q2_K RMSE:                            {mean_rmse:>12.6e} ║")
+        # Global RMSE = sqrt(Σ SE / Σ N): every quantized weight counts once,
+        # regardless of which tensor it lives in. This is the number to tune
+        # against; the unweighted per-tensor mean is kept below for continuity
+        # with previous runs but is dominated by small tensors.
+        global_rmse = np.sqrt(q2k_se_sum / max(q2k_elem_count, 1))
+        print(f"  ║  Global Q2_K RMSE (element-wtd):            {global_rmse:>12.6e} ║")
+        print(f"  ║  Mean Q2_K RMSE (per-tensor, unwtd):        {mean_rmse:>12.6e} ║")
     print(f"  ║  Total time:        {elapsed:>39.1f} sec ║")
     print("  ╚════════════════════════════════════════════════════════════════╝")
     print()
@@ -1335,4 +1532,4 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
+    main()