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