quantbenchx 0.3.0__py3-none-any.whl

quantbenchx/cli.py

@@ -0,0 +1,153 @@
+"""CLI for quantbenchx."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Optional
+
+
+def _build_cli():  # type: ignore[no-untyped-def]
+    try:
+        import click
+    except ImportError:
+        raise SystemExit("CLI dependencies required: pip install quantbenchx[cli]")
+
+    @click.group()
+    @click.version_option(package_name="quantbenchx")
+    def cli() -> None:
+        """quantbenchx — quantization quality analyzer."""
+
+    @cli.command()
+    @click.argument("model_path", type=click.Path(exists=True))
+    @click.option("--json-out", "-o", type=click.Path(), default=None)
+    @click.option("--markdown", "-m", is_flag=True)
+    def profile(model_path: str, json_out: Optional[str], markdown: bool) -> None:
+        """Profile a quantized model file (GGUF or safetensors)."""
+        from quantbenchx.predict import estimate_quality
+        from quantbenchx.profile import profile_gguf, profile_safetensors
+        from quantbenchx.report import (
+            format_markdown,
+            format_report_rich,
+            format_report_text,
+            save_json,
+        )
+
+        p = Path(model_path)
+        if p.suffix == ".gguf":
+            prof = profile_gguf(p)
+        elif p.suffix == ".safetensors":
+            prof = profile_safetensors(p)
+        else:
+            click.echo(f"Unknown format: {p.suffix}", err=True)
+            raise SystemExit(1)
+
+        quality = estimate_quality(prof)
+
+        if markdown:
+            click.echo(format_markdown(prof, quality))
+        else:
+            try:
+                click.echo(format_report_rich(prof, quality))
+            except Exception:
+                click.echo(format_report_text(prof, quality))
+
+        if json_out:
+            save_json(prof, json_out, quality)
+            click.echo(f"Report saved to {json_out}", err=True)
+
+    @cli.command()
+    @click.argument("model_a", type=click.Path(exists=True))
+    @click.argument("model_b", type=click.Path(exists=True))
+    def compare(model_a: str, model_b: str) -> None:
+        """Compare two quantized model files."""
+        from quantbenchx.compare import compare_profiles
+        from quantbenchx.profile import profile_gguf, profile_safetensors
+
+        def _load(path: str):  # type: ignore[no-untyped-def]
+            p = Path(path)
+            if p.suffix == ".gguf":
+                return profile_gguf(p)
+            elif p.suffix == ".safetensors":
+                return profile_safetensors(p)
+            else:
+                click.echo(f"Unknown format: {p.suffix}", err=True)
+                raise SystemExit(1)
+
+        prof_a = _load(model_a)
+        prof_b = _load(model_b)
+        result = compare_profiles(prof_a, prof_b)
+        click.echo(json.dumps(result, indent=2, default=str))
+
+    @cli.command()
+    @click.argument("model_path", type=click.Path(exists=True))
+    def layers(model_path: str) -> None:
+        """Show layerwise analysis for a quantized model."""
+        from quantbenchx.layerwise import analyze_layers
+        from quantbenchx.profile import profile_gguf, profile_safetensors
+
+        p = Path(model_path)
+        if p.suffix == ".gguf":
+            prof = profile_gguf(p)
+        elif p.suffix == ".safetensors":
+            prof = profile_safetensors(p)
+        else:
+            click.echo(f"Unknown format: {p.suffix}", err=True)
+            raise SystemExit(1)
+
+        analysis = analyze_layers(prof)
+        try:
+            from rich.console import Console
+            from rich.table import Table
+
+            console = Console()
+            table = Table(title="Layerwise Analysis")
+            table.add_column("Layer", style="bold")
+            table.add_column("Params", justify="right")
+            table.add_column("BPW", justify="right")
+            table.add_column("Dtype")
+            table.add_column("Sensitivity", justify="right")
+
+            for row in analysis:
+                sens = row["sensitivity"]
+                sens_style = "green" if sens < 0.5 else "yellow" if sens < 0.7 else "red"
+                table.add_row(
+                    row["name"],
+                    f"{row['n_params']:,}",
+                    f"{row['avg_bits_per_weight']:.2f}",
+                    row["dominant_dtype"],
+                    f"[{sens_style}]{sens:.3f}[/{sens_style}]",
+                )
+            console.print(table)
+        except ImportError:
+            for row in analysis:
+                click.echo(
+                    f"{row['name']:30s} {row['n_params']:>12,} "
+                    f"{row['avg_bits_per_weight']:>6.2f} {row['dominant_dtype']:8s} "
+                    f"{row['sensitivity']:.3f}"
+                )
+
+    @cli.command()
+    @click.argument("model_path", type=click.Path(exists=True))
+    @click.option("--target-bpw", type=float, default=4.5)
+    def recommend(model_path: str, target_bpw: float) -> None:
+        """Recommend mixed-precision quantization strategy."""
+        from quantbenchx.layerwise import recommend_mixed_quant
+        from quantbenchx.profile import profile_gguf, profile_safetensors
+
+        p = Path(model_path)
+        if p.suffix == ".gguf":
+            prof = profile_gguf(p)
+        elif p.suffix == ".safetensors":
+            prof = profile_safetensors(p)
+        else:
+            click.echo(f"Unknown format: {p.suffix}", err=True)
+            raise SystemExit(1)
+
+        result = recommend_mixed_quant(prof, target_bpw=target_bpw)
+        click.echo(json.dumps(result, indent=2, default=str))
+
+    return cli
+
+
+cli = _build_cli()

quantbenchx/compare.py

@@ -0,0 +1,101 @@
+"""Compare quantized model profiles across formats and methods."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+from quantbenchx._types import ModelProfile
+
+
+def compare_profiles(profile_a: ModelProfile, profile_b: ModelProfile) -> Dict[str, Any]:
+    """Compare two model profiles and return a structured diff."""
+    size_delta = profile_b.total_size_bytes - profile_a.total_size_bytes
+    size_ratio = (
+        profile_b.total_size_bytes / profile_a.total_size_bytes
+        if profile_a.total_size_bytes > 0
+        else 0.0
+    )
+
+    bpw_a = profile_a.quant.avg_bits_per_weight
+    bpw_b = profile_b.quant.avg_bits_per_weight
+    bpw_delta = bpw_b - bpw_a
+
+    # Find tensors that differ in dtype
+    tensors_a = {t.name: t for t in profile_a.tensors}
+    tensors_b = {t.name: t for t in profile_b.tensors}
+    common = set(tensors_a.keys()) & set(tensors_b.keys())
+
+    dtype_changes: List[Dict[str, str]] = []
+    for name in sorted(common):
+        ta, tb = tensors_a[name], tensors_b[name]
+        if ta.dtype != tb.dtype:
+            dtype_changes.append({
+                "tensor": name,
+                "dtype_a": ta.dtype.value,
+                "dtype_b": tb.dtype.value,
+                "bpw_a": ta.bits_per_weight,
+                "bpw_b": tb.bits_per_weight,
+            })
+
+    only_a = sorted(set(tensors_a.keys()) - set(tensors_b.keys()))
+    only_b = sorted(set(tensors_b.keys()) - set(tensors_a.keys()))
+
+    return {
+        "model_a": profile_a.name,
+        "model_b": profile_b.name,
+        "format_a": profile_a.format.value,
+        "format_b": profile_b.format.value,
+        "size_a_gb": round(profile_a.size_gb, 3),
+        "size_b_gb": round(profile_b.size_gb, 3),
+        "size_delta_bytes": size_delta,
+        "size_ratio": round(size_ratio, 3),
+        "bpw_a": round(bpw_a, 2),
+        "bpw_b": round(bpw_b, 2),
+        "bpw_delta": round(bpw_delta, 2),
+        "method_a": profile_a.quant.method.value,
+        "method_b": profile_b.quant.method.value,
+        "n_dtype_changes": len(dtype_changes),
+        "dtype_changes": dtype_changes[:50],  # limit output
+        "tensors_only_in_a": only_a[:20],
+        "tensors_only_in_b": only_b[:20],
+        "params_a": profile_a.total_params,
+        "params_b": profile_b.total_params,
+    }
+
+
+def compare_formats(profiles: List[ModelProfile]) -> Dict[str, Any]:
+    """Compare multiple quantization formats of the same base model.
+
+    Returns a summary table and ranking by size vs quality trade-off.
+    """
+    if not profiles:
+        return {"models": [], "ranking": []}
+
+    rows: List[Dict[str, Any]] = []
+    for p in profiles:
+        rows.append({
+            "name": p.name,
+            "format": p.format.value,
+            "method": p.quant.method.value,
+            "size_gb": round(p.size_gb, 3),
+            "avg_bpw": round(p.quant.avg_bits_per_weight, 2),
+            "compression_ratio": round(p.compression_ratio, 2),
+            "n_params": p.total_params,
+            "n_quantized": p.quant.n_quantized_layers,
+            "n_full_precision": p.quant.n_full_precision_layers,
+        })
+
+    # Rank by compression efficiency (higher compression at higher bpw = better quality)
+    ranked = sorted(
+        rows,
+        key=lambda r: (r["avg_bpw"], -r["compression_ratio"]),
+        reverse=True,
+    )
+    for i, r in enumerate(ranked):
+        r["rank"] = i + 1
+
+    return {
+        "models": rows,
+        "ranking": ranked,
+        "n_models": len(profiles),
+    }

quantbenchx/imatrix.py

@@ -0,0 +1,201 @@
+"""GGUF importance matrix (imatrix) parsing and analysis."""
+
+from __future__ import annotations
+
+import math
+import struct
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Optional, Union
+
+from quantbenchx._types import QuantbenchError
+
+
+@dataclass
+class ImatrixEntry:
+    """A single tensor entry in an importance matrix."""
+
+    name: str
+    num_values: int
+    num_calls: int
+    values: list[float]
+
+
+@dataclass
+class ImatrixData:
+    """Parsed importance matrix data."""
+
+    entries: list[ImatrixEntry]
+    total_calls: int
+
+    def by_name(self, name: str) -> Optional[ImatrixEntry]:
+        """Look up an entry by tensor name."""
+        for entry in self.entries:
+            if entry.name == name:
+                return entry
+        return None
+
+    def top_k(self, k: int) -> list[ImatrixEntry]:
+        """Return the top-k entries ranked by mean importance value."""
+        def _mean(entry: ImatrixEntry) -> float:
+            if not entry.values:
+                return 0.0
+            return sum(entry.values) / len(entry.values)
+
+        return sorted(self.entries, key=_mean, reverse=True)[:k]
+
+
+@dataclass
+class ImatrixAnalysis:
+    """Result of analyzing an importance matrix."""
+
+    total_tensors: int
+    mean_importance_per_layer: Dict[str, float]
+    variance_per_layer: Dict[str, float]
+    outlier_layers: list[str]
+
+
+def parse_imatrix(path: Union[str, Path]) -> ImatrixData:
+    """Parse a binary imatrix file.
+
+    Format per tensor:
+        name_length (int32) + name (bytes) + num_values (int32)
+        + num_calls (int32) + values (float32 * num_values)
+    """
+    path = Path(path)
+    if not path.exists():
+        raise QuantbenchError(f"imatrix file not found: {path}")
+
+    entries: list[ImatrixEntry] = []
+    total_calls = 0
+
+    with open(path, "rb") as f:
+        data = f.read()
+
+    offset = 0
+    while offset < len(data):
+        # name_length (int32)
+        if offset + 4 > len(data):
+            break
+        (name_length,) = struct.unpack_from("<i", data, offset)
+        offset += 4
+
+        if name_length <= 0 or offset + name_length > len(data):
+            raise QuantbenchError(
+                f"Invalid name length {name_length} at offset {offset - 4}"
+            )
+
+        # name (bytes)
+        name = data[offset : offset + name_length].decode("utf-8")
+        offset += name_length
+
+        # num_values (int32)
+        if offset + 4 > len(data):
+            raise QuantbenchError(f"Unexpected EOF reading num_values for '{name}'")
+        (num_values,) = struct.unpack_from("<i", data, offset)
+        offset += 4
+
+        if num_values < 0:
+            raise QuantbenchError(f"Invalid num_values {num_values} for '{name}'")
+
+        # num_calls (int32)
+        if offset + 4 > len(data):
+            raise QuantbenchError(f"Unexpected EOF reading num_calls for '{name}'")
+        (num_calls,) = struct.unpack_from("<i", data, offset)
+        offset += 4
+
+        # values (float32 * num_values)
+        values_size = num_values * 4
+        if offset + values_size > len(data):
+            raise QuantbenchError(
+                f"Unexpected EOF reading {num_values} values for '{name}'"
+            )
+        values = list(struct.unpack_from(f"<{num_values}f", data, offset))
+        offset += values_size
+
+        total_calls = max(total_calls, num_calls)
+        entries.append(
+            ImatrixEntry(
+                name=name,
+                num_values=num_values,
+                num_calls=num_calls,
+                values=values,
+            )
+        )
+
+    return ImatrixData(entries=entries, total_calls=total_calls)
+
+
+def analyze_imatrix(data: ImatrixData) -> ImatrixAnalysis:
+    """Analyze an importance matrix and identify outlier layers."""
+    mean_importance: Dict[str, float] = {}
+    variance: Dict[str, float] = {}
+
+    for entry in data.entries:
+        if not entry.values:
+            mean_importance[entry.name] = 0.0
+            variance[entry.name] = 0.0
+            continue
+        n = len(entry.values)
+        mu = sum(entry.values) / n
+        var = sum((v - mu) ** 2 for v in entry.values) / n
+        mean_importance[entry.name] = mu
+        variance[entry.name] = var
+
+    # Identify outliers: layers with mean importance > 2 std above global mean
+    all_means = list(mean_importance.values())
+    outlier_layers: list[str] = []
+    if all_means:
+        global_mean = sum(all_means) / len(all_means)
+        global_var = sum((m - global_mean) ** 2 for m in all_means) / len(all_means)
+        global_std = math.sqrt(global_var)
+        threshold = global_mean + 2 * global_std
+        outlier_layers = [
+            name
+            for name, mu in mean_importance.items()
+            if mu > threshold
+        ]
+
+    return ImatrixAnalysis(
+        total_tensors=len(data.entries),
+        mean_importance_per_layer=mean_importance,
+        variance_per_layer=variance,
+        outlier_layers=outlier_layers,
+    )
+
+
+def format_imatrix_report(analysis: ImatrixAnalysis) -> str:
+    """Format the analysis as a human-readable text report."""
+    lines: list[str] = []
+    lines.append("=" * 60)
+    lines.append("Importance Matrix Analysis Report")
+    lines.append("=" * 60)
+    lines.append(f"Total tensors: {analysis.total_tensors}")
+    lines.append("")
+
+    # Sort layers by importance descending
+    sorted_layers = sorted(
+        analysis.mean_importance_per_layer.items(),
+        key=lambda kv: kv[1],
+        reverse=True,
+    )
+
+    lines.append("Layer Importance (descending):")
+    lines.append("-" * 60)
+    lines.append(f"{'Layer':<40} {'Mean':>10} {'Variance':>10}")
+    lines.append("-" * 60)
+    for name, mu in sorted_layers:
+        var = analysis.variance_per_layer.get(name, 0.0)
+        tag = " [OUTLIER]" if name in analysis.outlier_layers else ""
+        lines.append(f"{name:<40} {mu:>10.6f} {var:>10.6f}{tag}")
+
+    lines.append("")
+    if analysis.outlier_layers:
+        lines.append(f"Outlier layers ({len(analysis.outlier_layers)}):")
+        for name in analysis.outlier_layers:
+            lines.append(f"  - {name}")
+    else:
+        lines.append("No outlier layers detected.")
+
+    lines.append("=" * 60)
+    return "\n".join(lines)

quantbenchx/layerwise.py

@@ -0,0 +1,167 @@
+"""Layerwise analysis — sensitivity estimation and mixed-quant recommendations."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+from quantbenchx._types import DType, ModelProfile
+
+# Sensitivity heuristics based on layer position and type
+_SENSITIVITY_KEYWORDS = {
+    "embed": 0.9,
+    "lm_head": 0.9,
+    "output": 0.8,
+    "norm": 0.85,
+    "ln": 0.85,
+    "layernorm": 0.85,
+    "rmsnorm": 0.85,
+    "attn_q": 0.7,
+    "attn_k": 0.7,
+    "attn_v": 0.65,
+    "attn_output": 0.6,
+    "q_proj": 0.7,
+    "k_proj": 0.7,
+    "v_proj": 0.65,
+    "o_proj": 0.6,
+    "gate": 0.5,
+    "up": 0.5,
+    "down": 0.45,
+    "ffn": 0.45,
+    "mlp": 0.45,
+}
+
+# Position-based sensitivity multiplier (first and last layers are more sensitive)
+_POSITION_CURVE = {
+    "first_10pct": 1.3,
+    "last_10pct": 1.2,
+    "middle": 1.0,
+}
+
+
+def analyze_layers(profile: ModelProfile) -> List[Dict[str, Any]]:
+    """Analyze each layer's quantization characteristics.
+
+    Returns a list of dicts with layer name, param count, avg bpw,
+    dominant dtype, and estimated sensitivity.
+    """
+    results: List[Dict[str, Any]] = []
+    n_layers = len(profile.layers)
+
+    for i, layer in enumerate(profile.layers):
+        sensitivity = _estimate_layer_sensitivity(layer.name, i, n_layers)
+        results.append({
+            "name": layer.name,
+            "n_params": layer.n_params,
+            "n_tensors": len(layer.tensors),
+            "size_bytes": layer.size_bytes,
+            "avg_bits_per_weight": round(layer.avg_bits_per_weight, 2),
+            "dominant_dtype": layer.dominant_dtype.value,
+            "sensitivity": round(sensitivity, 3),
+        })
+
+    return results
+
+
+def layer_sensitivity(profile: ModelProfile) -> Dict[str, float]:
+    """Return a mapping of layer name → estimated sensitivity score (0-1)."""
+    n_layers = len(profile.layers)
+    return {
+        layer.name: round(_estimate_layer_sensitivity(layer.name, i, n_layers), 3)
+        for i, layer in enumerate(profile.layers)
+    }
+
+
+def recommend_mixed_quant(
+    profile: ModelProfile,
+    target_bpw: float = 4.5,
+    high_quant: DType = DType.Q5_K_M,
+    low_quant: DType = DType.Q4_K_M,
+) -> Dict[str, Any]:
+    """Recommend a mixed-precision quantization strategy.
+
+    Assigns higher precision to more sensitive layers to meet a target
+    average bits-per-weight while maximizing quality.
+    """
+    n_layers = len(profile.layers)
+    if n_layers == 0:
+        return {"strategy": [], "estimated_avg_bpw": 0.0}
+
+    # Score layers by sensitivity
+    scored: List[tuple] = []  # (sensitivity, layer_name, n_params)
+    for i, layer in enumerate(profile.layers):
+        sens = _estimate_layer_sensitivity(layer.name, i, n_layers)
+        scored.append((sens, layer.name, layer.n_params))
+
+    scored.sort(key=lambda x: x[0], reverse=True)
+
+    # Binary search for the cutoff: how many layers get high_quant
+    total_params = sum(s[2] for s in scored)
+    if total_params == 0:
+        return {"strategy": [], "estimated_avg_bpw": 0.0}
+
+    best_k = 0
+    for k in range(len(scored) + 1):
+        high_params = sum(scored[j][2] for j in range(k))
+        low_params = total_params - high_params
+        avg = (high_params * high_quant.bits_per_weight + low_params * low_quant.bits_per_weight) / total_params
+        if avg <= target_bpw:
+            best_k = k
+            break
+    else:
+        best_k = len(scored)
+
+    # Build strategy
+    high_set = {scored[j][1] for j in range(best_k)}
+    strategy = []
+    for sens, name, n_params in scored:
+        chosen = high_quant if name in high_set else low_quant
+        strategy.append({
+            "layer": name,
+            "dtype": chosen.value,
+            "sensitivity": round(sens, 3),
+            "n_params": n_params,
+        })
+
+    # Compute final average
+    high_params = sum(s[2] for s in scored[:best_k])
+    low_params = total_params - high_params
+    est_avg = (high_params * high_quant.bits_per_weight + low_params * low_quant.bits_per_weight) / total_params
+
+    return {
+        "strategy": strategy,
+        "estimated_avg_bpw": round(est_avg, 2),
+        "target_bpw": target_bpw,
+        "high_quant": high_quant.value,
+        "low_quant": low_quant.value,
+        "n_high_precision_layers": best_k,
+        "n_low_precision_layers": len(scored) - best_k,
+    }
+
+
+def _estimate_layer_sensitivity(name: str, position: int, total_layers: int) -> float:
+    """Estimate how sensitive a layer is to quantization.
+
+    Score 0-1 where 1 = most sensitive (should keep high precision).
+    Uses name-based heuristics and position in the network.
+    """
+    # Base sensitivity from name
+    base = 0.5
+    name_lower = name.lower()
+    for keyword, score in _SENSITIVITY_KEYWORDS.items():
+        if keyword in name_lower:
+            base = max(base, score)
+            break
+
+    # Position multiplier
+    if total_layers <= 1:
+        mult = 1.0
+    else:
+        frac = position / (total_layers - 1) if total_layers > 1 else 0.0
+        if frac < 0.1:
+            mult = _POSITION_CURVE["first_10pct"]
+        elif frac > 0.9:
+            mult = _POSITION_CURVE["last_10pct"]
+        else:
+            mult = _POSITION_CURVE["middle"]
+
+    return min(base * mult, 1.0)