File size: 5,577 Bytes
c99f13f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | import gguf
import numpy as np
from typing import List
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.
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
|