quantbenchx 0.3.0__py3-none-any.whl

quantbenchx/__init__.py

@@ -0,0 +1,132 @@
+"""quantbenchx — Quantization quality analyzer."""
+from __future__ import annotations
+
+__version__ = "0.3.0"
+
+from quantbenchx._types import (
+    DType,
+    LayerInfo,
+    ModelProfile,
+    QualityEstimate,
+    QuantbenchError,
+    QuantFormat,
+    QuantMethod,
+    QuantProfile,
+    TensorInfo,
+)
+from quantbenchx.bandwidth import (
+    KNOWN_GPUS,
+    BandwidthEstimate,
+    BandwidthEstimator,
+    GPUSpec,
+    compare_gpus,
+    format_bandwidth_report,
+)
+from quantbenchx.compare import (
+    compare_formats,
+    compare_profiles,
+)
+from quantbenchx.imatrix import (
+    ImatrixAnalysis,
+    ImatrixData,
+    ImatrixEntry,
+    analyze_imatrix,
+    format_imatrix_report,
+    parse_imatrix,
+)
+from quantbenchx.layerwise import (
+    analyze_layers,
+    layer_sensitivity,
+    recommend_mixed_quant,
+)
+from quantbenchx.matrix import (
+    KNOWN_FORMATS,
+    ComparisonMatrix,
+    FormatComparison,
+    QuantFormatSpec,
+    format_comparison_table,
+)
+from quantbenchx.perplexity import (
+    PerplexityDelta,
+    PerplexityScore,
+    estimate_perplexity_delta,
+    format_quality_report,
+    perplexity_from_logprobs,
+)
+from quantbenchx.perplexity import (
+    quality_score as perplexity_quality_score,
+)
+from quantbenchx.predict import (
+    estimate_quality,
+    perplexity_delta,
+)
+from quantbenchx.profile import profile_from_dict, profile_gguf, profile_safetensors
+from quantbenchx.recommend import (
+    Recommendation,
+    format_recommendation,
+    recommend,
+)
+from quantbenchx.report import (
+    format_markdown,
+    format_report_rich,
+    format_report_text,
+    load_json,
+    report_to_dict,
+    save_json,
+)
+
+__all__ = [
+    "DType",
+    "ImatrixAnalysis",
+    "ImatrixData",
+    "ImatrixEntry",
+    "LayerInfo",
+    "ModelProfile",
+    "QuantFormat",
+    "QuantMethod",
+    "QuantProfile",
+    "QuantbenchError",
+    "QualityEstimate",
+    "PerplexityDelta",
+    "PerplexityScore",
+    "Recommendation",
+    "TensorInfo",
+    "analyze_imatrix",
+    "analyze_layers",
+    "compare_formats",
+    "compare_profiles",
+    "estimate_perplexity_delta",
+    "estimate_quality",
+    "format_imatrix_report",
+    "format_markdown",
+    "format_quality_report",
+    "format_recommendation",
+    "format_report_rich",
+    "format_report_text",
+    "layer_sensitivity",
+    "load_json",
+    "parse_imatrix",
+    "perplexity_delta",
+    "perplexity_from_logprobs",
+    "perplexity_quality_score",
+    "profile_from_dict",
+    "profile_gguf",
+    "profile_safetensors",
+    "recommend",
+    "recommend_mixed_quant",
+    "report_to_dict",
+    "save_json",
+    # matrix
+    "ComparisonMatrix",
+    "FormatComparison",
+    "KNOWN_FORMATS",
+    "QuantFormatSpec",
+    "format_comparison_table",
+    # bandwidth
+    "BandwidthEstimate",
+    "BandwidthEstimator",
+    "GPUSpec",
+    "KNOWN_GPUS",
+    "compare_gpus",
+    "format_bandwidth_report",
+]

quantbenchx/_types.py

@@ -0,0 +1,220 @@
+"""Core types for quantbenchx."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List
+
+
+class QuantbenchError(Exception):
+    """Base exception for quantbenchx."""
+
+
+class QuantFormat(Enum):
+    """Quantization file format."""
+    GGUF = "gguf"
+    SAFETENSORS = "safetensors"
+    PYTORCH = "pytorch"
+    UNKNOWN = "unknown"
+
+
+class QuantMethod(Enum):
+    """Quantization method/algorithm."""
+    GPTQ = "gptq"
+    AWQ = "awq"
+    GGML = "ggml"
+    BITSANDBYTES = "bitsandbytes"
+    FP16 = "fp16"
+    BF16 = "bf16"
+    FP32 = "fp32"
+    INT8 = "int8"
+    INT4 = "int4"
+    UNKNOWN = "unknown"
+
+
+class DType(Enum):
+    """Data type for tensors."""
+    F32 = "f32"
+    F16 = "f16"
+    BF16 = "bf16"
+    Q8_0 = "q8_0"
+    Q6_K = "q6_k"
+    Q5_K_M = "q5_k_m"
+    Q5_K_S = "q5_k_s"
+    Q5_1 = "q5_1"
+    Q5_0 = "q5_0"
+    Q4_K_M = "q4_k_m"
+    Q4_K_S = "q4_k_s"
+    Q4_1 = "q4_1"
+    Q4_0 = "q4_0"
+    Q3_K_M = "q3_k_m"
+    Q3_K_S = "q3_k_s"
+    Q3_K_L = "q3_k_l"
+    Q2_K = "q2_k"
+    IQ4_XS = "iq4_xs"
+    IQ3_XXS = "iq3_xxs"
+    IQ2_XXS = "iq2_xxs"
+    IQ1_S = "iq1_s"
+    UNKNOWN = "unknown"
+
+    @property
+    def bits_per_weight(self) -> float:
+        """Approximate bits per weight for this dtype."""
+        _bpw = {
+            "f32": 32.0, "f16": 16.0, "bf16": 16.0,
+            "q8_0": 8.5, "q6_k": 6.5625, "q5_k_m": 5.5,
+            "q5_k_s": 5.5, "q5_1": 5.5, "q5_0": 5.5,
+            "q4_k_m": 4.85, "q4_k_s": 4.5, "q4_1": 4.5,
+            "q4_0": 4.5, "q3_k_m": 3.9, "q3_k_s": 3.5,
+            "q3_k_l": 4.3, "q2_k": 3.35,
+            "iq4_xs": 4.25, "iq3_xxs": 3.06,
+            "iq2_xxs": 2.06, "iq1_s": 1.56,
+            "unknown": 16.0,
+        }
+        return _bpw.get(self.value, 16.0)
+
+
+@dataclass
+class TensorInfo:
+    """Information about a single tensor in a model."""
+    name: str
+    shape: List[int]
+    dtype: DType
+    n_elements: int = 0
+    size_bytes: int = 0
+
+    def __post_init__(self) -> None:
+        if self.n_elements == 0 and self.shape:
+            prod = 1
+            for s in self.shape:
+                prod *= s
+            self.n_elements = prod
+        if self.size_bytes == 0 and self.n_elements > 0:
+            self.size_bytes = int(self.n_elements * self.dtype.bits_per_weight / 8)
+
+    @property
+    def bits_per_weight(self) -> float:
+        return self.dtype.bits_per_weight
+
+    @property
+    def compression_ratio(self) -> float:
+        """Compression ratio vs FP32."""
+        return 32.0 / self.bits_per_weight if self.bits_per_weight > 0 else 1.0
+
+
+@dataclass
+class LayerInfo:
+    """Information about a model layer (group of tensors)."""
+    name: str
+    tensors: List[TensorInfo] = field(default_factory=list)
+
+    @property
+    def n_params(self) -> int:
+        return sum(t.n_elements for t in self.tensors)
+
+    @property
+    def size_bytes(self) -> int:
+        return sum(t.size_bytes for t in self.tensors)
+
+    @property
+    def avg_bits_per_weight(self) -> float:
+        total_elements = sum(t.n_elements for t in self.tensors)
+        if total_elements == 0:
+            return 0.0
+        weighted = sum(t.n_elements * t.bits_per_weight for t in self.tensors)
+        return weighted / total_elements
+
+    @property
+    def dominant_dtype(self) -> DType:
+        if not self.tensors:
+            return DType.UNKNOWN
+        dtype_counts: Dict[DType, int] = {}
+        for t in self.tensors:
+            dtype_counts[t.dtype] = dtype_counts.get(t.dtype, 0) + t.n_elements
+        return max(dtype_counts, key=lambda d: dtype_counts[d])
+
+
+@dataclass
+class QuantProfile:
+    """Quantization profile summary."""
+    method: QuantMethod = QuantMethod.UNKNOWN
+    avg_bits_per_weight: float = 0.0
+    dtype_distribution: Dict[str, float] = field(default_factory=dict)
+    n_quantized_layers: int = 0
+    n_full_precision_layers: int = 0
+    group_size: int = 0
+
+
+@dataclass
+class ModelProfile:
+    """Full model profile with all tensor/layer information."""
+    name: str
+    format: QuantFormat
+    total_params: int = 0
+    total_size_bytes: int = 0
+    tensors: List[TensorInfo] = field(default_factory=list)
+    layers: List[LayerInfo] = field(default_factory=list)
+    quant: QuantProfile = field(default_factory=QuantProfile)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def size_gb(self) -> float:
+        return self.total_size_bytes / (1024**3) if self.total_size_bytes > 0 else 0.0
+
+    @property
+    def compression_ratio(self) -> float:
+        fp32_size = self.total_params * 4
+        return fp32_size / self.total_size_bytes if self.total_size_bytes > 0 else 1.0
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "name": self.name,
+            "format": self.format.value,
+            "total_params": self.total_params,
+            "total_size_bytes": self.total_size_bytes,
+            "size_gb": round(self.size_gb, 3),
+            "compression_ratio": round(self.compression_ratio, 2),
+            "quant": {
+                "method": self.quant.method.value,
+                "avg_bits_per_weight": round(self.quant.avg_bits_per_weight, 2),
+                "dtype_distribution": self.quant.dtype_distribution,
+                "n_quantized_layers": self.quant.n_quantized_layers,
+                "n_full_precision_layers": self.quant.n_full_precision_layers,
+            },
+            "n_tensors": len(self.tensors),
+            "n_layers": len(self.layers),
+            "metadata": self.metadata,
+        }
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> ModelProfile:
+        quant_d = d.get("quant", {})
+        qp = QuantProfile(
+            method=QuantMethod(quant_d.get("method", "unknown")),
+            avg_bits_per_weight=quant_d.get("avg_bits_per_weight", 0.0),
+            dtype_distribution=quant_d.get("dtype_distribution", {}),
+            n_quantized_layers=quant_d.get("n_quantized_layers", 0),
+            n_full_precision_layers=quant_d.get("n_full_precision_layers", 0),
+        )
+        return cls(
+            name=d["name"],
+            format=QuantFormat(d.get("format", "unknown")),
+            total_params=d.get("total_params", 0),
+            total_size_bytes=d.get("total_size_bytes", 0),
+            quant=qp,
+            metadata=d.get("metadata", {}),
+        )
+
+
+@dataclass
+class QualityEstimate:
+    """Estimated quality loss from quantization."""
+    model_name: str
+    method: str
+    avg_bits_per_weight: float
+    estimated_perplexity_delta: float
+    quality_score: float  # 0-1, higher = better
+    risk_level: str  # low, medium, high, critical
+    sensitive_layers: List[str] = field(default_factory=list)
+    recommendations: List[str] = field(default_factory=list)

quantbenchx/bandwidth.py

@@ -0,0 +1,290 @@
+"""Memory bandwidth estimation and roofline analysis for quantized models."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+
+@dataclass
+class GPUSpec:
+    """Specification of a GPU's memory and compute capabilities."""
+
+    name: str
+    memory_bandwidth_gbps: float  # GB/s
+    memory_gb: float  # Total VRAM in GB
+    compute_tflops: float  # Peak FP16/BF16 TFLOPS
+
+
+@dataclass
+class BandwidthEstimate:
+    """Result of a bandwidth estimation."""
+
+    model_size_gb: float
+    gpu: str
+    transfer_time_ms: float
+    is_memory_bound: bool
+    arithmetic_intensity: float  # FLOP/byte
+    achievable_tflops: float
+
+
+# Pre-defined GPU specs (peak values from vendor datasheets)
+KNOWN_GPUS: Dict[str, GPUSpec] = {
+    "H100_SXM": GPUSpec(
+        name="H100 SXM",
+        memory_bandwidth_gbps=3350.0,
+        memory_gb=80.0,
+        compute_tflops=989.5,
+    ),
+    "H100_PCIe": GPUSpec(
+        name="H100 PCIe",
+        memory_bandwidth_gbps=2039.0,
+        memory_gb=80.0,
+        compute_tflops=756.5,
+    ),
+    "A100_80GB": GPUSpec(
+        name="A100 80GB SXM",
+        memory_bandwidth_gbps=2039.0,
+        memory_gb=80.0,
+        compute_tflops=312.0,
+    ),
+    "A100_40GB": GPUSpec(
+        name="A100 40GB",
+        memory_bandwidth_gbps=1555.0,
+        memory_gb=40.0,
+        compute_tflops=312.0,
+    ),
+    "A10G": GPUSpec(
+        name="A10G",
+        memory_bandwidth_gbps=600.0,
+        memory_gb=24.0,
+        compute_tflops=125.0,
+    ),
+    "L4": GPUSpec(
+        name="L4",
+        memory_bandwidth_gbps=300.0,
+        memory_gb=24.0,
+        compute_tflops=121.0,
+    ),
+    "T4": GPUSpec(
+        name="T4",
+        memory_bandwidth_gbps=300.0,
+        memory_gb=16.0,
+        compute_tflops=65.0,
+    ),
+    "RTX_4090": GPUSpec(
+        name="RTX 4090",
+        memory_bandwidth_gbps=1008.0,
+        memory_gb=24.0,
+        compute_tflops=330.0,
+    ),
+    "V100": GPUSpec(
+        name="V100",
+        memory_bandwidth_gbps=900.0,
+        memory_gb=16.0,
+        compute_tflops=125.0,
+    ),
+}
+
+
+class BandwidthEstimator:
+    """Estimate memory bandwidth requirements for model inference."""
+
+    def __init__(self, gpu: GPUSpec) -> None:
+        self.gpu = gpu
+
+    def estimate_transfer(self, model_size_gb: float) -> BandwidthEstimate:
+        """Estimate time to transfer model weights from HBM once.
+
+        This is the lower bound for a single forward pass in a memory-bound
+        regime — every weight must be read at least once.
+        """
+        if model_size_gb <= 0:
+            raise ValueError("model_size_gb must be positive")
+
+        transfer_time_s = model_size_gb / self.gpu.memory_bandwidth_gbps
+        transfer_time_ms = transfer_time_s * 1000.0
+
+        # A pure weight-read has zero compute → always memory-bound
+        return BandwidthEstimate(
+            model_size_gb=model_size_gb,
+            gpu=self.gpu.name,
+            transfer_time_ms=round(transfer_time_ms, 4),
+            is_memory_bound=True,
+            arithmetic_intensity=0.0,
+            achievable_tflops=0.0,
+        )
+
+    def estimate_inference(
+        self,
+        model_size_gb: float,
+        batch_size: int = 1,
+        seq_length: int = 512,
+    ) -> BandwidthEstimate:
+        """Estimate inference bandwidth characteristics.
+
+        Uses a simplified model:
+        - Bytes transferred ≈ model_size (weights read once per token)
+        - FLOPs per token ≈ 2 × parameters (one multiply + one add per param)
+        - Parameters estimated from model_size assuming 2 bytes/param (fp16 baseline).
+        """
+        if model_size_gb <= 0:
+            raise ValueError("model_size_gb must be positive")
+        if batch_size < 1:
+            raise ValueError("batch_size must be >= 1")
+        if seq_length < 1:
+            raise ValueError("seq_length must be >= 1")
+
+        model_size_bytes = model_size_gb * (1024 ** 3)
+
+        # Estimate parameter count (assume 2 bytes / param as fp16 baseline)
+        est_params = model_size_bytes / 2.0
+
+        # FLOPs per token ≈ 2 × params  (matmul dominated)
+        flops_per_token = 2.0 * est_params
+        total_flops = flops_per_token * batch_size
+
+        # Bytes read from memory per token generation step
+        bytes_read = model_size_bytes  # weights read once
+
+        # Arithmetic intensity (FLOP / byte)
+        arithmetic_intensity = total_flops / bytes_read if bytes_read > 0 else 0.0
+
+        # Roofline: compute-bound when AI > ridge point
+        ridge_point = (self.gpu.compute_tflops * 1e12) / (
+            self.gpu.memory_bandwidth_gbps * 1e9
+        )
+        is_memory_bound = arithmetic_intensity < ridge_point
+
+        if is_memory_bound:
+            # Achievable throughput limited by bandwidth
+            achievable_tflops = (
+                arithmetic_intensity * self.gpu.memory_bandwidth_gbps * 1e9
+            ) / 1e12
+        else:
+            achievable_tflops = self.gpu.compute_tflops
+
+        # Time for one decoding step (generating one token)
+        if is_memory_bound:
+            time_s = bytes_read / (self.gpu.memory_bandwidth_gbps * 1e9)
+        else:
+            time_s = total_flops / (self.gpu.compute_tflops * 1e12)
+
+        transfer_time_ms = round(time_s * 1000.0, 4)
+
+        return BandwidthEstimate(
+            model_size_gb=model_size_gb,
+            gpu=self.gpu.name,
+            transfer_time_ms=transfer_time_ms,
+            is_memory_bound=is_memory_bound,
+            arithmetic_intensity=round(arithmetic_intensity, 4),
+            achievable_tflops=round(achievable_tflops, 4),
+        )
+
+    def roofline_analysis(
+        self,
+        model_size_gb: float,
+        flops_per_token: float,
+    ) -> Dict[str, object]:
+        """Perform roofline analysis for a given workload.
+
+        Returns a dict with ridge point, operational intensity, and
+        whether the workload is compute- or memory-bound.
+        """
+        if model_size_gb <= 0:
+            raise ValueError("model_size_gb must be positive")
+        if flops_per_token <= 0:
+            raise ValueError("flops_per_token must be positive")
+
+        model_size_bytes = model_size_gb * (1024 ** 3)
+        arithmetic_intensity = flops_per_token / model_size_bytes
+
+        ridge_point = (self.gpu.compute_tflops * 1e12) / (
+            self.gpu.memory_bandwidth_gbps * 1e9
+        )
+
+        is_memory_bound = arithmetic_intensity < ridge_point
+
+        if is_memory_bound:
+            achievable_tflops = (
+                arithmetic_intensity * self.gpu.memory_bandwidth_gbps * 1e9
+            ) / 1e12
+        else:
+            achievable_tflops = self.gpu.compute_tflops
+
+        # Time per token
+        if is_memory_bound:
+            time_s = model_size_bytes / (self.gpu.memory_bandwidth_gbps * 1e9)
+        else:
+            time_s = flops_per_token / (self.gpu.compute_tflops * 1e12)
+
+        tokens_per_second = 1.0 / time_s if time_s > 0 else 0.0
+
+        return {
+            "gpu": self.gpu.name,
+            "model_size_gb": model_size_gb,
+            "ridge_point": round(ridge_point, 4),
+            "arithmetic_intensity": round(arithmetic_intensity, 4),
+            "is_memory_bound": is_memory_bound,
+            "achievable_tflops": round(achievable_tflops, 4),
+            "time_per_token_ms": round(time_s * 1000.0, 4),
+            "tokens_per_second": round(tokens_per_second, 2),
+        }
+
+    def fits_in_memory(self, model_size_gb: float) -> bool:
+        """Check if the model fits in a single GPU's memory."""
+        return model_size_gb <= self.gpu.memory_gb
+
+    def required_gpus(self, model_size_gb: float) -> int:
+        """Minimum number of GPUs required to hold the model weights."""
+        if model_size_gb <= 0:
+            raise ValueError("model_size_gb must be positive")
+        return math.ceil(model_size_gb / self.gpu.memory_gb)
+
+
+def compare_gpus(
+    model_size_gb: float,
+    gpus: Optional[List[GPUSpec]] = None,
+) -> List[BandwidthEstimate]:
+    """Compare bandwidth estimates across multiple GPUs.
+
+    If *gpus* is ``None``, all :data:`KNOWN_GPUS` are used.
+    """
+    if gpus is None:
+        gpus = list(KNOWN_GPUS.values())
+
+    results: List[BandwidthEstimate] = []
+    for gpu in gpus:
+        estimator = BandwidthEstimator(gpu)
+        est = estimator.estimate_inference(model_size_gb)
+        results.append(est)
+    return results
+
+
+def format_bandwidth_report(estimates: List[BandwidthEstimate]) -> str:
+    """Format a list of bandwidth estimates as a human-readable report."""
+    if not estimates:
+        return "No estimates to report."
+
+    lines: List[str] = []
+    lines.append("Bandwidth Estimation Report")
+    lines.append("=" * 60)
+
+    model_size = estimates[0].model_size_gb
+    lines.append(f"Model size: {model_size:.2f} GB")
+    lines.append("")
+
+    header = f"{'GPU':<20} {'Time/tok(ms)':>13} {'Bound':>8} {'AI':>8} {'TFLOPS':>8}"
+    lines.append(header)
+    lines.append("-" * 60)
+
+    for est in sorted(estimates, key=lambda e: e.transfer_time_ms):
+        bound = "MEM" if est.is_memory_bound else "COMP"
+        lines.append(
+            f"{est.gpu:<20} {est.transfer_time_ms:>13.4f} {bound:>8} "
+            f"{est.arithmetic_intensity:>8.2f} {est.achievable_tflops:>8.2f}"
+        )
+
+    lines.append("")
+    return "\n".join(lines)