zeroquantz 0.1.0__py3-none-any.whl

zeroquantz/utils/env.py

@@ -0,0 +1,78 @@
+"""Capture the software/hardware environment for reproducible reports."""
+
+from __future__ import annotations
+
+import importlib.metadata
+import importlib.util
+import os
+import platform
+from datetime import UTC, datetime
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from zeroquantz.hardware.capabilities import HardwareProfile
+
+
+def _pkg_version(name: str) -> str | None:
+    try:
+        return importlib.metadata.version(name)
+    except importlib.metadata.PackageNotFoundError:
+        return None
+
+
+def enable_fast_downloads() -> bool:
+    """Turn on the Rust ``hf_transfer`` accelerator for Hub downloads, if present.
+
+    ``hf_transfer`` is HuggingFace's Rust, multi-threaded, chunked downloader; it
+    can multiply throughput on fast links where a single HTTPS connection is the
+    bottleneck (on a bandwidth-capped connection it makes little difference).
+
+    Crucially, ``huggingface_hub`` *raises* at download time if
+    ``HF_HUB_ENABLE_HF_TRANSFER=1`` is set but the package is missing, so we only
+    set the flag when ``hf_transfer`` is actually importable — and use
+    ``setdefault`` so an explicit ``HF_HUB_ENABLE_HF_TRANSFER=0`` from the user's
+    environment is always respected. Returns whether acceleration is now active.
+    """
+    if importlib.util.find_spec("hf_transfer") is None:
+        return False
+    os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+    return os.environ.get("HF_HUB_ENABLE_HF_TRANSFER") == "1"
+
+
+def fast_downloads_active() -> bool:
+    """Whether ``hf_transfer`` acceleration is currently enabled in this process."""
+    return (
+        os.environ.get("HF_HUB_ENABLE_HF_TRANSFER") == "1"
+        and importlib.util.find_spec("hf_transfer") is not None
+    )
+
+
+def utc_timestamp() -> str:
+    return datetime.now(UTC).isoformat(timespec="seconds")
+
+
+def capture_environment(hardware: HardwareProfile | None = None) -> dict[str, Any]:
+    """A JSON-serializable snapshot of versions, platform, and (optional) GPU."""
+    from zeroquantz import __version__
+
+    env: dict[str, Any] = {
+        "zeroquantz_version": __version__,
+        "timestamp": utc_timestamp(),
+        "python": platform.python_version(),
+        "platform": platform.platform(),
+        "torch": _pkg_version("torch"),
+        "transformers": _pkg_version("transformers"),
+        "bitsandbytes": _pkg_version("bitsandbytes"),
+        "torchao": _pkg_version("torchao"),
+    }
+    if hardware is not None:
+        env["gpu"] = hardware.gpu_name
+        env["vram_gb"] = hardware.total_vram_gb
+        env["compute_capability"] = (
+            ".".join(map(str, hardware.compute_capability))
+            if hardware.compute_capability
+            else None
+        )
+        env["cuda_version"] = hardware.cuda_version
+        env["driver_version"] = hardware.driver_version
+    return env

zeroquantz/utils/logging.py

@@ -0,0 +1,73 @@
+"""Structured logging setup.
+
+The interactive UI shows concise messages; detailed logs are written to
+``~/.zeroquantz/logs/zeroquantz.log``. ``--verbose`` raises the console level to
+INFO and ``--debug`` to DEBUG.
+"""
+
+from __future__ import annotations
+
+import logging
+from logging.handlers import RotatingFileHandler
+from typing import Literal
+
+from zeroquantz.utils.paths import paths
+
+_CONFIGURED = False
+
+Verbosity = Literal["quiet", "verbose", "debug"]
+
+
+def get_logger(name: str) -> logging.Logger:
+    """Return a namespaced child of the ``zeroquantz`` logger."""
+    if not name.startswith("zeroquantz"):
+        name = f"zeroquantz.{name}"
+    return logging.getLogger(name)
+
+
+def configure_logging(verbosity: Verbosity = "quiet", *, log_to_file: bool = True) -> None:
+    """Configure the root ``zeroquantz`` logger. Safe to call more than once."""
+    global _CONFIGURED
+    root = logging.getLogger("zeroquantz")
+    root.setLevel(logging.DEBUG)
+    root.handlers.clear()
+
+    console_level = {
+        "quiet": logging.WARNING,
+        "verbose": logging.INFO,
+        "debug": logging.DEBUG,
+    }[verbosity]
+
+    console = logging.StreamHandler()
+    console.setLevel(console_level)
+    console.setFormatter(logging.Formatter("%(levelname)s %(name)s: %(message)s"))
+    root.addHandler(console)
+
+    if log_to_file:
+        try:
+            log_dir = paths().ensure().logs
+            file_handler = RotatingFileHandler(
+                log_dir / "zeroquantz.log",
+                maxBytes=2 * 1024 * 1024,
+                backupCount=3,
+                encoding="utf-8",
+            )
+            file_handler.setLevel(logging.DEBUG)
+            file_handler.setFormatter(
+                logging.Formatter(
+                    "%(asctime)s %(levelname)-7s %(name)s %(message)s",
+                    datefmt="%Y-%m-%dT%H:%M:%S",
+                )
+            )
+            root.addHandler(file_handler)
+        except OSError:
+            # A read-only home directory should never take the app down.
+            pass
+
+    root.propagate = False
+    _CONFIGURED = True
+
+
+def ensure_configured() -> None:
+    if not _CONFIGURED:
+        configure_logging("quiet")

zeroquantz/utils/metrics.py

@@ -0,0 +1,98 @@
+"""Pure numeric comparison metrics (NumPy only).
+
+Shared by sensitivity profiling and verification so the definitions live in one
+place. Everything accepts array-likes (Python lists, NumPy arrays, or detached
+tensors converted to NumPy) and returns plain floats.
+
+NumPy is imported lazily inside each function rather than at module import time:
+these metrics are only used on the profiling/verification paths, so keeping the
+import local shaves NumPy (~100 ms) off ``zeroquantz`` cold-start, which imports
+this module transitively via the dispatcher. The import is cached after first
+use, so repeated calls pay only a dict lookup.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.typing import ArrayLike
+
+
+def _arr(x: ArrayLike) -> np.ndarray:
+    import numpy as np
+
+    return np.asarray(x, dtype=np.float64)
+
+
+def mse(a: ArrayLike, b: ArrayLike) -> float:
+    """Mean squared error between two equally-shaped arrays."""
+    import numpy as np
+
+    av, bv = _arr(a), _arr(b)
+    return float(np.mean((av - bv) ** 2))
+
+
+def cosine_similarity(a: ArrayLike, b: ArrayLike) -> float:
+    """Cosine similarity of two arrays, flattened to vectors."""
+    import numpy as np
+
+    av, bv = _arr(a).ravel(), _arr(b).ravel()
+    na = float(np.linalg.norm(av))
+    nb = float(np.linalg.norm(bv))
+    if na == 0.0 or nb == 0.0:
+        return 0.0
+    return float(np.dot(av, bv) / (na * nb))
+
+
+def softmax(logits: ArrayLike, axis: int = -1) -> np.ndarray:
+    import numpy as np
+
+    x = _arr(logits)
+    x = x - np.max(x, axis=axis, keepdims=True)
+    e = np.exp(x)
+    return e / np.sum(e, axis=axis, keepdims=True)
+
+
+def kl_divergence(p_logits: ArrayLike, q_logits: ArrayLike, *, eps: float = 1e-12) -> float:
+    """Mean KL(P || Q) where P, Q are softmaxes of the given logits.
+
+    Works for 1D (single distribution) or 2D (a batch of rows, averaged) input.
+    """
+    import numpy as np
+
+    p = softmax(p_logits)
+    q = softmax(q_logits)
+    p = np.clip(p, eps, 1.0)
+    q = np.clip(q, eps, 1.0)
+    kl = np.sum(p * (np.log(p) - np.log(q)), axis=-1)
+    return float(np.mean(kl))
+
+
+def top1_agreement(p_logits: ArrayLike, q_logits: ArrayLike) -> float:
+    """Fraction of rows whose argmax agrees between the two logit arrays."""
+    import numpy as np
+
+    p = _arr(p_logits)
+    q = _arr(q_logits)
+    if p.ndim == 1:
+        return float(np.argmax(p) == np.argmax(q))
+    return float(np.mean(np.argmax(p, axis=-1) == np.argmax(q, axis=-1)))
+
+
+def perplexity(nlls: ArrayLike) -> float:
+    """Perplexity from a sequence of per-token negative log-likelihoods (nats)."""
+    import numpy as np
+
+    v = _arr(nlls)
+    if v.size == 0:
+        return float("nan")
+    return float(np.exp(np.mean(v)))
+
+
+def relative_delta(baseline: float, candidate: float) -> float:
+    """Signed relative change ``(candidate - baseline) / baseline``."""
+    if baseline == 0:
+        return float("nan")
+    return (candidate - baseline) / baseline

zeroquantz/utils/paths.py

@@ -0,0 +1,57 @@
+"""Filesystem locations ZeroQuantz reads and writes.
+
+The layout follows the spec's ``~/.zeroquantz`` convention, but the root is
+overridable with the ``ZEROQUANTZ_HOME`` environment variable (useful for tests
+and CI, which point it at a temp directory).
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+
+def _default_home() -> Path:
+    override = os.environ.get("ZEROQUANTZ_HOME")
+    if override:
+        return Path(override).expanduser()
+    return Path.home() / ".zeroquantz"
+
+
+@dataclass(frozen=True)
+class ZeroQuantzPaths:
+    """Resolved application directories. Created lazily via :meth:`ensure`."""
+
+    home: Path
+
+    @classmethod
+    def resolve(cls) -> ZeroQuantzPaths:
+        return cls(home=_default_home())
+
+    @property
+    def sessions(self) -> Path:
+        return self.home / "sessions"
+
+    @property
+    def logs(self) -> Path:
+        return self.home / "logs"
+
+    @property
+    def cache(self) -> Path:
+        return self.home / "cache"
+
+    @property
+    def config_file(self) -> Path:
+        return self.home / "config.toml"
+
+    def ensure(self) -> ZeroQuantzPaths:
+        """Create the standard directories if they do not already exist."""
+        for directory in (self.home, self.sessions, self.logs, self.cache):
+            directory.mkdir(parents=True, exist_ok=True)
+        return self
+
+
+def paths() -> ZeroQuantzPaths:
+    """Return freshly-resolved paths (re-reads ``ZEROQUANTZ_HOME`` each call)."""
+    return ZeroQuantzPaths.resolve()

zeroquantz/utils/units.py

@@ -0,0 +1,134 @@
+"""Byte/parameter/precision math and the one place ZeroQuantz defines its units.
+
+**Unit convention.** ZeroQuantz reports memory in *binary gigabytes* (GiB,
+``2**30`` bytes) but labels them ``GB`` in the UI, because that is how GPU VRAM
+and user memory budgets are universally stated ("a 24 GB 4090", "fit it under
+8 GB"). Keeping VRAM, model size, and user budgets in one consistent binary unit
+is what lets ``estimated_vram_gb <= max_vram_gb`` mean something. Every
+conversion in the codebase goes through :func:`bytes_to_gb` / :func:`gb_to_bytes`
+so the convention lives in exactly one file.
+"""
+
+from __future__ import annotations
+
+from typing import Final
+
+BYTES_PER_GIB: Final[int] = 1024**3
+BYTES_PER_MIB: Final[int] = 1024**2
+
+# Effective *bits per stored weight* for each precision. Sub-byte quantization
+# schemes carry group-wise metadata (scales, and often zero-points), so their
+# effective footprint is higher than the nominal bit-width. These defaults
+# assume a group size of 128 with fp16 scales, which is the common case; the
+# candidate generator can override them per-config.
+BITS_PER_PARAM: Final[dict[str, float]] = {
+    "float32": 32.0,
+    "fp32": 32.0,
+    "float16": 16.0,
+    "fp16": 16.0,
+    "bfloat16": 16.0,
+    "bf16": 16.0,
+    "int8": 8.125,  # 8 + fp16 scale per group of 128
+    "fp8": 8.0,
+    "fp8_e4m3": 8.0,
+    "fp8_e5m2": 8.0,
+    "int4": 4.25,  # 4 + fp16 scale + int4 zero-point over a group of 128
+    "nf4": 4.127,  # NF4 with double quantization (bitsandbytes)
+    "fp4": 4.25,
+    "int3": 3.25,
+    "int2": 2.25,
+}
+
+# Canonical display/short names for the precisions we reason about.
+DTYPE_ALIASES: Final[dict[str, str]] = {
+    "torch.float32": "float32",
+    "torch.float16": "float16",
+    "torch.bfloat16": "bfloat16",
+    "float32": "float32",
+    "float16": "float16",
+    "bfloat16": "bfloat16",
+    "f32": "float32",
+    "f16": "float16",
+    "bf16": "bfloat16",
+    "fp32": "float32",
+    "fp16": "float16",
+    "half": "float16",
+    "int8": "int8",
+    "i8": "int8",
+    "uint8": "int8",
+    "int4": "int4",
+    "i4": "int4",
+    "nf4": "nf4",
+    "fp4": "fp4",
+    "fp8": "fp8",
+}
+
+
+def normalize_dtype(dtype: str) -> str:
+    """Return a canonical dtype name (e.g. ``"bf16"`` -> ``"bfloat16"``)."""
+    key = str(dtype).strip().lower()
+    return DTYPE_ALIASES.get(key, key)
+
+
+def bits_per_param(precision: str) -> float:
+    """Effective stored bits per parameter for ``precision``.
+
+    Raises:
+        KeyError: if the precision is unknown. Callers that want a soft failure
+            should catch this and treat the config as unsupported.
+    """
+    key = normalize_dtype(precision)
+    if key in BITS_PER_PARAM:
+        return BITS_PER_PARAM[key]
+    # Also accept the already-canonical short forms used as dict keys above.
+    if precision.lower() in BITS_PER_PARAM:
+        return BITS_PER_PARAM[precision.lower()]
+    raise KeyError(f"unknown precision: {precision!r}")
+
+
+def bytes_per_param(precision: str) -> float:
+    """Effective stored bytes per parameter for ``precision``."""
+    return bits_per_param(precision) / 8.0
+
+
+def bytes_to_gb(num_bytes: float) -> float:
+    """Bytes -> binary gigabytes (GiB), the ZeroQuantz canonical unit."""
+    return num_bytes / BYTES_PER_GIB
+
+
+def gb_to_bytes(gb: float) -> float:
+    """Binary gigabytes (GiB) -> bytes."""
+    return gb * BYTES_PER_GIB
+
+
+def param_bytes(num_params: float, precision: str) -> float:
+    """Total bytes to store ``num_params`` parameters at ``precision``."""
+    return num_params * bytes_per_param(precision)
+
+
+def param_gb(num_params: float, precision: str) -> float:
+    """Total GiB to store ``num_params`` parameters at ``precision``."""
+    return bytes_to_gb(param_bytes(num_params, precision))
+
+
+def humanize_params(num_params: float) -> str:
+    """Render a parameter count as e.g. ``"8.2B"`` / ``"350M"`` / ``"1.3K"``."""
+    n = float(num_params)
+    for divisor, suffix in ((1e12, "T"), (1e9, "B"), (1e6, "M"), (1e3, "K")):
+        if abs(n) >= divisor:
+            return f"{n / divisor:.1f}{suffix}"
+    return str(int(n))
+
+
+def humanize_bytes(num_bytes: float) -> str:
+    """Render a byte count with a binary suffix (``"5.3 GB"``, ``"812 MB"``)."""
+    n = float(num_bytes)
+    for divisor, suffix in (
+        (1024**4, "TB"),
+        (1024**3, "GB"),
+        (1024**2, "MB"),
+        (1024, "KB"),
+    ):
+        if abs(n) >= divisor:
+            return f"{n / divisor:.2f} {suffix}"
+    return f"{int(n)} B"

zeroquantz/verification/__init__.py

@@ -0,0 +1,17 @@
+"""Verification: compare a quantized model against its full-precision baseline."""
+
+from __future__ import annotations
+
+from zeroquantz.verification.report import (
+    VerificationReport,
+    VerificationThresholds,
+    Verifier,
+    evaluate_thresholds,
+)
+
+__all__ = [
+    "VerificationReport",
+    "VerificationThresholds",
+    "Verifier",
+    "evaluate_thresholds",
+]

zeroquantz/verification/logits.py

@@ -0,0 +1,55 @@
+"""Compare logits between a baseline and a quantized model (lazy torch)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from zeroquantz.utils import metrics
+
+
+@dataclass
+class LogitComparison:
+    samples: int
+    logit_cosine: float
+    mean_kl: float
+    top1_agreement: float
+
+
+def compare_logits(
+    base_model: Any,
+    quant_model: Any,
+    tokenizer: Any,
+    prompts: list[str],
+    *,
+    max_length: int = 128,
+) -> LogitComparison:
+    """Run both models on ``prompts`` and compare their output logits."""
+    import numpy as np
+    import torch
+
+    base_logits: list[np.ndarray] = []
+    quant_logits: list[np.ndarray] = []
+
+    base_device = next(base_model.parameters()).device
+    quant_device = next(quant_model.parameters()).device
+
+    with torch.no_grad():
+        for prompt in prompts:
+            enc = tokenizer(
+                prompt, return_tensors="pt", truncation=True, max_length=max_length
+            )
+            b = base_model(**{k: v.to(base_device) for k, v in enc.items()})
+            q = quant_model(**{k: v.to(quant_device) for k, v in enc.items()})
+            # last-token logits
+            base_logits.append(b.logits[0, -1].float().cpu().numpy())
+            quant_logits.append(q.logits[0, -1].float().cpu().numpy())
+
+    base_arr = np.stack(base_logits)
+    quant_arr = np.stack(quant_logits)
+    return LogitComparison(
+        samples=len(prompts),
+        logit_cosine=round(metrics.cosine_similarity(base_arr, quant_arr), 6),
+        mean_kl=round(metrics.kl_divergence(base_arr, quant_arr), 6),
+        top1_agreement=round(metrics.top1_agreement(base_arr, quant_arr), 6),
+    )