splitbit_llm/model/quantization.py

"""SplitBit weight quantization — sub-byte quantization for model weights.

Reuses the SplitBit precision math concepts from inc_llm_v1.
Quantizes NumPy weight arrays to tier-appropriate format.
Dequantizes on-the-fly during matrix multiply (zero extra memory).
"""

from __future__ import annotations

import logging
import math
from typing import Any

import numpy as np

logger = logging.getLogger(__name__)

# Bits per weight for each quantization format
BPW_TABLE: dict[str, float] = {
    "ternary": math.log2(3),   # 1.585
    "q2_k": 2.0,
    "q3_k_s": 3.0,
    "q3_k_m": 3.0,
    "q4_k_s": 4.0,
    "q4_k_m": 4.0,
    "q5_k_s": 5.0,
    "q5_k_m": 5.0,
    "q6_k": 6.0,
    "q8_0": 8.0,
    "fp16": 16.0,
    "fp32": 32.0,
}


def bits_per_weight(quant_format: str) -> float:
    return BPW_TABLE.get(quant_format.lower(), 16.0)


def compression_ratio(quant_format: str, reference_bpw: float = 16.0) -> float:
    bpw = bits_per_weight(quant_format)
    return reference_bpw / bpw if bpw > 0 else 1.0


def compute_memory_footprint(param_count: int, bpw: float) -> float:
    """Memory in GB: params * bpw / 8 / 1024^3."""
    if param_count <= 0 or bpw <= 0:
        return 0.0
    return (param_count * bpw) / 8.0 / (1024 ** 3)


class SplitBitQuantizer:
    """Quantizes/dequantizes weight arrays using SplitBit sub-byte encoding.

    Supported formats:
    - ternary: weights → {-1, 0, +1} (1.585 bpw) via absmean scheme
    - q2_k: 2-bit quantization with block scaling
    - q4_k_m: 4-bit quantization with mixed block sizes
    - q8_0: 8-bit quantization with block scaling
    - fp16: no quantization (passthrough)
    """

    def __init__(self, format: str = "q4_k_m", block_size: int = 32) -> None:
        self.format = format.lower()
        self.bpw = bits_per_weight(self.format)
        self.block_size = block_size
        self._stats = {"quantized": 0, "dequantized": 0, "bytes_saved": 0}

    def quantize(self, weights: np.ndarray) -> dict[str, Any]:
        """Quantize a weight array. Returns packed data + metadata for dequantization.

        Returns dict with:
        - 'data': quantized bytes/array
        - 'shape': original shape
        - 'scale': per-block scale factors
        - 'format': quant format used
        """
        self._stats["quantized"] += 1
        original_bytes = weights.nbytes

        if self.format in ("fp16", "fp32"):
            return {
                "data": weights.astype(np.float16 if self.format == "fp16" else np.float32),
                "shape": weights.shape,
                "scale": None,
                "format": self.format,
            }

        if self.format == "ternary":
            packed = self._quantize_ternary(weights)
        elif self.bpw <= 2.0:
            packed = self._quantize_int(weights, bits=2)
        elif self.bpw <= 4.0:
            packed = self._quantize_int(weights, bits=4)
        elif self.bpw <= 8.0:
            packed = self._quantize_int(weights, bits=8)
        else:
            packed = {"data": weights.astype(np.float16), "scale": None}

        packed["shape"] = weights.shape
        packed["format"] = self.format

        quantized_bytes = packed["data"].nbytes if hasattr(packed["data"], "nbytes") else len(packed["data"])
        self._stats["bytes_saved"] += max(0, original_bytes - quantized_bytes)

        return packed

    def dequantize(self, packed: dict[str, Any]) -> np.ndarray:
        """Dequantize packed weights back to float32."""
        self._stats["dequantized"] += 1
        fmt = packed["format"]
        shape = packed["shape"]

        if fmt in ("fp16", "fp32"):
            return packed["data"].astype(np.float32).reshape(shape)

        if fmt == "ternary":
            return self._dequantize_ternary(packed, shape)

        # Integer quantization
        bits = packed.get("bits", 4)
        return self._dequantize_int(packed, shape, bits)

    def _quantize_ternary(self, weights: np.ndarray) -> dict[str, Any]:
        """Ternary quantization: weights → {-1, 0, +1} using absmean scheme.

        Based on BitNet b1.58:
          γ = average(|W|)
          W_q = RoundClip(W / γ, -1, 1)
        """
        gamma = np.mean(np.abs(weights))
        if gamma == 0:
            return {"data": np.zeros_like(weights, dtype=np.int8), "scale": 1.0}

        scaled = weights / gamma
        quantized = np.clip(np.round(scaled), -1, 1).astype(np.int8)
        # Pack as 2-bit values (-1=0, 0=1, 1=2) → 4 values per byte
        packed = (quantized + 1).astype(np.uint8)
        return {"data": packed, "scale": float(gamma)}

    def _dequantize_ternary(self, packed: dict[str, Any], shape: tuple) -> np.ndarray:
        gamma = packed["scale"]
        packed_data = packed["data"].astype(np.int8) - 1
        return (packed_data.astype(np.float32) * gamma).reshape(shape)

    def _quantize_int(self, weights: np.ndarray, bits: int = 4) -> dict[str, Any]:
        """Symmetric integer quantization with block scaling.

        Block size = 32. Each block has its own scale factor.
        """
        flat = weights.flatten().astype(np.float32)
        n = len(flat)
        block_size = self.block_size
        n_blocks = (n + block_size - 1) // block_size

        # Pad to block boundary
        pad_len = n_blocks * block_size - n
        if pad_len > 0:
            flat = np.pad(flat, (0, pad_len))

        blocks = flat.reshape(n_blocks, block_size)

        # Per-block scale: max_abs / (2^(bits-1) - 1)
        max_levels = (1 << (bits - 1)) - 1  # e.g., 7 for 4-bit, 127 for 8-bit
        scales = np.max(np.abs(blocks), axis=1, keepdims=True)
        scales = np.where(scales == 0, 1.0, scales)
        scales = scales / max_levels

        # Quantize
        quantized = np.clip(np.round(blocks / scales), -max_levels, max_levels).astype(np.int8)

        return {
            "data": quantized,
            "scale": scales.flatten().astype(np.float32),
            "bits": bits,
            "n_blocks": n_blocks,
            "block_size": block_size,
            "pad_len": pad_len,
        }

    def _dequantize_int(self, packed: dict[str, Any], shape: tuple, bits: int) -> np.ndarray:
        quantized = packed["data"].astype(np.float32)
        scales = packed["scale"]
        n_blocks = packed["n_blocks"]
        block_size = packed["block_size"]
        pad_len = packed["pad_len"]

        blocks = quantized.reshape(n_blocks, block_size)
        scales = scales.reshape(n_blocks, 1)
        dequant = blocks * scales
        flat = dequant.flatten()

        if pad_len > 0:
            flat = flat[:-pad_len]

        return flat.reshape(shape)

    def get_stats(self) -> dict[str, Any]:
        return {
            **self._stats,
            "format": self.format,
            "bpw": round(self.bpw, 4),
            "compression_ratio": round(compression_ratio(self.format), 2),
        }