File size: 2,095 Bytes
5634ba9 | 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 | #!/usr/bin/env python3
"""Rank H3 transformer matrices by mean activation importance."""
from __future__ import annotations
import argparse
import json
import re
import struct
from pathlib import Path
import numpy as np
FC1 = re.compile(r"model\.diffusion_model\.blocks\.(\d+)\.mlp\.fc1\.weight$")
def read_entries(path: Path) -> list[dict[str, object]]:
rows: list[dict[str, object]] = []
with path.open("rb") as stream:
(entry_count,) = struct.unpack("<i", stream.read(4))
for _ in range(entry_count):
(name_length,) = struct.unpack("<i", stream.read(4))
name = stream.read(name_length).decode("utf-8")
calls, value_count = struct.unpack("<ii", stream.read(8))
values = np.frombuffer(stream.read(value_count * 4), dtype="<f4")
match = FC1.fullmatch(name)
if match:
rows.append(
{
"block": int(match.group(1)),
"calls": calls,
"values": value_count,
"mean": float(values.mean()),
"median": float(np.median(values)),
"p95": float(np.quantile(values, 0.95)),
"max": float(values.max()),
}
)
# stable-diffusion.cpp appends the collector's final call index.
trailer = stream.read()
if len(trailer) != 4:
raise ValueError(f"unexpected imatrix trailer length: {len(trailer)}")
struct.unpack("<i", trailer)
return sorted(rows, key=lambda row: float(row["mean"]), reverse=True)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("imatrix", type=Path)
parser.add_argument("--output", type=Path)
args = parser.parse_args()
rows = read_entries(args.imatrix)
rendered = json.dumps(rows, indent=2) + "\n"
if args.output:
args.output.write_text(rendered, encoding="utf-8")
else:
print(rendered, end="")
if __name__ == "__main__":
main()
|