"""Singularity weight quantization — sub-byte quantization for model weights. Reuses the Singularity 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 SingularityQuantizer: """Quantizes/dequantizes weight arrays using Singularity 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), }