| |
| """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" |
|
|
| |
| 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 |
| 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()) |
|
|
