| """Command-line entry point for HPC-Quantize V2. |
| |
| Examples:: |
| |
| python -m hpc_quantize_v2 --demo |
| python -m hpc_quantize_v2 --input tensor.npy --output tensor_q2k.npz |
| python -m hpc_quantize_v2 --gguf model.gguf --list-tensors |
| python -m hpc_quantize_v2 --gguf model.gguf --tensor blk.0.ffn_down.weight \ |
| --imatrix-gguf imatrix.gguf --output ffn_down_q2k.npz |
| python -m hpc_quantize_v2 --gguf model.gguf --all-q2k \ |
| --imatrix-gguf imatrix.gguf --output model_q2k_arrays.npz |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from .config import HPCConfig |
| from .factored import factor_only_solve |
| from .gguf import ( |
| GGUFReader, |
| expand_importance, |
| is_q2k_eligible, |
| lookup_imatrix, |
| read_imatrix, |
| target_quant_type, |
| write_quantized_gguf, |
| ) |
| from .pipeline import Q2KQuantizer |
| from .reference import reference_solve, synthetic_cases |
|
|
|
|
| def _config_from_args(args) -> HPCConfig: |
| cfg = HPCConfig(verbose=args.verbose) |
| if args.cand_d is not None: |
| cfg.n_cand_d = args.cand_d |
| if args.cand_m is not None: |
| cfg.n_cand_m = args.cand_m |
| if args.workers is not None: |
| cfg.workers = args.workers |
| if args.device is not None: |
| cfg.device = args.device |
| if args.mse: |
| cfg.optimizer = "mse" |
| if args.no_spectral: |
| cfg.use_spectral = False |
| if args.thermal is not None: |
| cfg.therm_passes = int(args.thermal) |
| if args.thermal_tilt is not None: |
| cfg.therm_lambda_tilt = float(args.thermal_tilt) |
| if args.pin_dc: |
| cfg.pin_dc = True |
| if args.no_q4_fallback: |
| cfg.enable_q4_fallback = False |
| if args.keep_embd: |
| cfg.quantize_embd_q4 = False |
| if args.promote_attn_q4: |
| cfg.promote_attn_q4 = True |
| if args.format is not None: |
| cfg.quant_format = args.format |
| if args.format in ("q2_k", "q3_k"): |
| cfg.enable_q4_fallback = False |
| if getattr(args, "tensor_type", None): |
| custom: dict[str, str] = {} |
| for item in args.tensor_type: |
| if "=" not in item: |
| raise SystemExit(f"--tensor-type expects TYPE=QUANT, got {item!r}") |
| key, value = (part.strip() for part in item.split("=", 1)) |
| if not key: |
| raise SystemExit(f"--tensor-type has an empty tensor type: {item!r}") |
| canonical = { |
| "q1_0": "Q1_0", |
| "q2_k": "Q2_K", |
| "q3_k": "Q3_K", |
| "q4_0": "Q4_0", |
| "copy": "Copy", |
| "auto": "Auto", |
| } |
| norm = canonical.get(value.lower()) |
| if norm is None: |
| raise SystemExit( |
| f"--tensor-type {key!r} has invalid quant {value!r}; " |
| "use Q1_0, Q2_K, Q3_K, Q4_0, Copy, or Auto" |
| ) |
| custom[key] = norm |
| cfg.custom_quant_map = custom |
| return cfg.validate() |
|
|
|
|
| def _importance_for(imat, name, n_elements, row_width): |
| if imat is None: |
| return None |
| raw = lookup_imatrix(imat, name) |
| if raw is None: |
| return None |
| return expand_importance(raw, n_elements, row_width) |
|
|
|
|
| def _save_result(path: str, res, prefix: str = "", extra: dict | None = None) -> None: |
| payload = { |
| f"{prefix}scales": res.scales, |
| f"{prefix}qs": res.qs, |
| f"{prefix}dm": res.dm, |
| f"{prefix}dequantized": res.dequantize(), |
| f"{prefix}rmse": np.float32(res.stats.rmse), |
| } |
| if extra: |
| payload.update(extra) |
| np.savez_compressed(path, **payload) |
|
|
|
|
| def _run_gguf(args, cfg: HPCConfig) -> None: |
| reader = GGUFReader(args.gguf) |
| imat = None |
| imatrix_path = args.imatrix_gguf or args.imatrix |
| if imatrix_path: |
| imat = read_imatrix(imatrix_path) |
| print(f"iMatrix: {len(imat)} tensors from {imatrix_path}") |
|
|
| if args.list_tensors or (not args.tensor and not args.all_q2k): |
| print(f"{reader.path}: {len(reader.tensors)} tensors, {len(reader.metadata)} metadata keys") |
| for ti in reader.tensors: |
| tgt = target_quant_type( |
| ti.name, |
| ti.dims, |
| enable_q4_fallback=cfg.enable_q4_fallback, |
| quantize_embd_q4=cfg.quantize_embd_q4, |
| promote_attn_q4=cfg.promote_attn_q4, |
| quant_format=cfg.quant_format, |
| source_type=ti.type, |
| custom_quant_map=cfg.custom_quant_map or None, |
| ) |
| tag = tgt if tgt else "COPY" |
| print( |
| f" {tag:4s} {ti.name:48s} {ti.type_name:6s} " |
| f"dims={ti.dims} row={ti.row_width}" |
| ) |
| return |
|
|
| quantizer = Q2KQuantizer(cfg) |
| wants_gguf = bool(args.output and str(args.output).lower().endswith(".gguf")) |
| if wants_gguf: |
| if not args.all_q2k and not args.tensor: |
| args.all_q2k = True |
| if args.rows: |
| raise SystemExit("--rows is only valid for .npz sampling output") |
| names = {args.tensor} if args.tensor else None |
| stats = write_quantized_gguf( |
| args.gguf, |
| args.output, |
| quantizer, |
| imat=imat, |
| quantize_all=args.all_q2k, |
| tensor_names=names, |
| limit=args.limit, |
| verbose=True, |
| ) |
| print( |
| f"wrote {args.output}: {int(stats.get('q2_k', 0))} Q2_K + {int(stats.get('q3_k', 0))} Q3_K + " |
| f"{int(stats.get('q4_0', 0))} Q4_0 + {int(stats.get('q1_0', 0))} Q1_0 " |
| f"({int(stats['quantized'])}/{int(stats['tensors'])} tensors), rmse={stats['rmse']:.6e}" |
| ) |
| return |
|
|
| if args.all_q2k: |
| payload: dict[str, np.ndarray] = {} |
| done = 0 |
| for ti in reader.tensors: |
| if not is_q2k_eligible(ti.name, ti.dims): |
| continue |
| if args.limit and done >= args.limit: |
| break |
| x, row = reader.read_tensor(ti.name) |
| if args.rows: |
| x = x[: args.rows * row] |
| imp = _importance_for(imat, ti.name, x.size, row) |
| res = quantizer.quantize(x, importance=imp, row_width=row) |
| prefix = ti.name.replace("/", "_") + "__" |
| payload[f"{prefix}scales"] = res.scales |
| payload[f"{prefix}qs"] = res.qs |
| payload[f"{prefix}dm"] = res.dm |
| payload[f"{prefix}rmse"] = np.float32(res.stats.rmse) |
| print( |
| f" {ti.name:48s} row={row:6d} rmse={res.stats.rmse:.6e} " |
| f"imatrix={'yes' if imp is not None else 'no'}" |
| ) |
| done += 1 |
| if args.output: |
| np.savez_compressed(args.output, **payload) |
| print(f"wrote {args.output} ({done} tensors)") |
| return |
|
|
| ti = reader.info(args.tensor) |
| x, row = reader.read_tensor(ti.name) |
| if args.rows: |
| x = x[: args.rows * row] |
| imp = _importance_for(imat, ti.name, x.size, row) |
| res = quantizer.quantize(x, importance=imp, row_width=row) |
| print( |
| f"{ti.name}: type={ti.type_name} row={row} blocks={len(res.blocks)} " |
| f"rmse={res.stats.rmse:.6e} imatrix={'yes' if imp is not None else 'no'}" |
| ) |
| if args.output: |
| _save_result( |
| args.output, |
| res, |
| extra={ |
| "name": np.array(ti.name), |
| "row_width": np.int64(row), |
| "source_type": np.array(ti.type_name), |
| }, |
| ) |
| print(f"wrote {args.output}") |
|
|
|
|
| def _run_numpy(args, cfg: HPCConfig) -> None: |
| weights = np.load(args.input) |
| importance = np.load(args.importance) if args.importance else None |
| if importance is None and (args.imatrix_gguf or args.imatrix): |
| if not args.tensor_name: |
| raise SystemExit("--tensor-name is required to select an iMatrix entry for .npy input") |
| imat = read_imatrix(args.imatrix_gguf or args.imatrix) |
| raw = lookup_imatrix(imat, args.tensor_name) |
| if raw is None: |
| raise SystemExit(f"iMatrix entry not found: {args.tensor_name}") |
| row = int(weights.shape[-1]) if weights.ndim >= 2 else int(weights.size) |
| importance = expand_importance(raw, weights.size, row) |
| res = Q2KQuantizer(cfg).quantize(weights, importance=importance) |
| if args.output: |
| _save_result(args.output, res) |
| print(f"wrote {args.output}") |
| else: |
| _ = res.dequantize() |
|
|
|
|
| def _bench(per_kind: int) -> None: |
| from collections import defaultdict |
|
|
| cases = synthetic_cases(123, per_kind) |
| for use_spectral in (False, True): |
| cfg = HPCConfig(use_spectral=use_spectral) |
| gains = defaultdict(list) |
| for kind, x in cases: |
| w = np.ones(x.size, dtype=np.float32) |
| v2 = factor_only_solve(x, w, cfg) |
| ref_obj, _ = reference_solve(x, w, cfg, use_spectral) |
| gains[kind].append((ref_obj - v2.objective) / max(abs(ref_obj), 1.0e-30)) |
| print(f"spectral={use_spectral}") |
| for kind, vals in gains.items(): |
| arr = np.asarray(vals) |
| print( |
| f" {kind:14s} mean={arr.mean()*100:6.2f}% " |
| f"median={np.median(arr)*100:6.2f}% min={arr.min()*100:6.2f}% " |
| f"worse={int(np.sum(arr < -1e-9))}/{arr.size}" |
| ) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="HPC-Quantize V2 factored Q2_K") |
| source = parser.add_mutually_exclusive_group() |
| source.add_argument("--input", help=".npy file containing float weights") |
| source.add_argument("--gguf", help="source GGUF file") |
| parser.add_argument("--tensor", help="GGUF tensor name to quantize") |
| parser.add_argument("--tensor-name", help="name used to look up an iMatrix entry for .npy input") |
| parser.add_argument("--list-tensors", action="store_true", help="list GGUF tensors and exit") |
| parser.add_argument("--all-q2k", action="store_true", help="quantize all eligible GGUF tensors") |
| parser.add_argument("--limit", type=int, default=0, help="limit tensors in --all-q2k") |
| parser.add_argument("--rows", type=int, default=0, help="limit rows for a single GGUF tensor") |
| parser.add_argument("--importance", help="optional .npy importance array") |
| parser.add_argument("--imatrix-gguf", help="llama.cpp iMatrix .gguf file") |
| parser.add_argument("--imatrix", help="alias for --imatrix-gguf (GGUF or legacy .dat)") |
| parser.add_argument("--output", help="output .gguf (full model) or .npz (block arrays)") |
| parser.add_argument("--demo", action="store_true", help="quantize synthetic weights") |
| parser.add_argument("--blocks", type=int, default=64, help="demo block count") |
| parser.add_argument("--std", type=float, default=0.02, help="demo weight std") |
| parser.add_argument("--seed", type=int, default=1234, help="demo RNG seed") |
| parser.add_argument("--cand-d", type=int, default=None, help="candidate d count") |
| parser.add_argument("--cand-m", type=int, default=None, help="candidate min count") |
| parser.add_argument("--workers", type=int, default=None, help="worker processes (0=auto CPU cores, 1=single)") |
| parser.add_argument("--device", default=None, help="compute device: auto (default), cuda, or cpu") |
| parser.add_argument("--mse", action="store_true", help="disable HPC sequence search") |
| parser.add_argument("--no-spectral", action="store_true", help="disable spectral/DC objective") |
| parser.add_argument( |
| "--thermal", |
| type=int, |
| default=None, |
| metavar="PASSES", |
| help="thermal refinery heat-bath passes after Phase 4 (default 0 = off; seeded, bit-exact per seed)", |
| ) |
| parser.add_argument( |
| "--thermal-tilt", |
| type=float, |
| default=None, |
| metavar="TILT", |
| help="zero-residual anneal: multiplier (>=1, 1=off) on dc/vw lambdas during thermal passes; " |
| "~100-1e3 drives the coherent residual (row DC + vesica) toward zero, spending weighted SSE", |
| ) |
| parser.add_argument( |
| "--pin-dc", |
| action="store_true", |
| help="zero-DC pin (Q2_K free lunch): re-solve (d, m) on the exact dc=0 line " |
| "after assembly/thermal; strictly lowers SSE while driving row-DC residual to ~0", |
| ) |
| parser.add_argument( |
| "--format", |
| choices=("mixed", "mixed_q1", "custom", "q1_0", "q3_k", "q4_0", "q2_k"), |
| default=None, |
| help="quantization format mode: mixed (Q2_K+Q4_0), mixed_q1 (Q1_0+Q4_0), custom (per-tensor-type overrides), q1_0 (pure Q1_0), q3_k (pure Q3_K), q4_0 (pure Q4_0), or q2_k (pure Q2_K)", |
| ) |
| parser.add_argument( |
| "--tensor-type", |
| action="append", |
| default=[], |
| metavar="TYPE=QUANT", |
| help="per-tensor-type override, e.g. ffn_up=Q1_0 or attn_output=Q1_0 " |
| "(QUANT: Q1_0, Q2_K, Q3_K, Q4_0, Copy, Auto; repeatable)", |
| ) |
| parser.add_argument("--no-q4-fallback", action="store_true", help="disable Q4_0 fallback for non-256-aligned tensors") |
| parser.add_argument("--keep-embd", action="store_true", help="keep token_embd/output.weight at source precision") |
| parser.add_argument("--promote-attn-q4", action="store_true", help="promote attention Q/K/V/O projections to Q4_0") |
| parser.add_argument("--verbose", action="store_true") |
| parser.add_argument("--bench", action="store_true", help="measure solver gains vs C port") |
| parser.add_argument("--bench-per-kind", type=int, default=12) |
| args = parser.parse_args() |
|
|
| cfg = _config_from_args(args) |
|
|
| if args.bench: |
| _bench(args.bench_per_kind) |
| return |
|
|
| if args.gguf: |
| if ( |
| not args.output |
| and not args.list_tensors |
| and (args.all_q2k or args.tensor) |
| ): |
| stem = Path(args.gguf).stem |
| args.output = str(Path(args.gguf).with_name(f"{stem}-Q2K.gguf")) |
| _run_gguf(args, cfg) |
| return |
|
|
| if args.demo: |
| rng = np.random.default_rng(args.seed) |
| weights = rng.normal(0.0, args.std, size=(args.blocks, 256)).astype(np.float32) |
| importance = None |
| if args.imatrix_gguf or args.imatrix: |
| imat = read_imatrix(args.imatrix_gguf or args.imatrix) |
| name = args.tensor_name or next(iter(imat)) |
| importance = expand_importance(lookup_imatrix(imat, name), weights.size, int(weights.shape[-1])) |
| res = Q2KQuantizer(cfg).quantize(weights, importance=importance) |
| if args.output: |
| _save_result(args.output, res) |
| print(f"wrote {args.output}") |
| return |
|
|
| if args.input: |
| _run_numpy(args, cfg) |
| return |
|
|
| parser.error("pass --demo, --input, --gguf or --bench") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
