quantbench-engine 0.1.0__py3-none-any.whl

quantbench/eval/perplexity.py

@@ -0,0 +1,143 @@
+# Copyright 2026 Prajwal7-2002
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Sliding-window causal language model perplexity evaluation engine.
+
+Mathematical Foundation:
+    Perplexity (PPL) is the standard metric for language model evaluation:
+        PPL = exp( - (1 / N) * sum_{t=1}^N log P(w_t | w_{<t}) )
+
+    Standard non-overlapping chunk evaluation splits text into discrete blocks of
+    length L, discarding past context for the initial tokens of every block and
+    artificially inflating measured perplexity.
+
+    The sliding-window approach uses a context window of length L and strides forward
+    by S < L tokens:
+        - The model sees the full preceding context of length L.
+        - Cross-entropy loss is evaluated strictly over the trailing S tokens.
+        - Boundary effects are eliminated, guaranteeing statistically rigorous
+          comparisons across FP16, INT8, INT4, and sparse model checkpoints.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def compute_sliding_window_perplexity(
+    model: nn.Module,
+    input_ids: torch.Tensor,
+    window_size: int = 512,
+    stride: int = 256,
+    device: Optional[torch.device] = None,
+) -> float:
+    """Compute exact sliding-window perplexity for a causal language model.
+
+    Args:
+        model: Causal language model producing logits of shape (B, S, V).
+        input_ids: 1D or 2D tensor of token IDs.
+        window_size: Maximum context window size L (e.g. 512 or 1024).
+        stride: Stride step S < L between consecutive evaluation windows (e.g. 256).
+        device: Device to run evaluation on.
+
+    Returns:
+        Perplexity value as a positive float.
+    """
+    if stride >= window_size:
+        raise ValueError(
+            f"stride ({stride}) must be strictly less than window_size ({window_size})"
+        )
+
+    model.eval()
+    dev = device or next(model.parameters()).device
+
+    # Ensure 1D flat token vector
+    flat_tokens = input_ids.detach().view(-1).to(dev)
+    seq_len = flat_tokens.size(0)
+
+    if seq_len <= 1:
+        raise ValueError("input_ids must contain at least 2 tokens for causal language modeling.")
+
+    nlls: list[torch.Tensor] = []
+    total_eval_tokens = 0
+
+    prev_end_loc = 0
+    for begin_loc in range(0, seq_len, stride):
+        end_loc = min(begin_loc + window_size, seq_len)
+        trg_len = end_loc - prev_end_loc  # Number of newly evaluated tokens in this window
+
+        input_chunk = flat_tokens[begin_loc:end_loc].unsqueeze(0)  # (1, chunk_len)
+        target_chunk = input_chunk.clone()
+
+        # Mask out context tokens that were already evaluated in previous windows
+        target_chunk[:, :-trg_len] = -100
+
+        with torch.no_grad():
+            outputs = model(input_chunk)
+            logits = outputs.logits if hasattr(outputs, "logits") else outputs
+
+            # Shift logits and targets for next-token causal prediction:
+            # logits[:, :-1, :] predicts targets[:, 1:]
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_targets = target_chunk[:, 1:].contiguous()
+
+            # Cross entropy loss with ignore_index=-100
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_targets.view(-1),
+                ignore_index=-100,
+                reduction="sum",
+            )
+
+        nlls.append(loss)
+        # Count non-masked target tokens
+        num_valid = (shift_targets != -100).sum().item()
+        total_eval_tokens += num_valid
+
+        prev_end_loc = end_loc
+        if end_loc == seq_len:
+            break
+
+    if total_eval_tokens == 0:
+        return float("inf")
+
+    total_nll = torch.stack(nlls).sum().item()
+    mean_nll = total_nll / total_eval_tokens
+    perplexity = math.exp(mean_nll)
+
+    return float(perplexity)
+
+
+def generate_synthetic_corpus(
+    vocab_size: int = 1000,
+    seq_length: int = 1024,
+    seed: int = 42,
+) -> torch.Tensor:
+    """Generate a deterministic synthetic token sequence for offline testing.
+
+    Args:
+        vocab_size: Vocabulary cardinality.
+        seq_length: Number of tokens to generate.
+        seed: Random seed for repeatability.
+
+    Returns:
+        Tensor of shape (seq_length,) containing token IDs.
+    """
+    g = torch.Generator().manual_seed(seed)
+    return torch.randint(low=0, high=vocab_size, size=(seq_length,), generator=g)

quantbench/eval/sqnr.py

@@ -0,0 +1,66 @@
+# Copyright 2026 Prajwal7-2002
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Signal-to-Quantization-Noise Ratio (SQNR) evaluation metrics."""
+
+from __future__ import annotations
+
+import torch
+
+
+def compute_sqnr(
+    original: torch.Tensor,
+    quantized: torch.Tensor,
+    eps: float = 1e-12,
+) -> float:
+    """Compute Signal-to-Quantization-Noise Ratio (SQNR) in decibels (dB).
+
+    Mathematical Formulation:
+        SQNR (dB) = 10 * log10( P_signal / P_noise )
+                  = 10 * log10( sum(|W|^2) / (sum(|W - W_hat|^2) + eps) )
+
+    Interpretation:
+        - Higher SQNR indicates superior numerical fidelity.
+        - FP16 -> INT8 typically achieves 40 dB to 55 dB.
+        - FP16 -> INT4 blockwise typically achieves 22 dB to 32 dB.
+        - Severe breakdown or NaN yields <= 0 dB.
+
+    Args:
+        original: Reference unquantized floating-point tensor.
+        quantized: Reconstructed/dequantized floating-point tensor of matching shape.
+        eps: Numerical stability constant to prevent division by zero for exact equality.
+
+    Returns:
+        Float value of SQNR in decibels (dB).
+    """
+    if original.shape != quantized.shape:
+        raise ValueError(
+            f"Shape mismatch in SQNR calculation: {original.shape} vs {quantized.shape}"
+        )
+
+    orig_f = original.detach().float()
+    quant_f = quantized.detach().float()
+
+    signal_power = torch.sum(orig_f.pow(2)).item()
+    noise_power = torch.sum((orig_f - quant_f).pow(2)).item()
+
+    if noise_power <= eps:
+        return 100.0  # Cap at +100 dB for near-lossless / exact reconstruction
+
+    ratio = signal_power / (noise_power + eps)
+    if ratio <= 0.0:
+        return 0.0
+
+    sqnr_db = 10.0 * torch.log10(torch.tensor(ratio)).item()
+    return float(sqnr_db)

quantbench/nn/__init__.py

@@ -0,0 +1,9 @@
+"""Drop-in quantized neural network modules and model conversion runtimes."""
+
+from quantbench.nn.model_convert import convert_model_to_quant
+from quantbench.nn.quant_linear import QuantLinear
+
+__all__ = [
+    "QuantLinear",
+    "convert_model_to_quant",
+]

quantbench/nn/model_convert.py

@@ -0,0 +1,109 @@
+"""In-place model conversion runtime for PyTorch transformer architectures.
+
+Recursively scans any PyTorch neural network (such as Hugging Face CausalLM models)
+and replaces target torch.nn.Linear layers with QuantLinear modules.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import torch.nn as nn
+
+from quantbench.nn.quant_linear import QuantLinear
+
+
+def convert_model_to_quant(
+    model: nn.Module,
+    target_modules: list[str] | None = None,
+    exclude_modules: list[str] | None = None,
+    scheme: str = "int4",
+    group_size: int = 64,
+    symmetric: bool = True,
+) -> tuple[nn.Module, dict[str, Any]]:
+    """In-place recursive conversion of model linear layers to QuantLinear.
+
+    Args:
+        model: PyTorch model (e.g. Hugging Face model or custom transformer).
+        target_modules: Optional list of substring names to convert (e.g. ['q_proj', 'v_proj']).
+            If None, all linear layers are targeted except those in exclude_modules.
+        exclude_modules: Module names to explicitly preserve unquantized (defaults to ['lm_head']).
+        scheme: Target quantization scheme ('int4', 'int8', 'ternary', 'sparse24').
+        group_size: Group size for blockwise quantization (-1 for per-channel).
+        symmetric: Whether to use symmetric quantization.
+
+    Returns:
+        Tuple of (converted_model, conversion_statistics_dictionary).
+    """
+    if exclude_modules is None:
+        exclude_modules = ["lm_head"]
+
+    t0 = time.perf_counter()
+
+    # Calculate pre-conversion memory
+    orig_bytes = sum(p.nelement() * p.element_size() for p in model.parameters())
+
+    converted_count = 0
+    skipped_count = 0
+    layer_names: list[str] = []
+
+    # Helper function to recursively replace layers
+    def _replace_layers(parent: nn.Module, prefix: str = ""):
+        nonlocal converted_count, skipped_count
+        for name, child in list(parent.named_children()):
+            full_name = f"{prefix}.{name}" if prefix else name
+
+            # Check if this child is an explicit exclusion
+            if any(exc in full_name for exc in exclude_modules):
+                skipped_count += 1
+                continue
+
+            if isinstance(child, nn.Linear):
+                # Check if target matches
+                if target_modules is not None:
+                    is_target = any(tgt in full_name for tgt in target_modules)
+                else:
+                    is_target = True
+
+                if is_target:
+                    # Convert to QuantLinear
+                    quant_child = QuantLinear.from_linear(
+                        child,
+                        scheme=scheme,
+                        group_size=group_size,
+                        symmetric=symmetric,
+                    )
+                    setattr(parent, name, quant_child)
+                    converted_count += 1
+                    layer_names.append(full_name)
+                else:
+                    skipped_count += 1
+            else:
+                _replace_layers(child, full_name)
+
+    _replace_layers(model)
+
+    dt = time.perf_counter() - t0
+
+    # Calculate post-conversion memory (buffers + parameters)
+    quant_bytes = sum(p.nelement() * p.element_size() for p in model.parameters())
+    quant_bytes += sum(b.nelement() * b.element_size() for b in model.buffers())
+
+    orig_mb = orig_bytes / (1024 * 1024)
+    quant_mb = quant_bytes / (1024 * 1024)
+    ratio = orig_mb / max(quant_mb, 1e-6)
+
+    stats = {
+        "converted_layers": converted_count,
+        "skipped_layers": skipped_count,
+        "scheme": scheme,
+        "group_size": group_size,
+        "orig_memory_mb": round(orig_mb, 2),
+        "quant_memory_mb": round(quant_mb, 2),
+        "compression_ratio": round(ratio, 2),
+        "duration_sec": round(dt, 4),
+        "layer_names": layer_names,
+    }
+
+    return model, stats

quantbench/nn/quant_linear.py

@@ -0,0 +1,296 @@
+"""QuantLinear: Drop-in PyTorch Linear layer with sub-byte packed weights.
+
+Replaces standard torch.nn.Linear by storing weights as packed integer buffers:
+- INT4: Vectorized uint8 packed nibbles (2 weights/byte) + blockwise/channel scales.
+- BitNet Ternary: 2-bit offset-binary packed uint8 buffers (4 weights/byte) + scale.
+- 2:4 Sparsity: 50% condensed float values + 2-bit position metadata nibbles.
+- INT8: Signed 8-bit integers + per-channel or per-tensor scales.
+
+Weights are registered as persistent buffers rather than trainable nn.Parameters,
+eliminating optimizer memory overhead and enabling direct serialization into safetensors.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from quantbench.quant.affine import (
+    dequantize_absmax,
+    dequantize_affine,
+    quantize_absmax,
+    quantize_affine,
+)
+from quantbench.quant.blockwise import dequantize_blockwise, quantize_blockwise
+from quantbench.quant.packing import (
+    decompress_2_4,
+    pack_int4,
+    pack_ternary,
+    prune_and_compress_2_4,
+    unpack_int4,
+    unpack_ternary,
+)
+
+
+class QuantLinear(nn.Module):
+    """Drop-in quantized linear layer storing packed weights in memory buffers.
+
+    Args:
+        in_features: Size of each input sample.
+        out_features: Size of each output sample.
+        scheme: Quantization scheme: 'int4', 'int8', 'ternary', 'sparse24'.
+        group_size: Group size for blockwise quantization (-1 for per-channel).
+        bias: Whether to include an additive bias parameter.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        scheme: str = "int4",
+        group_size: int = 64,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.scheme = scheme.lower()
+        self.group_size = group_size
+
+        if self.scheme not in ("int4", "int8", "ternary", "sparse24"):
+            raise ValueError(
+                f"Unknown scheme '{scheme}'. Supported: 'int4', 'int8', 'ternary', 'sparse24'."
+            )
+
+        # Register bias as parameter if present
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_features))
+        else:
+            self.register_parameter("bias", None)
+
+        # Buffers are allocated upon calling from_linear or loading state_dict
+        self.register_buffer("packed_weight", None)
+        self.register_buffer("scales", None)
+        self.register_buffer("zero_points", None)
+        self.register_buffer("condensed_weight", None)
+        self.register_buffer("packed_metadata", None)
+
+    def _load_from_state_dict(
+        self,
+        state_dict: dict[str, Any],
+        prefix: str,
+        local_metadata: dict[str, Any],
+        strict: bool,
+        missing_keys: list[str],
+        unexpected_keys: list[str],
+        error_msgs: list[str],
+    ) -> None:
+        """Register incoming serialized buffer tensors dynamically during load_state_dict."""
+        for buf_name in (
+            "packed_weight",
+            "scales",
+            "zero_points",
+            "condensed_weight",
+            "packed_metadata",
+        ):
+            full_key = prefix + buf_name
+            if full_key in state_dict and state_dict[full_key] is not None:
+                self.register_buffer(buf_name, state_dict[full_key])
+        super()._load_from_state_dict(
+            state_dict,
+            prefix,
+            local_metadata,
+            strict,
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+        )
+
+    @classmethod
+    def from_linear(
+        cls,
+        linear: nn.Linear,
+        scheme: str = "int4",
+        group_size: int = 64,
+        symmetric: bool = True,
+    ) -> QuantLinear:
+        """Create a QuantLinear module from an unquantized torch.nn.Linear layer.
+
+        Args:
+            linear: Source torch.nn.Linear layer.
+            scheme: Target quantization scheme: 'int4', 'int8', 'ternary', 'sparse24'.
+            group_size: Group size for blockwise quantization (e.g. 32, 64, 128, or -1).
+            symmetric: Whether to use symmetric quantization (zero-point = 0).
+
+        Returns:
+            QuantLinear instance with compressed weight buffers.
+        """
+        in_features = linear.in_features
+        out_features = linear.out_features
+        has_bias = linear.bias is not None
+
+        quant_layer = cls(
+            in_features=in_features,
+            out_features=out_features,
+            scheme=scheme,
+            group_size=group_size,
+            bias=has_bias,
+        )
+
+        w = linear.weight.data.clone().to(torch.float32)
+        if has_bias:
+            quant_layer.bias.data.copy_(linear.bias.data)
+
+        scheme_lower = scheme.lower()
+
+        if scheme_lower == "int4":
+            if group_size > 0:
+                # Group-wise INT4
+                q_blocks, s, zp, _ = quantize_blockwise(
+                    w, group_size=group_size, num_bits=4, symmetric=symmetric
+                )
+                # Pack each block along dim=-1
+                packed_w, _ = pack_int4(q_blocks, signed=symmetric, dim=-1)
+                quant_layer.register_buffer("packed_weight", packed_w)
+                quant_layer.register_buffer("scales", s)
+                if zp is not None:
+                    quant_layer.register_buffer("zero_points", zp)
+            else:
+                # Per-channel INT4 along output channels (reduce across input dim 1)
+                if symmetric:
+                    q_w, s = quantize_absmax(w, num_bits=4, dim=1)
+                    zp = None
+                else:
+                    q_w, s, zp = quantize_affine(w, num_bits=4, dim=1)
+                packed_w, _ = pack_int4(q_w, signed=symmetric, dim=-1)
+                quant_layer.register_buffer("packed_weight", packed_w)
+                quant_layer.register_buffer("scales", s)
+                if zp is not None:
+                    quant_layer.register_buffer("zero_points", zp)
+
+        elif scheme_lower == "int8":
+            if symmetric:
+                q_w, s = quantize_absmax(w, num_bits=8, dim=1)
+                zp = None
+            else:
+                q_w, s, zp = quantize_affine(w, num_bits=8, dim=1)
+            quant_layer.register_buffer("packed_weight", q_w.to(torch.int8))
+            quant_layer.register_buffer("scales", s)
+            if zp is not None:
+                quant_layer.register_buffer("zero_points", zp)
+
+        elif scheme_lower == "ternary":
+            # BitNet ternary scaling: gamma = mean(|W|)
+            gamma = w.abs().mean().clamp(min=1e-8)
+            scaled_w = w / gamma
+            ternary_w = torch.clamp(torch.round(scaled_w), -1.0, 1.0)
+            packed_w, _ = pack_ternary(ternary_w, mode="2bit", dim=-1)
+            quant_layer.register_buffer("packed_weight", packed_w)
+            quant_layer.register_buffer("scales", gamma.view(1))
+
+        elif scheme_lower in ("sparse24", "2:4"):
+            condensed, metadata = prune_and_compress_2_4(w)
+            quant_layer.register_buffer("condensed_weight", condensed)
+            quant_layer.register_buffer("packed_metadata", metadata)
+
+        return quant_layer
+
+    def dequantize_weight(self) -> torch.Tensor:
+        """Dequantize packed buffers back to a standard floating-point weight matrix."""
+        if self.scheme == "int4":
+            is_signed = self.zero_points is None
+            if self.group_size > 0:
+                # Group-wise reconstruction
+                unpacked_q = unpack_int4(
+                    self.packed_weight, original_len=self.group_size, signed=is_signed, dim=-1
+                )
+                n_elements = self.out_features * self.in_features
+                remainder = n_elements % self.group_size
+                pad_len = (self.group_size - remainder) % self.group_size
+                meta = {
+                    "original_shape": (self.out_features, self.in_features),
+                    "n_elements": n_elements,
+                    "pad_len": pad_len,
+                    "group_size": self.group_size,
+                    "symmetric": is_signed,
+                    "dtype": torch.float32,
+                    "device": self.packed_weight.device,
+                }
+                zp = self.zero_points if self.zero_points is not None else None
+                return dequantize_blockwise(unpacked_q, self.scales, zp, meta)
+            else:
+                # Per-channel reconstruction
+                unpacked_q = unpack_int4(
+                    self.packed_weight, original_len=self.in_features, signed=is_signed, dim=-1
+                )
+                if self.zero_points is not None:
+                    return dequantize_affine(unpacked_q, self.scales, self.zero_points)
+                else:
+                    return dequantize_absmax(unpacked_q, self.scales)
+
+        elif self.scheme == "int8":
+            if self.zero_points is not None:
+                return dequantize_affine(self.packed_weight, self.scales, self.zero_points)
+            else:
+                return dequantize_absmax(self.packed_weight, self.scales)
+
+        elif self.scheme == "ternary":
+            unpacked_t = unpack_ternary(
+                self.packed_weight, original_len=self.in_features, mode="2bit", dim=-1
+            )
+            return unpacked_t.to(torch.float32) * self.scales
+
+        elif self.scheme in ("sparse24", "2:4"):
+            return decompress_2_4(
+                self.condensed_weight,
+                self.packed_metadata,
+                original_shape=(self.out_features, self.in_features),
+            )
+
+        raise RuntimeError(f"Unsupported scheme '{self.scheme}'.")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass with dynamic vector dequantization.
+
+        Args:
+            x: Input activation tensor with shape (..., in_features).
+
+        Returns:
+            Projected activation tensor with shape (..., out_features).
+        """
+        w_hat = self.dequantize_weight().to(x.dtype)
+        return F.linear(x, w_hat, self.bias)
+
+    def memory_footprint_bytes(self) -> dict[str, float]:
+        """Compute exact memory footprint in bytes and compression factor relative to FP16.
+
+        Returns:
+            Dictionary containing 'compressed_bytes', 'fp16_bytes', and 'compression_ratio'.
+        """
+        compressed_bytes = 0
+        for buf in self.buffers():
+            if buf is not None:
+                compressed_bytes += buf.nelement() * buf.element_size()
+        if self.bias is not None:
+            compressed_bytes += self.bias.nelement() * self.bias.element_size()
+
+        fp16_bytes = (self.in_features * self.out_features * 2) + (
+            (self.out_features * 2) if self.bias is not None else 0
+        )
+        ratio = fp16_bytes / max(compressed_bytes, 1)
+
+        return {
+            "compressed_bytes": compressed_bytes,
+            "fp16_bytes": fp16_bytes,
+            "compression_ratio": ratio,
+        }
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"scheme={self.scheme}, group_size={self.group_size}, "
+            f"bias={self.bias is not None}"
+        )

quantbench/quant/__init__.py

@@ -0,0 +1,58 @@
+"""First-principles quantization primitives, algorithms, and bit-packing runtimes."""
+
+from quantbench.quant.affine import (
+    dequantize_absmax,
+    dequantize_affine,
+    quantize_absmax,
+    quantize_affine,
+)
+from quantbench.quant.algorithms import (
+    awq_quantize_weights,
+    compute_activation_channel_means,
+    compute_activation_hessian,
+    compute_smoothquant_scale,
+    extract_activation_channel_scales,
+    gptq_quantize_weights,
+    search_best_awq_scale,
+    smooth_linear_weights,
+)
+from quantbench.quant.blockwise import dequantize_blockwise, quantize_blockwise
+from quantbench.quant.fp8 import FP8Format, dequantize_fp8, quantize_fp8
+from quantbench.quant.packing import (
+    decompress_2_4,
+    pack_2_4_indices,
+    pack_int4,
+    pack_ternary,
+    prune_and_compress_2_4,
+    unpack_2_4_indices,
+    unpack_int4,
+    unpack_ternary,
+)
+
+__all__ = [
+    "quantize_affine",
+    "dequantize_affine",
+    "quantize_absmax",
+    "dequantize_absmax",
+    "quantize_blockwise",
+    "dequantize_blockwise",
+    "quantize_fp8",
+    "dequantize_fp8",
+    "FP8Format",
+    "pack_int4",
+    "unpack_int4",
+    "pack_ternary",
+    "unpack_ternary",
+    "pack_2_4_indices",
+    "unpack_2_4_indices",
+    "prune_and_compress_2_4",
+    "decompress_2_4",
+    "compute_smoothquant_scale",
+    "extract_activation_channel_scales",
+    "smooth_linear_weights",
+    "compute_activation_channel_means",
+    "search_best_awq_scale",
+    "awq_quantize_weights",
+    "compute_activation_hessian",
+    "gptq_quantize_weights",
+]