File size: 3,354 Bytes
111541a | 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 | #!/usr/bin/env python3
"""Requantize selected GGUF tensors (default: refiner FFN/attention weights) to Q4_K.
The reference Z-Image-Turbo Q4_K GGUF stores the two noise-refiner and two
context-refiner blocks as Q8_0. On a memory-bound CPU GEMM those 8-bit weights
make the refiner matmuls ~2x slower than the Q4_K main layers even though the
token count is smaller. This tool rewrites just the matching Q8_0 tensors as
Q4_K, leaving every other tensor byte-for-byte untouched.
Requires the small `q4k` helper (build with build.sh or see README) which wraps
ggml's own Q4_K quantizer.
"""
import argparse
import os
import subprocess
import sys
import tempfile
import numpy as np
import gguf
from gguf import GGUFReader, GGUFWriter, GGUFValueType
Q4K = os.environ.get("Q4K", "/tmp/q4k")
def q4k_quantize(f32: np.ndarray, helper: str) -> np.ndarray:
nrows, n_per_row = f32.shape
p = subprocess.run([helper, str(nrows), str(n_per_row)],
input=np.ascontiguousarray(f32, dtype=np.float32).tobytes(),
stdout=subprocess.PIPE, check=True)
return np.frombuffer(p.stdout, dtype=np.uint8)
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("src")
ap.add_argument("dst")
ap.add_argument("--pattern", default="refiner")
ap.add_argument("--only-kinds", default="attention,feed_forward",
help="comma list of substrings a tensor name must contain (empty = any)")
ap.add_argument("--helper", default=Q4K)
ap.add_argument("--dry-run", action="store_true")
args = ap.parse_args()
r = GGUFReader(args.src)
arch = str(r.fields["general.architecture"].contents()) if "general.architecture" in r.fields else "zimage"
# pick targets
kinds = [k for k in args.only_kinds.split(",") if k]
targets = []
for t in r.tensors:
if args.pattern in t.name and t.tensor_type == gguf.GGMLQuantizationType.Q8_0:
if kinds and not any(k in t.name for k in kinds):
continue
targets.append(t.name)
print(f"targets ({len(targets)}):")
for n in targets:
print(" ", n)
if args.dry_run or not targets:
return 0
w = GGUFWriter(args.dst, arch, use_temp_file=True)
for name, field in r.fields.items():
if name == "general.architecture" or name.startswith("GGUF."):
continue
v = field.contents()
w.add_key_value(name, v, field.types[0],
sub_type=field.types[-1] if len(field.types) > 1 else None)
tset = set(targets)
for t in r.tensors:
if t.name in tset:
f32 = gguf.quants.dequantize(t.data, t.tensor_type)
raw = q4k_quantize(f32, args.helper)
tgt = int(t.shape[0]) // 256 * 144 # Q4_K bytes per row (ne0 = input width)
w.add_tensor(t.name, raw.reshape(-1, tgt),
raw_dtype=gguf.GGMLQuantizationType.Q4_K)
print(f" requantized {t.name} {tuple(t.shape)} Q8_0->Q4_K "
f"{t.n_bytes} -> {raw.nbytes} bytes")
else:
w.add_tensor(t.name, t.data, raw_dtype=t.tensor_type)
w.write_header_to_file()
w.write_kv_data_to_file()
w.write_tensors_to_file()
w.close()
print("wrote", args.dst)
return 0
if __name__ == "__main__":
sys.exit(main())
|