quantbenchx 0.3.0__py3-none-any.whl

quantbenchx/profile.py

@@ -0,0 +1,301 @@
+"""Profile quantized models from GGUF, safetensors, or dict metadata."""
+
+from __future__ import annotations
+
+import json
+import struct
+from pathlib import Path
+from typing import Any, BinaryIO, Dict, List
+
+from quantbenchx._types import (
+    DType,
+    LayerInfo,
+    ModelProfile,
+    QuantbenchError,
+    QuantFormat,
+    QuantMethod,
+    QuantProfile,
+    TensorInfo,
+)
+
+# ── GGUF constants ──
+GGUF_MAGIC = 0x46475547  # "GGUF" in little-endian
+GGUF_DTYPE_MAP = {
+    0: DType.F32, 1: DType.F16, 2: DType.Q4_0, 3: DType.Q4_1,
+    6: DType.Q5_0, 7: DType.Q5_1, 8: DType.Q8_0,
+    10: DType.Q2_K, 11: DType.Q3_K_S, 12: DType.Q3_K_M,
+    13: DType.Q3_K_L, 14: DType.Q4_K_S, 15: DType.Q4_K_M,
+    16: DType.Q5_K_S, 17: DType.Q5_K_M, 18: DType.Q6_K,
+    19: DType.IQ2_XXS, 20: DType.IQ3_XXS, 26: DType.IQ4_XS,
+    24: DType.IQ1_S, 28: DType.BF16,
+}
+
+# Type tags for GGUF metadata values
+_GGUF_TYPE_UINT8 = 0
+_GGUF_TYPE_INT8 = 1
+_GGUF_TYPE_UINT16 = 2
+_GGUF_TYPE_INT16 = 3
+_GGUF_TYPE_UINT32 = 4
+_GGUF_TYPE_INT32 = 5
+_GGUF_TYPE_FLOAT32 = 6
+_GGUF_TYPE_BOOL = 7
+_GGUF_TYPE_STRING = 8
+_GGUF_TYPE_ARRAY = 9
+_GGUF_TYPE_UINT64 = 10
+_GGUF_TYPE_INT64 = 11
+_GGUF_TYPE_FLOAT64 = 12
+
+# safetensors dtype sizes
+_SAFETENSORS_DTYPE_BPW = {
+    "F32": 32, "F16": 16, "BF16": 16, "I64": 64, "I32": 32,
+    "I16": 16, "I8": 8, "U8": 8, "BOOL": 1, "F8_E4M3": 8, "F8_E5M2": 8,
+}
+_SAFETENSORS_DTYPE_MAP = {
+    "F32": DType.F32, "F16": DType.F16, "BF16": DType.BF16,
+    "I8": DType.Q8_0, "U8": DType.Q8_0,
+}
+
+
+def _read_gguf_string(f: BinaryIO) -> str:
+    """Read a GGUF string (uint64 length + bytes)."""
+    length = struct.unpack("<Q", f.read(8))[0]
+    return f.read(length).decode("utf-8", errors="replace")
+
+
+def _read_gguf_value(f: BinaryIO, vtype: int) -> Any:
+    """Read a single GGUF metadata value."""
+    if vtype == _GGUF_TYPE_UINT8:
+        return struct.unpack("<B", f.read(1))[0]
+    elif vtype == _GGUF_TYPE_INT8:
+        return struct.unpack("<b", f.read(1))[0]
+    elif vtype == _GGUF_TYPE_UINT16:
+        return struct.unpack("<H", f.read(2))[0]
+    elif vtype == _GGUF_TYPE_INT16:
+        return struct.unpack("<h", f.read(2))[0]
+    elif vtype == _GGUF_TYPE_UINT32:
+        return struct.unpack("<I", f.read(4))[0]
+    elif vtype == _GGUF_TYPE_INT32:
+        return struct.unpack("<i", f.read(4))[0]
+    elif vtype == _GGUF_TYPE_FLOAT32:
+        return struct.unpack("<f", f.read(4))[0]
+    elif vtype == _GGUF_TYPE_BOOL:
+        return struct.unpack("<B", f.read(1))[0] != 0
+    elif vtype == _GGUF_TYPE_STRING:
+        return _read_gguf_string(f)
+    elif vtype == _GGUF_TYPE_UINT64:
+        return struct.unpack("<Q", f.read(8))[0]
+    elif vtype == _GGUF_TYPE_INT64:
+        return struct.unpack("<q", f.read(8))[0]
+    elif vtype == _GGUF_TYPE_FLOAT64:
+        return struct.unpack("<d", f.read(8))[0]
+    elif vtype == _GGUF_TYPE_ARRAY:
+        elem_type = struct.unpack("<I", f.read(4))[0]
+        count = struct.unpack("<Q", f.read(8))[0]
+        return [_read_gguf_value(f, elem_type) for _ in range(count)]
+    else:
+        raise QuantbenchError(f"Unknown GGUF value type: {vtype}")
+
+
+def profile_gguf(path: str | Path) -> ModelProfile:
+    """Parse a GGUF file and return a model profile.
+
+    This is a pure-Python parser — no external dependencies required.
+    Reads only the header (metadata + tensor info), not the actual weights.
+    """
+    path = Path(path)
+    if not path.exists():
+        raise QuantbenchError(f"File not found: {path}")
+
+    with open(path, "rb") as f:
+        # Read magic number
+        magic = struct.unpack("<I", f.read(4))[0]
+        if magic != GGUF_MAGIC:
+            raise QuantbenchError(f"Not a GGUF file (magic: {magic:#x})")
+
+        # Read version
+        version = struct.unpack("<I", f.read(4))[0]
+        if version not in (2, 3):
+            raise QuantbenchError(f"Unsupported GGUF version: {version}")
+
+        # Read counts
+        n_tensors = struct.unpack("<Q", f.read(8))[0]
+        n_kv = struct.unpack("<Q", f.read(8))[0]
+
+        # Read metadata
+        metadata: Dict[str, Any] = {}
+        for _ in range(n_kv):
+            key = _read_gguf_string(f)
+            vtype = struct.unpack("<I", f.read(4))[0]
+            value = _read_gguf_value(f, vtype)
+            metadata[key] = value
+
+        # Read tensor info
+        tensors: List[TensorInfo] = []
+        for _ in range(n_tensors):
+            name = _read_gguf_string(f)
+            n_dims = struct.unpack("<I", f.read(4))[0]
+            shape = [struct.unpack("<Q", f.read(8))[0] for _ in range(n_dims)]
+            dtype_id = struct.unpack("<I", f.read(4))[0]
+            _offset = struct.unpack("<Q", f.read(8))[0]  # tensor data offset
+
+            dtype = GGUF_DTYPE_MAP.get(dtype_id, DType.UNKNOWN)
+            tensors.append(TensorInfo(name=name, shape=shape, dtype=dtype))
+
+    # Build layers from tensor names
+    layers = _group_tensors_into_layers(tensors)
+
+    # Build quant profile
+    quant = _build_quant_profile(tensors, metadata)
+
+    total_params = sum(t.n_elements for t in tensors)
+    total_size = sum(t.size_bytes for t in tensors)
+
+    model_name = metadata.get("general.name", path.stem)
+    if not isinstance(model_name, str):
+        model_name = str(model_name)
+
+    return ModelProfile(
+        name=model_name,
+        format=QuantFormat.GGUF,
+        total_params=total_params,
+        total_size_bytes=total_size,
+        tensors=tensors,
+        layers=layers,
+        quant=quant,
+        metadata={k: v for k, v in metadata.items() if isinstance(v, (str, int, float, bool))},
+    )
+
+
+def profile_safetensors(path: str | Path) -> ModelProfile:
+    """Parse a safetensors file header and return a model profile.
+
+    Only reads the JSON header — does not load weights into memory.
+    """
+    path = Path(path)
+    if not path.exists():
+        raise QuantbenchError(f"File not found: {path}")
+
+    with open(path, "rb") as f:
+        # First 8 bytes: header size as uint64
+        header_size = struct.unpack("<Q", f.read(8))[0]
+        if header_size > 100_000_000:  # sanity check: 100MB max header
+            raise QuantbenchError(f"Header too large: {header_size}")
+
+        header_bytes = f.read(header_size)
+        header = json.loads(header_bytes)
+
+    # Extract tensor metadata (skip __metadata__ key)
+    tensors: List[TensorInfo] = []
+    meta = header.pop("__metadata__", {})
+
+    for tensor_name, info in header.items():
+        dtype_str = info.get("dtype", "F32")
+        shape = info.get("shape", [])
+        offsets = info.get("data_offsets", [0, 0])
+
+        dtype = _SAFETENSORS_DTYPE_MAP.get(dtype_str, DType.UNKNOWN)
+        size_bytes = offsets[1] - offsets[0] if len(offsets) == 2 else 0
+
+        ti = TensorInfo(name=tensor_name, shape=shape, dtype=dtype, size_bytes=size_bytes)
+        tensors.append(ti)
+
+    layers = _group_tensors_into_layers(tensors)
+    quant = _build_quant_profile(tensors, meta)
+
+    total_params = sum(t.n_elements for t in tensors)
+    total_size = sum(t.size_bytes for t in tensors)
+
+    return ModelProfile(
+        name=path.stem,
+        format=QuantFormat.SAFETENSORS,
+        total_params=total_params,
+        total_size_bytes=total_size,
+        tensors=tensors,
+        layers=layers,
+        quant=quant,
+        metadata=meta if isinstance(meta, dict) else {},
+    )
+
+
+def profile_from_dict(data: Dict[str, Any]) -> ModelProfile:
+    """Create a ModelProfile from a dictionary (e.g. loaded from JSON)."""
+    return ModelProfile.from_dict(data)
+
+
+def _group_tensors_into_layers(tensors: List[TensorInfo]) -> List[LayerInfo]:
+    """Group tensors into logical layers based on naming conventions."""
+    layer_map: Dict[str, List[TensorInfo]] = {}
+
+    for t in tensors:
+        # Extract layer identifier from tensor name
+        # Common patterns: "blk.0.attn_q.weight", "model.layers.0.self_attn.q_proj.weight"
+        parts = t.name.split(".")
+        layer_name = "other"
+
+        for i, part in enumerate(parts):
+            if part.isdigit() and i > 0:
+                layer_name = ".".join(parts[: i + 1])
+                break
+            if part in ("blk", "layers", "block", "h"):
+                if i + 1 < len(parts) and parts[i + 1].isdigit():
+                    layer_name = ".".join(parts[: i + 2])
+                    break
+
+        if layer_name not in layer_map:
+            layer_map[layer_name] = []
+        layer_map[layer_name].append(t)
+
+    return [LayerInfo(name=name, tensors=ts) for name, ts in sorted(layer_map.items())]
+
+
+def _build_quant_profile(tensors: List[TensorInfo], metadata: Dict[str, Any]) -> QuantProfile:
+    """Build quantization profile from tensor list and metadata."""
+    if not tensors:
+        return QuantProfile()
+
+    # Count dtype distribution
+    total_elements = sum(t.n_elements for t in tensors)
+    dtype_counts: Dict[str, int] = {}
+    for t in tensors:
+        key = t.dtype.value
+        dtype_counts[key] = dtype_counts.get(key, 0) + t.n_elements
+
+    dtype_dist = {}
+    if total_elements > 0:
+        dtype_dist = {k: round(v / total_elements, 4) for k, v in dtype_counts.items()}
+
+    # Compute average bits per weight
+    avg_bpw = 0.0
+    if total_elements > 0:
+        avg_bpw = sum(t.n_elements * t.bits_per_weight for t in tensors) / total_elements
+
+    # Count quantized vs full precision layers
+    fp_dtypes = {DType.F32, DType.F16, DType.BF16}
+    n_quant = sum(1 for t in tensors if t.dtype not in fp_dtypes)
+    n_fp = sum(1 for t in tensors if t.dtype in fp_dtypes)
+
+    # Detect method from metadata
+    method = QuantMethod.UNKNOWN
+    meta_str = str(metadata).lower()
+    if "gptq" in meta_str:
+        method = QuantMethod.GPTQ
+    elif "awq" in meta_str:
+        method = QuantMethod.AWQ
+    elif "ggml" in meta_str or "llama.cpp" in meta_str:
+        method = QuantMethod.GGML
+    elif "bitsandbytes" in meta_str:
+        method = QuantMethod.BITSANDBYTES
+    elif avg_bpw > 15:
+        method = QuantMethod.FP16
+    elif avg_bpw > 7:
+        method = QuantMethod.INT8
+    elif avg_bpw > 3:
+        method = QuantMethod.INT4
+
+    return QuantProfile(
+        method=method,
+        avg_bits_per_weight=round(avg_bpw, 2),
+        dtype_distribution=dtype_dist,
+        n_quantized_layers=n_quant,
+        n_full_precision_layers=n_fp,
+    )

quantbenchx/recommend.py

@@ -0,0 +1,240 @@
+"""Quantization recommendation engine — suggest optimal format based on model profile."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+from quantbenchx._types import DType, ModelProfile
+from quantbenchx.layerwise import layer_sensitivity
+
+
+@dataclass
+class Recommendation:
+    """Quantization recommendation for a model."""
+
+    format: str
+    estimated_size_gb: float
+    estimated_quality: float  # 0-1 quality retention score
+    per_layer: Dict[str, str]  # layer name -> recommended format
+    explanation: str
+
+
+# Candidate formats ordered from highest quality to most compressed.
+# Each entry: (label, target DType, approx bpw, quality baseline)
+_CANDIDATES: List[tuple] = [
+    ("Q8_0", DType.Q8_0, 8.5, 0.99),
+    ("Q6_K", DType.Q6_K, 6.5625, 0.97),
+    ("Q5_K_M", DType.Q5_K_M, 5.5, 0.95),
+    ("Q5_K_S", DType.Q5_K_S, 5.5, 0.94),
+    ("Q4_K_M", DType.Q4_K_M, 4.85, 0.90),
+    ("Q4_K_S", DType.Q4_K_S, 4.5, 0.87),
+    ("Q3_K_M", DType.Q3_K_M, 3.9, 0.80),
+    ("Q3_K_S", DType.Q3_K_S, 3.5, 0.74),
+    ("Q2_K", DType.Q2_K, 3.35, 0.60),
+]
+
+# Layers matching these keywords are kept at a higher precision tier.
+_HIGH_PRECISION_KEYWORDS = {"embed", "lm_head", "output", "norm", "layernorm", "rmsnorm"}
+
+
+def recommend(
+    profile: ModelProfile,
+    target_bits: Optional[float] = None,
+    max_quality_loss: float = 0.05,
+) -> Recommendation:
+    """Recommend a quantization strategy for *profile*.
+
+    Parameters
+    ----------
+    profile:
+        A :class:`ModelProfile` produced by the profiling helpers.
+    target_bits:
+        If given, pick the format closest to this bits-per-weight.
+        Overrides *max_quality_loss*.
+    max_quality_loss:
+        Maximum tolerable quality loss (0-1).  ``0.05`` means we want
+        ≥ 95 % quality retention.  Ignored when *target_bits* is set.
+
+    Returns
+    -------
+    Recommendation
+    """
+    total_params = profile.total_params
+    if total_params == 0:
+        total_params = sum(l.n_params for l in profile.layers)
+
+    # --- pick the overall format ----------------------------------------
+    if target_bits is not None:
+        chosen = _closest_candidate(target_bits)
+    else:
+        min_quality = 1.0 - max_quality_loss
+        chosen = _best_candidate_for_quality(min_quality)
+
+    label, dtype, bpw, quality_base = chosen
+
+    # --- estimate size --------------------------------------------------
+    if total_params > 0:
+        estimated_size_bytes = total_params * bpw / 8
+        estimated_size_gb = round(estimated_size_bytes / (1024**3), 3)
+    else:
+        estimated_size_gb = 0.0
+
+    # --- per-layer recommendations (mixed-quant) ------------------------
+    sensitivities = layer_sensitivity(profile) if profile.layers else {}
+    high_dtype = _one_tier_higher(dtype)
+    per_layer: Dict[str, str] = {}
+    for layer_name, sens in sensitivities.items():
+        if _is_sensitive_layer(layer_name) or sens >= 0.75:
+            per_layer[layer_name] = high_dtype.value
+        else:
+            per_layer[layer_name] = dtype.value
+
+    # --- adjust quality estimate for mixed-quant benefit -----------------
+    estimated_quality = _adjust_quality(quality_base, per_layer, sensitivities)
+
+    # --- human-readable explanation -------------------------------------
+    explanation = _build_explanation(
+        profile, label, bpw, estimated_size_gb, estimated_quality,
+        per_layer, sensitivities, target_bits, max_quality_loss,
+    )
+
+    return Recommendation(
+        format=label,
+        estimated_size_gb=estimated_size_gb,
+        estimated_quality=round(estimated_quality, 4),
+        per_layer=per_layer,
+        explanation=explanation,
+    )
+
+
+def format_recommendation(rec: Recommendation) -> str:
+    """Return a human-readable multi-line summary of *rec*."""
+    lines = [
+        f"Recommended format : {rec.format}",
+        f"Estimated size     : {rec.estimated_size_gb:.2f} GB",
+        f"Quality retention  : {rec.estimated_quality * 100:.1f}%",
+        "",
+    ]
+    if rec.per_layer:
+        lines.append("Per-layer strategy:")
+        for layer, fmt in rec.per_layer.items():
+            lines.append(f"  {layer}: {fmt}")
+        lines.append("")
+    lines.append(rec.explanation)
+    return "\n".join(lines)
+
+
+# ── internal helpers ─────────────────────────────────────────────────────
+
+
+def _closest_candidate(target_bpw: float) -> tuple:
+    best = _CANDIDATES[0]
+    best_dist = abs(best[2] - target_bpw)
+    for c in _CANDIDATES[1:]:
+        d = abs(c[2] - target_bpw)
+        if d < best_dist:
+            best = c
+            best_dist = d
+    return best
+
+
+def _best_candidate_for_quality(min_quality: float) -> tuple:
+    # Walk from most compressed to least; pick the most compressed that
+    # still meets the quality threshold.
+    for c in reversed(_CANDIDATES):
+        if c[3] >= min_quality:
+            return c
+    # Nothing meets the bar — fall back to highest quality.
+    return _CANDIDATES[0]
+
+
+def _one_tier_higher(dtype: DType) -> DType:
+    """Return the next-higher-quality DType tier."""
+    order = [c[1] for c in _CANDIDATES]
+    try:
+        idx = order.index(dtype)
+    except ValueError:
+        return dtype
+    return order[max(0, idx - 1)]
+
+
+def _is_sensitive_layer(name: str) -> bool:
+    name_lower = name.lower()
+    return any(kw in name_lower for kw in _HIGH_PRECISION_KEYWORDS)
+
+
+def _adjust_quality(
+    base: float,
+    per_layer: Dict[str, str],
+    sensitivities: Dict[str, float],
+) -> float:
+    """Bump quality estimate when sensitive layers get higher precision."""
+    if not per_layer or not sensitivities:
+        return base
+
+    total_sens = sum(sensitivities.values())
+    if total_sens == 0:
+        return base
+
+    # Fraction of total sensitivity weight that got upgraded.
+    upgraded_sens = sum(
+        sensitivities.get(ln, 0.0)
+        for ln, fmt in per_layer.items()
+        # "upgraded" if the format differs from the majority
+        if fmt != _mode_value(per_layer)
+    )
+    benefit = 0.02 * (upgraded_sens / total_sens)
+    return min(base + benefit, 1.0)
+
+
+def _mode_value(d: Dict[str, str]) -> str:
+    """Most common value in a dict."""
+    counts: Dict[str, int] = {}
+    for v in d.values():
+        counts[v] = counts.get(v, 0) + 1
+    return max(counts, key=lambda k: counts[k]) if counts else ""
+
+
+def _build_explanation(
+    profile: ModelProfile,
+    label: str,
+    bpw: float,
+    size_gb: float,
+    quality: float,
+    per_layer: Dict[str, str],
+    sensitivities: Dict[str, float],
+    target_bits: Optional[float],
+    max_quality_loss: float,
+) -> str:
+    parts: List[str] = []
+    param_b = profile.total_params / 1e9 if profile.total_params else 0
+    if param_b > 0:
+        parts.append(f"For a ~{param_b:.1f}B-parameter model, ")
+    else:
+        parts.append("Based on the model profile, ")
+
+    if target_bits is not None:
+        parts.append(
+            f"{label} (≈{bpw:.1f} bpw) is the closest match to the "
+            f"requested {target_bits:.1f} bits-per-weight target."
+        )
+    else:
+        parts.append(
+            f"{label} (≈{bpw:.1f} bpw) is recommended to stay within "
+            f"{max_quality_loss * 100:.0f}% quality loss."
+        )
+
+    n_upgraded = sum(
+        1 for fmt in per_layer.values() if fmt != _mode_value(per_layer)
+    )
+    if n_upgraded:
+        parts.append(
+            f" {n_upgraded} sensitive layer(s) are assigned higher precision "
+            f"for mixed-quantization."
+        )
+
+    parts.append(
+        f" Estimated size: {size_gb:.2f} GB, quality retention: {quality * 100:.1f}%."
+    )
+    return "".join(parts)