Blame for __main__.py

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