MERNIK / imatrix_reader.py
wepiqx's picture
Upload imatrix_reader.py with huggingface_hub
3566461 verified
Raw History Blame Contribute Delete
6.93 kB
import gguf
import numpy as np
from typing import List
def imatrix_divergence(imatrix_list: List[dict]) -> List[tuple]:
"""Per-tensor relative spread of importance across several imatrix.
max-vs-first, relative to the first file's value. Used to decide whether
passing only one imatrix to both MERNIK and llama-quantize is honest
(they agree) or a silent truncation of the lens (they do not).
"""
if len(imatrix_list) < 2:
return []
base = imatrix_list[0]["tensors"]
out = []
for im in imatrix_list[1:]:
for name, t in im["tensors"].items():
if name not in base:
continue
b = base[name]["importance_mean"]
v = t["importance_mean"]
denom = max(abs(b), 1e-12)
out.append((abs(v - b) / denom, name, b, v))
out.sort(reverse=True)
return out
def read_imatrix(path: str) -> dict:
"""Parse imatrix GGUF, return per-tensor importance data."""
r = gguf.GGUFReader(path)
raw = {}
meta = {}
for k, v in r.fields.items():
try:
meta[k] = v.data
except:
meta[k] = str(v)
for t in r.tensors:
name = t.name
arr = np.array(t.data, dtype=np.float64)
if name.endswith(".in_sum2"):
base = name[:-8]
if base not in raw:
raw[base] = {}
raw[base]["in_sum2"] = arr
elif name.endswith(".counts"):
base = name[:-7]
if base not in raw:
raw[base] = {}
raw[base]["counts"] = float(np.mean(arr))
result = {}
for base, data in raw.items():
if "in_sum2" not in data:
continue
arr = data["in_sum2"]
result[base] = {
"importance_mean": float(np.mean(arr)),
"importance_sum": float(np.sum(arr)),
"importance_max": float(np.max(arr)),
"importance_min": float(np.min(arr)),
"n_elements": arr.size,
"in_sum2_raw": arr,
}
return {
"path": path,
"tensors": result,
"n_tensors": len(result),
"meta": meta,
}
def combine_imatrix(imatrix_list: List[dict], method: str = "max") -> dict:
"""Combine multiple imatrix into one by aggregating importance.
NOTE (scar 2026-09-28): this aggregates PER-TENSOR MEANS (max of means),
which is NOT the same operator as merging the per-column arrays and then
averaging (mean of elementwise max) — mean(max) >= max(mean), strictly.
The binary only ever reads one imatrix file, so the build's real
importance is whatever that single file holds. main.py therefore writes
the merged file first and reads the queue's numbers back from it; this
function is for exploration and for the explicit --no-merge-imatrix path.
Args:
imatrix_list: List of imatrix dicts from read_imatrix()
method: "max", "mean", or "weighted_mean"
"""
if not imatrix_list:
return {}
if len(imatrix_list) == 1:
return imatrix_list[0]
# Get all tensor names across all imatrix
all_names = set()
for im in imatrix_list:
all_names.update(im["tensors"].keys())
combined_tensors = {}
for name in all_names:
vals = []
in_sum2_raw = None
n_elements = 0
for im in imatrix_list:
if name in im["tensors"]:
t = im["tensors"][name]
vals.append(t["importance_mean"])
if in_sum2_raw is None and "in_sum2_raw" in t:
in_sum2_raw = t["in_sum2_raw"]
n_elements = max(n_elements, t["n_elements"])
if not vals:
continue
if method == "max":
imp_mean = max(vals)
elif method == "mean":
imp_mean = sum(vals) / len(vals)
elif method == "weighted_mean":
# Could weight by n_elements or dataset size
imp_mean = sum(vals) / len(vals)
else:
imp_mean = max(vals)
combined_tensors[name] = {
"importance_mean": imp_mean,
"importance_sum": imp_mean * n_elements,
"importance_max": max(v.get("importance_max", 0) for im in imatrix_list if name in im["tensors"] for v in [im["tensors"][name]]),
"importance_min": min(v.get("importance_min", float('inf')) for im in imatrix_list if name in im["tensors"] for v in [im["tensors"][name]]),
"n_elements": n_elements,
"in_sum2_raw": in_sum2_raw,
}
# Merge metadata
combined_meta = {}
for im in imatrix_list:
for k, v in im["meta"].items():
if k not in combined_meta:
combined_meta[k] = v
return {
"path": "+".join(im["path"] for im in imatrix_list),
"tensors": combined_tensors,
"n_tensors": len(combined_tensors),
"meta": combined_meta,
}
def detect_tied_groups(imatrix: dict, atol: float = 1e-5) -> list:
"""Find tied tensor groups (identical importance arrays)."""
names = sorted(imatrix["tensors"].keys())
tied_groups = []
visited = set()
for i, n1 in enumerate(names):
if n1 in visited:
continue
group = [n1]
arr1 = imatrix["tensors"][n1].get("in_sum2_raw")
if arr1 is None:
tied_groups.append(group)
visited.add(n1)
continue
for j in range(i + 1, len(names)):
n2 = names[j]
if n2 in visited:
continue
arr2 = imatrix["tensors"][n2].get("in_sum2_raw")
if arr2 is None:
continue
if arr1.shape == arr2.shape and np.allclose(arr1, arr2, atol=atol):
group.append(n2)
visited.add(n2)
tied_groups.append(group)
visited.add(n1)
return tied_groups
def build_importance_table(imatrix: dict, model: dict) -> dict:
"""Build unified importance table, merging imatrix with model tensor info."""
table = {}
for tname, info in imatrix["tensors"].items():
ttype = _imatrix_type(tname)
table[tname] = {
"importance_mean": info["importance_mean"],
"importance_sum": info["importance_sum"],
"importance_max": info["importance_max"],
"importance_min": info["importance_min"],
"n_elements": info["n_elements"],
"type": ttype,
}
# Also index by name without trailing dot (safety for old code)
for tname, info in list(table.items()):
if tname.endswith("."):
alt = tname.rstrip(".")
table[alt] = info
return table
def _imatrix_type(name: str) -> str:
parts = name.split(".")
if len(parts) >= 3 and parts[0] == "blk":
return parts[2]
return name