zeroquantz 0.1.0__py3-none-any.whl

zeroquantz/profiling/sensitivity.py

@@ -0,0 +1,234 @@
+"""Layer-sensitivity profiling.
+
+Estimates how much each transformer sub-module (attention / MLP) degrades model
+outputs when pushed to low precision, so the mixed-precision planner can keep the
+sensitive ones higher. Two paths:
+
+* :meth:`SensitivityProfiler.heuristic` — a zero-cost prior (edge layers are most
+  sensitive). Always available; used offline and as a fallback.
+* :meth:`SensitivityProfiler.profile` — the real measurement (spec §11): baseline
+  outputs, then per-module simulated quantization, measure deviation, restore.
+  Requires torch + transformers and is imported lazily.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, Field
+
+from zeroquantz.core.exceptions import DependencyError
+from zeroquantz.utils import metrics
+from zeroquantz.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from zeroquantz.models.metadata import ModelProfile
+    from zeroquantz.profiling.calibration import CalibrationDataset
+
+log = get_logger(__name__)
+
+
+def layer_sensitivity_prior(layer_idx: int, num_layers: int, kind: str = "attn") -> float:
+    """Prior sensitivity in [0, 1]: layers near input/output are most sensitive,
+    attention slightly more than MLP."""
+    if num_layers <= 1:
+        base = 1.0
+    else:
+        edge_distance = min(layer_idx, num_layers - 1 - layer_idx)
+        base = 1.0 - edge_distance / ((num_layers - 1) / 2)
+    base = max(0.0, min(1.0, base))
+    return round(base * (1.0 if kind == "attn" else 0.9), 4)
+
+
+class LayerSensitivity(BaseModel):
+    layer: str
+    score: float  # normalized 0..1, higher = more sensitive
+    raw: float | None = None  # raw deviation metric, if measured
+    metric: str = "prior"
+
+
+class SensitivityProfile(BaseModel):
+    model_id: str
+    metric: str = "prior"
+    target_precision: str = "int4"
+    layers: list[LayerSensitivity] = Field(default_factory=list)
+
+    def as_lookup(self) -> dict[str, float]:
+        return {ls.layer: ls.score for ls in self.layers}
+
+    def ranked(self) -> list[LayerSensitivity]:
+        return sorted(self.layers, key=lambda ls: ls.score, reverse=True)
+
+    def top(self, n: int) -> list[LayerSensitivity]:
+        return self.ranked()[:n]
+
+
+def _normalize_scores(raw: dict[str, float]) -> dict[str, float]:
+    if not raw:
+        return {}
+    lo, hi = min(raw.values()), max(raw.values())
+    if hi <= lo:
+        return {k: 0.5 for k in raw}
+    return {k: round((v - lo) / (hi - lo), 4) for k, v in raw.items()}
+
+
+class SensitivityProfiler:
+    """Produce a :class:`SensitivityProfile` for a model."""
+
+    @staticmethod
+    def heuristic(model: ModelProfile) -> SensitivityProfile:
+        layers = model.num_layers or 0
+        entries: list[LayerSensitivity] = []
+        for i in range(layers):
+            for kind, mod in (("attn", "self_attn"), ("mlp", "mlp")):
+                entries.append(
+                    LayerSensitivity(
+                        layer=f"model.layers.{i}.{mod}",
+                        score=layer_sensitivity_prior(i, layers, kind),
+                        metric="prior",
+                    )
+                )
+        return SensitivityProfile(model_id=model.model_id, metric="prior", layers=entries)
+
+    @staticmethod
+    def profile(
+        model_ref: str,
+        calibration: CalibrationDataset | None = None,
+        *,
+        metric: str = "mse",
+        target_bits: int = 4,
+        group_size: int = 128,
+        max_layers: int | None = None,
+        max_length: int = 128,
+        progress: Callable[[str, float], None] | None = None,
+    ) -> SensitivityProfile:
+        """Measure per-module sensitivity (requires torch + transformers).
+
+        For each layer's attention and MLP block, all linear weights are replaced
+        with a simulated-quantized round-trip, the deviation of the output logits
+        from baseline is measured, and the original weights are restored.
+        """
+        try:
+            import torch
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+        except ImportError as exc:
+            raise DependencyError.for_extra(
+                "transformers", "torch", purpose="profile layer sensitivity"
+            ) from exc
+
+        from zeroquantz.profiling.calibration import default_calibration
+
+        calibration = calibration or default_calibration()
+
+        if progress:
+            progress("Loading model", 0.02)
+        tokenizer = AutoTokenizer.from_pretrained(model_ref)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_ref, torch_dtype=torch.float16, device_map="auto"
+        )
+        model.eval()
+        device = next(model.parameters()).device
+
+        prompts = calibration.sample(min(len(calibration), 8))
+        enc = tokenizer(
+            prompts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=max_length,
+        ).to(device)
+
+        with torch.no_grad():
+            baseline = model(**enc).logits.float().cpu().numpy()
+
+        blocks = SensitivityProfiler._decoder_blocks(model)
+        if max_layers is not None:
+            blocks = blocks[:max_layers]
+
+        raw: dict[str, float] = {}
+        total = len(blocks) * 2 or 1
+        step = 0
+        with torch.no_grad():
+            for layer_idx, block in blocks:
+                for kind, submodule_name in (("attn", "self_attn"), ("mlp", "mlp")):
+                    submodule = getattr(block, submodule_name, None)
+                    if submodule is None:
+                        step += 1
+                        continue
+                    saved = SensitivityProfiler._fake_quantize_submodule(
+                        submodule, torch, target_bits, group_size
+                    )
+                    logits = model(**enc).logits.float().cpu().numpy()
+                    raw[f"model.layers.{layer_idx}.{submodule_name}"] = _deviation(
+                        baseline, logits, metric
+                    )
+                    SensitivityProfiler._restore(submodule, saved)
+                    step += 1
+                    if progress:
+                        progress(f"layer {layer_idx} {kind}", step / total)
+
+        normalized = _normalize_scores(raw)
+        entries = [
+            LayerSensitivity(layer=k, score=v, raw=raw[k], metric=metric)
+            for k, v in normalized.items()
+        ]
+        return SensitivityProfile(
+            model_id=model_ref,
+            metric=metric,
+            target_precision=f"int{target_bits}",
+            layers=entries,
+        )
+
+    # ---- torch helpers (only reached on the real path) ----------------------
+
+    @staticmethod
+    def _decoder_blocks(model: Any) -> list[tuple[int, Any]]:
+        for attr in ("model", "transformer", "gpt_neox"):
+            inner = getattr(model, attr, None)
+            layers = getattr(inner, "layers", None) if inner is not None else None
+            if layers is not None:
+                return list(enumerate(layers))
+        layers = getattr(model, "layers", None)
+        return list(enumerate(layers)) if layers is not None else []
+
+    @staticmethod
+    def _fake_quantize_submodule(submodule, torch, bits: int, group_size: int) -> dict:
+        saved: dict[str, Any] = {}
+        for name, param in submodule.named_parameters(recurse=True):
+            if param.dim() != 2:  # only quantize linear weight matrices
+                continue
+            saved[name] = param.detach().clone()
+            param.copy_(_fake_quantize_tensor(param.detach(), torch, bits, group_size))
+        return saved
+
+    @staticmethod
+    def _restore(submodule, saved: dict) -> None:
+        params = dict(submodule.named_parameters(recurse=True))
+        for name, value in saved.items():
+            params[name].copy_(value)
+
+
+def _fake_quantize_tensor(weight, torch, bits: int, group_size: int):
+    """Symmetric per-group fake quantization round-trip (dequantized result)."""
+    qmax = 2 ** (bits - 1) - 1
+    w = weight.float()
+    out_features, in_features = w.shape
+    gs = group_size if group_size and in_features % group_size == 0 else in_features
+    w = w.reshape(out_features, in_features // gs, gs)
+    scale = w.abs().amax(dim=-1, keepdim=True) / max(qmax, 1)
+    scale = torch.where(scale == 0, torch.ones_like(scale), scale)
+    q = torch.clamp(torch.round(w / scale), -qmax - 1, qmax)
+    deq = (q * scale).reshape(out_features, in_features)
+    return deq.to(weight.dtype)
+
+
+def _deviation(baseline, candidate, metric: str) -> float:
+    if metric == "cosine":
+        return 1.0 - metrics.cosine_similarity(baseline, candidate)
+    if metric == "kl":
+        b = baseline.reshape(-1, baseline.shape[-1])
+        c = candidate.reshape(-1, candidate.shape[-1])
+        return metrics.kl_divergence(b, c)
+    return metrics.mse(baseline, candidate)

zeroquantz/quantization/__init__.py

@@ -0,0 +1,17 @@
+"""Quantization backends, configs, and the backend registry."""
+
+from __future__ import annotations
+
+from zeroquantz.quantization.base import MemoryEstimate, QuantizationBackend, QuantizationResult
+from zeroquantz.quantization.config import QuantizationConfig, QuantMethod
+from zeroquantz.quantization.registry import BackendRegistry, default_registry
+
+__all__ = [
+    "BackendRegistry",
+    "MemoryEstimate",
+    "QuantMethod",
+    "QuantizationBackend",
+    "QuantizationConfig",
+    "QuantizationResult",
+    "default_registry",
+]

zeroquantz/quantization/backends/__init__.py

@@ -0,0 +1,8 @@
+"""Built-in quantization backends."""
+
+from __future__ import annotations
+
+from zeroquantz.quantization.backends.bitsandbytes import BitsAndBytesBackend
+from zeroquantz.quantization.backends.torchao import TorchAOBackend
+
+__all__ = ["BitsAndBytesBackend", "TorchAOBackend"]

zeroquantz/quantization/backends/bitsandbytes.py

@@ -0,0 +1,210 @@
+"""bitsandbytes backend: LLM.int8() weight/activation INT8 and 4-bit NF4.
+
+bitsandbytes quantizes *at load time* via ``transformers``' ``BitsAndBytesConfig``,
+so :meth:`quantize` takes a model id / local path rather than a live module.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import time
+from typing import TYPE_CHECKING, Any
+
+from zeroquantz.core.exceptions import DependencyError, QuantizationError, ZeroQuantzError
+from zeroquantz.quantization import memory
+from zeroquantz.quantization.base import (
+    MemoryEstimate,
+    QuantizationBackend,
+    QuantizationResult,
+)
+from zeroquantz.quantization.config import QuantizationConfig, QuantMethod
+from zeroquantz.utils import units
+from zeroquantz.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from zeroquantz.hardware.capabilities import HardwareProfile
+    from zeroquantz.models.metadata import ModelProfile
+    from zeroquantz.optimization.constraints import OptimizationGoal
+
+log = get_logger(__name__)
+
+
+class BitsAndBytesBackend(QuantizationBackend):
+    name = "bitsandbytes"
+    description = "LLM.int8() 8-bit and 4-bit NF4 quantization (quantize-on-load)."
+    methods = (QuantMethod.BNB_NF4.value, QuantMethod.BNB_INT8.value)
+
+    def is_available(self) -> bool:
+        return all(
+            importlib.util.find_spec(m) is not None
+            for m in ("bitsandbytes", "torch", "transformers")
+        )
+
+    def availability_error(self) -> ZeroQuantzError | None:
+        if self.is_available():
+            return None
+        missing = next(
+            m
+            for m in ("bitsandbytes", "torch", "transformers")
+            if importlib.util.find_spec(m) is None
+        )
+        extra = "quant" if missing == "bitsandbytes" else "torch"
+        return DependencyError.for_extra(
+            missing, extra, purpose="quantize with bitsandbytes"
+        )
+
+    def supports(
+        self,
+        model: ModelProfile,
+        hardware: HardwareProfile,
+        runtime: str | None = None,
+        config: QuantizationConfig | None = None,
+    ) -> bool:
+        if not hardware.cuda_available:
+            return False  # bitsandbytes is CUDA-only in practice
+        method = config.method if config else None
+        if method == QuantMethod.BNB_INT8.value:
+            return hardware.supports("int8")
+        if method == QuantMethod.BNB_NF4.value:
+            return hardware.supports("int4")
+        # No specific method requested: available if the GPU can do either.
+        return hardware.supports("int8") or hardware.supports("int4")
+
+    def default_config(
+        self,
+        method: str,
+        model: ModelProfile,
+        hardware: HardwareProfile,
+        goal: OptimizationGoal | None = None,
+    ) -> QuantizationConfig:
+        compute_dtype = "bfloat16" if hardware.supports_bf16 else "float16"
+        if method == QuantMethod.BNB_NF4.value:
+            return QuantizationConfig(
+                backend=self.name,
+                method=method,
+                weight_bits=4,
+                activation_bits=16,
+                group_size=64,
+                compute_dtype=compute_dtype,
+                double_quant=True,
+                quant_type="nf4",
+                skip_modules=["lm_head"],
+            )
+        if method == QuantMethod.BNB_INT8.value:
+            return QuantizationConfig(
+                backend=self.name,
+                method=method,
+                weight_bits=8,
+                activation_bits=8,
+                group_size=None,
+                compute_dtype=compute_dtype,
+                skip_modules=["lm_head"],
+                extra={"llm_int8_threshold": 6.0},
+            )
+        raise QuantizationError(f"bitsandbytes has no method '{method}'.")
+
+    def estimate_memory(
+        self, model: ModelProfile, config: QuantizationConfig
+    ) -> MemoryEstimate:
+        return MemoryEstimate(
+            weight_gb=memory.estimate_weight_gb(model, config),
+            vram_gb=memory.estimate_vram_gb(model, config),
+        )
+
+    def quantize(
+        self,
+        model: str | Any,
+        config: QuantizationConfig,
+        *,
+        progress: Callable[[str, float], None] | None = None,
+    ) -> QuantizationResult:
+        err = self.availability_error()
+        if err is not None:
+            raise err
+        if not isinstance(model, str):
+            raise QuantizationError(
+                "bitsandbytes quantizes at load time.",
+                detail="Pass a model id or local path, not an already-loaded module.",
+                suggestions=["zeroquantz quantize Qwen/Qwen3-8B --backend bitsandbytes --method nf4"],
+            )
+
+        import torch
+        from transformers import AutoModelForCausalLM, BitsAndBytesConfig
+
+        if progress:
+            progress("Building bitsandbytes config", 0.05)
+        compute_dtype = getattr(torch, config.compute_dtype, torch.float16)
+        if config.method == QuantMethod.BNB_NF4.value:
+            bnb_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type=config.quant_type or "nf4",
+                bnb_4bit_use_double_quant=config.double_quant,
+                bnb_4bit_compute_dtype=compute_dtype,
+            )
+        elif config.method == QuantMethod.BNB_INT8.value:
+            bnb_config = BitsAndBytesConfig(
+                load_in_8bit=True,
+                llm_int8_threshold=float(config.extra.get("llm_int8_threshold", 6.0)),
+                llm_int8_skip_modules=config.skip_modules or None,
+            )
+        else:
+            raise QuantizationError(f"bitsandbytes has no method '{config.method}'.")
+
+        if progress:
+            progress("Loading and quantizing weights", 0.1)
+        started = time.perf_counter()
+        try:
+            module = AutoModelForCausalLM.from_pretrained(
+                model,
+                quantization_config=bnb_config,
+                device_map="auto",
+                low_cpu_mem_usage=True,
+            )
+        except Exception as exc:  # pragma: no cover - requires GPU + weights
+            raise QuantizationError(
+                f"bitsandbytes quantization of '{model}' failed.",
+                detail=str(exc),
+                suggestions=["zeroquantz recommend " + str(model), "check available VRAM"],
+            ) from exc
+        elapsed = time.perf_counter() - started
+        if progress:
+            progress("Quantization complete", 1.0)
+
+        return QuantizationResult(
+            backend=self.name,
+            method=config.method,
+            config=config,
+            model=module,
+            elapsed_s=elapsed,
+            quantized_size_gb=_module_size_gb(module),
+        )
+
+    def export(self, result: QuantizationResult, output_dir: str) -> str:
+        if result.model is None:
+            raise QuantizationError("Nothing to export: no quantized model in result.")
+        result.model.save_pretrained(output_dir)
+        result.output_dir = output_dir
+        return output_dir
+
+
+def _module_size_gb(module: Any) -> float | None:
+    try:
+        import torch
+
+        total = 0
+        for p in module.parameters():
+            total += p.numel() * _element_size(p)
+        for b in module.buffers():
+            total += b.numel() * _element_size(b)
+        return round(units.bytes_to_gb(total), 2)
+    except Exception:  # pragma: no cover - torch dependent
+        return None
+
+
+def _element_size(tensor: Any) -> float:
+    try:
+        return tensor.element_size()
+    except Exception:
+        return 4.0

zeroquantz/quantization/backends/torchao.py

@@ -0,0 +1,198 @@
+"""TorchAO backend: INT8 and INT4 weight-only quantization.
+
+TorchAO quantizes a *live* module in place via ``torchao.quantization.quantize_``,
+so :meth:`quantize` accepts either a model id/path (which it loads) or an
+already-loaded module.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import time
+from typing import TYPE_CHECKING, Any
+
+from zeroquantz.core.exceptions import DependencyError, QuantizationError, ZeroQuantzError
+from zeroquantz.quantization import memory
+from zeroquantz.quantization.base import (
+    MemoryEstimate,
+    QuantizationBackend,
+    QuantizationResult,
+)
+from zeroquantz.quantization.config import QuantizationConfig, QuantMethod
+from zeroquantz.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from zeroquantz.hardware.capabilities import HardwareProfile
+    from zeroquantz.models.metadata import ModelProfile
+    from zeroquantz.optimization.constraints import OptimizationGoal
+
+log = get_logger(__name__)
+
+
+class TorchAOBackend(QuantizationBackend):
+    name = "torchao"
+    description = "PyTorch-native INT8/INT4 weight-only quantization (torchao)."
+    methods = (QuantMethod.TORCHAO_INT8.value, QuantMethod.TORCHAO_INT4.value)
+
+    def is_available(self) -> bool:
+        return all(
+            importlib.util.find_spec(m) is not None for m in ("torchao", "torch")
+        )
+
+    def availability_error(self) -> ZeroQuantzError | None:
+        if self.is_available():
+            return None
+        missing = "torchao" if importlib.util.find_spec("torchao") is None else "torch"
+        extra = "quant" if missing == "torchao" else "torch"
+        return DependencyError.for_extra(missing, extra, purpose="quantize with TorchAO")
+
+    def supports(
+        self,
+        model: ModelProfile,
+        hardware: HardwareProfile,
+        runtime: str | None = None,
+        config: QuantizationConfig | None = None,
+    ) -> bool:
+        if not hardware.cuda_available:
+            return False
+        method = config.method if config else None
+        if method == QuantMethod.TORCHAO_INT4.value:
+            return hardware.supports("int4")
+        if method == QuantMethod.TORCHAO_INT8.value:
+            return hardware.supports("int8")
+        return hardware.supports("int8")
+
+    def default_config(
+        self,
+        method: str,
+        model: ModelProfile,
+        hardware: HardwareProfile,
+        goal: OptimizationGoal | None = None,
+    ) -> QuantizationConfig:
+        compute_dtype = "bfloat16" if hardware.supports_bf16 else "float16"
+        if method == QuantMethod.TORCHAO_INT4.value:
+            return QuantizationConfig(
+                backend=self.name,
+                method=method,
+                weight_bits=4,
+                activation_bits=16,
+                group_size=128,
+                compute_dtype=compute_dtype,
+                skip_modules=["lm_head"],
+            )
+        if method == QuantMethod.TORCHAO_INT8.value:
+            return QuantizationConfig(
+                backend=self.name,
+                method=method,
+                weight_bits=8,
+                activation_bits=16,
+                group_size=None,
+                compute_dtype=compute_dtype,
+                skip_modules=["lm_head"],
+            )
+        raise QuantizationError(f"TorchAO has no method '{method}'.")
+
+    def estimate_memory(
+        self, model: ModelProfile, config: QuantizationConfig
+    ) -> MemoryEstimate:
+        return MemoryEstimate(
+            weight_gb=memory.estimate_weight_gb(model, config),
+            vram_gb=memory.estimate_vram_gb(model, config),
+        )
+
+    def quantize(
+        self,
+        model: str | Any,
+        config: QuantizationConfig,
+        *,
+        progress: Callable[[str, float], None] | None = None,
+    ) -> QuantizationResult:
+        err = self.availability_error()
+        if err is not None:
+            raise err
+
+        import torch
+
+        if progress:
+            progress("Loading model", 0.1)
+        module = model
+        if isinstance(model, str):
+            from transformers import AutoModelForCausalLM
+
+            compute_dtype = getattr(torch, config.compute_dtype, torch.bfloat16)
+            module = AutoModelForCausalLM.from_pretrained(
+                model, torch_dtype=compute_dtype, device_map="cuda"
+            )
+
+        if progress:
+            progress("Applying TorchAO quantization", 0.4)
+        started = time.perf_counter()
+        try:
+            self._apply_quantization(module, config)
+        except QuantizationError:
+            raise
+        except Exception as exc:  # pragma: no cover - requires GPU
+            raise QuantizationError(
+                "TorchAO quantization failed.",
+                detail=str(exc),
+                suggestions=["update torchao", "zeroquantz recommend " + str(model)],
+            ) from exc
+        elapsed = time.perf_counter() - started
+        if progress:
+            progress("Quantization complete", 1.0)
+
+        return QuantizationResult(
+            backend=self.name,
+            method=config.method,
+            config=config,
+            model=module,
+            elapsed_s=elapsed,
+        )
+
+    def _apply_quantization(self, module: Any, config: QuantizationConfig) -> None:
+        from torchao.quantization import quantize_
+
+        # Prefer the modern config-object API; fall back to the factory functions
+        # used by older torchao releases.
+        if config.method == QuantMethod.TORCHAO_INT4.value:
+            try:
+                from torchao.quantization import Int4WeightOnlyConfig
+
+                quantize_(module, Int4WeightOnlyConfig(group_size=config.group_size or 128))
+                return
+            except ImportError:
+                from torchao.quantization import int4_weight_only
+
+                quantize_(module, int4_weight_only(group_size=config.group_size or 128))
+                return
+        if config.method == QuantMethod.TORCHAO_INT8.value:
+            try:
+                from torchao.quantization import Int8WeightOnlyConfig
+
+                quantize_(module, Int8WeightOnlyConfig())
+                return
+            except ImportError:
+                from torchao.quantization import int8_weight_only
+
+                quantize_(module, int8_weight_only())
+                return
+        raise QuantizationError(f"TorchAO has no method '{config.method}'.")
+
+    def export(self, result: QuantizationResult, output_dir: str) -> str:
+        if result.model is None:
+            raise QuantizationError("Nothing to export: no quantized model in result.")
+        # TorchAO models save through the standard HF API when loaded via
+        # transformers; otherwise fall back to a state-dict checkpoint.
+        try:
+            result.model.save_pretrained(output_dir, safe_serialization=False)
+        except AttributeError:  # pragma: no cover - non-HF module
+            import torch
+
+            from pathlib import Path
+
+            Path(output_dir).mkdir(parents=True, exist_ok=True)
+            torch.save(result.model.state_dict(), Path(output_dir) / "model.pt")
+        result.output_dir = output_dir
+        return output_dir