zeroquantz 0.1.0__py3-none-any.whl

zeroquantz/optimization/candidate.py

@@ -0,0 +1,272 @@
+"""Quantization candidate generation, driven by the format catalog + fit math.
+
+Each :class:`QuantizationCandidate` pairs a catalog format with *estimated* size /
+VRAM / max-context / quality-risk / speed numbers (from :mod:`zeroquantz.optimization.fit`
+and the format's quality/speed tiers) plus runtime compatibility and how ZeroQuantz
+would execute it (in the main env, an isolated sub-env, or produce-only).
+
+Estimates are heuristics, always labelled as such. Whether a format is runnable
+*here* (hardware, execution path) is tracked separately so ZeroQuantz can show the
+full landscape while only *recommending* something it can actually run.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field
+
+from zeroquantz.optimization import fit as fitmod
+from zeroquantz.optimization.constraints import OptimizationGoal
+from zeroquantz.quantization import catalog
+from zeroquantz.quantization.catalog import Execution, QuantFormat
+from zeroquantz.quantization.config import QuantizationConfig
+from zeroquantz.runtimes.compatibility import RuntimeCompat
+
+if TYPE_CHECKING:
+    from zeroquantz.hardware.capabilities import HardwareProfile
+    from zeroquantz.models.metadata import ModelProfile
+    from zeroquantz.quantization.registry import BackendRegistry
+
+# quality_tier (1..5) -> base quality-degradation risk (0..1), before size scaling
+_TIER_RISK = {5: 0.05, 4: 0.15, 3: 0.30, 2: 0.50, 1: 0.75}
+# speed_tier (1..5) -> estimated decode throughput vs bf16 baseline
+_TIER_SPEED = {5: 1.8, 4: 1.5, 3: 1.2, 2: 1.0, 1: 0.85}
+
+
+def _size_factor(parameter_count: int) -> float:
+    p = parameter_count
+    if p < 1_000_000_000:
+        return 1.6
+    if p < 3_000_000_000:
+        return 1.3
+    if p < 7_000_000_000:
+        return 1.15
+    if p < 13_000_000_000:
+        return 1.0
+    if p < 30_000_000_000:
+        return 0.92
+    if p < 70_000_000_000:
+        return 0.85
+    return 0.8
+
+
+def estimate_quality_risk(model: ModelProfile, fmt: QuantFormat) -> float:
+    risk = _TIER_RISK.get(fmt.quality_tier, 0.3) * _size_factor(model.parameter_count)
+    return round(min(max(risk, 0.02), 0.95), 3)
+
+
+def _risk_label(risk: float) -> str:
+    if risk < 0.10:
+        return "very low"
+    if risk < 0.25:
+        return "low"
+    if risk < 0.45:
+        return "moderate"
+    if risk < 0.70:
+        return "high"
+    return "very high"
+
+
+def _runtime_compat(fmt: QuantFormat, runtime: str | None) -> RuntimeCompat:
+    if runtime is None:
+        return RuntimeCompat.COMPATIBLE
+    return RuntimeCompat.COMPATIBLE if runtime.strip().lower() in fmt.runtimes else RuntimeCompat.UNSUPPORTED
+
+
+class QuantizationCandidate(BaseModel):
+    """A candidate format with estimated costs, runtime fit, and execution info."""
+
+    format_id: str
+    label: str
+    family: str
+    execution: str
+    produced_by: str
+    bits_per_weight: float
+    calibration: str
+    runtimes: list[str] = Field(default_factory=list)
+
+    config: QuantizationConfig
+    estimated_size_gb: float
+    estimated_vram_gb: float
+    estimated_quality_risk: float
+    estimated_speedup: float
+    max_context: int | None = None
+
+    runtime_compat: RuntimeCompat = RuntimeCompat.COMPATIBLE
+    hardware_status: str = "supported"       # supported | limited | unsupported
+    backend_available: bool = True
+    fits_memory: bool = True
+    meets_quality: bool = True
+    notes: list[str] = Field(default_factory=list)
+
+    # ---- conveniences -------------------------------------------------------
+
+    @property
+    def method(self) -> str:
+        return self.config.method
+
+    @property
+    def backend(self) -> str:
+        return self.config.backend
+
+    @property
+    def weight_bits(self) -> int:
+        return self.config.weight_bits
+
+    @property
+    def activation_bits(self) -> int:
+        return self.config.activation_bits
+
+    @property
+    def quality_risk_label(self) -> str:
+        return _risk_label(self.estimated_quality_risk)
+
+    @property
+    def executable(self) -> bool:
+        return self.execution in (Execution.IN_ENV.value, Execution.ISOLATED.value)
+
+    @property
+    def implemented(self) -> bool:  # kept for API compatibility with the recommender
+        return self.executable
+
+    @property
+    def feasible(self) -> bool:
+        """Runnable here and satisfies the goal's hard constraints."""
+        return (
+            self.executable
+            and self.hardware_status != "unsupported"
+            and self.fits_memory
+            and self.meets_quality
+            and self.runtime_compat is not RuntimeCompat.UNSUPPORTED
+        )
+
+
+class CandidateGenerator:
+    """Enumerate quantization candidates for a model/hardware/goal from the catalog."""
+
+    @staticmethod
+    def generate(
+        model: ModelProfile,
+        hardware: HardwareProfile,
+        goal: OptimizationGoal | None = None,
+        *,
+        registry: BackendRegistry | None = None,
+        formats: tuple[QuantFormat, ...] | None = None,
+    ) -> list[QuantizationCandidate]:
+        goal = goal or OptimizationGoal()
+        if registry is None:
+            from zeroquantz.quantization.registry import default_registry
+
+            registry = default_registry()
+        formats = formats or catalog.ALL_FORMATS
+
+        candidates: list[QuantizationCandidate] = []
+        for fmt in formats:
+            report = fitmod.fit(
+                model, fmt, budget_gb=goal.max_vram_gb, context_length=goal.context_length
+            )
+            risk = estimate_quality_risk(model, fmt)
+            compat = _runtime_compat(fmt, goal.runtime)
+            hw_status = CandidateGenerator._hardware_status(fmt, hardware)
+            config = CandidateGenerator._config_for(fmt, model, hardware, goal, registry)
+            backend_available = CandidateGenerator._backend_available(fmt, registry)
+
+            fits_memory = report.fits
+            if goal.max_model_size_gb is not None:
+                fits_memory = fits_memory and report.weight_gb <= goal.max_model_size_gb
+            meets_quality = (
+                goal.max_quality_loss is None
+                or risk <= max(goal.max_quality_loss * 4.0, goal.max_quality_loss)
+            )
+
+            candidates.append(
+                QuantizationCandidate(
+                    format_id=fmt.id,
+                    label=fmt.label,
+                    family=fmt.family,
+                    execution=fmt.execution.value,
+                    produced_by=fmt.produced_by,
+                    bits_per_weight=fmt.bits_per_weight,
+                    calibration=fmt.calibration,
+                    runtimes=list(fmt.runtimes),
+                    config=config,
+                    estimated_size_gb=report.weight_gb,
+                    estimated_vram_gb=report.total_gb,
+                    estimated_quality_risk=risk,
+                    estimated_speedup=_TIER_SPEED.get(fmt.speed_tier, 1.0),
+                    max_context=report.max_context,
+                    runtime_compat=compat,
+                    hardware_status=hw_status,
+                    backend_available=backend_available,
+                    fits_memory=fits_memory,
+                    meets_quality=meets_quality,
+                    notes=CandidateGenerator._notes(fmt, hw_status, compat, backend_available),
+                )
+            )
+        return candidates
+
+    # ---- internals ----------------------------------------------------------
+
+    @staticmethod
+    def _hardware_status(fmt: QuantFormat, hardware: HardwareProfile) -> str:
+        if fmt.requires_precision is None:
+            return "supported"  # CPU-capable (e.g. GGUF)
+        return hardware.status(fmt.requires_precision).label.replace("yes", "supported").replace(
+            "no", "unsupported"
+        )
+
+    @staticmethod
+    def _backend_available(fmt: QuantFormat, registry: BackendRegistry) -> bool:
+        if fmt.execution is Execution.IN_ENV and fmt.backend in registry.names():
+            return registry.get(fmt.backend).is_available()
+        if fmt.execution is Execution.ISOLATED:
+            return True  # ZeroQuantz can build a sub-env for it on demand
+        return False
+
+    @staticmethod
+    def _config_for(
+        fmt: QuantFormat,
+        model: ModelProfile,
+        hardware: HardwareProfile,
+        goal: OptimizationGoal,
+        registry: BackendRegistry,
+    ) -> QuantizationConfig:
+        if fmt.backend and fmt.method and fmt.backend in registry.names():
+            backend = registry.get(fmt.backend)
+            if fmt.method in backend.methods:
+                return backend.default_config(fmt.method, model, hardware, goal)
+        compute_dtype = "bfloat16" if hardware.supports_bf16 else "float16"
+        return QuantizationConfig(
+            backend=fmt.family,
+            method=fmt.method or fmt.id,
+            weight_bits=max(1, min(16, round(fmt.bits_per_weight))),
+            activation_bits=fmt.activation_bits,
+            group_size=128 if fmt.bits_per_weight < 8 else None,
+            compute_dtype=compute_dtype,
+            skip_modules=["lm_head"],
+            extra={"format_id": fmt.id},
+        )
+
+    @staticmethod
+    def _notes(
+        fmt: QuantFormat, hw_status: str, compat: RuntimeCompat, backend_available: bool
+    ) -> list[str]:
+        notes: list[str] = []
+        if fmt.execution is Execution.ISOLATED:
+            notes.append("runs in an isolated sub-env")
+        elif fmt.execution is Execution.EXTERNAL:
+            notes.append(f"produced via {fmt.produced_by}")
+        elif fmt.execution is Execution.PLANNED:
+            notes.append("roadmap")
+        if fmt.execution is Execution.IN_ENV and fmt.backend and not backend_available:
+            notes.append(f"needs {fmt.backend}")
+        if fmt.calibration == "required":
+            notes.append("needs calibration data")
+        if hw_status == "unsupported":
+            notes.append(f"{fmt.requires_precision} not supported on this GPU")
+        elif hw_status == "limited":
+            notes.append(f"{fmt.requires_precision} limited on this GPU")
+        if compat is RuntimeCompat.UNSUPPORTED:
+            notes.append("not supported by the target runtime")
+        return notes

zeroquantz/optimization/constraints.py

@@ -0,0 +1,70 @@
+"""The :class:`OptimizationGoal` — ZeroQuantz's internal representation of *what
+the user is asking for*.
+
+This object is the boundary between intent and planning: natural-language
+requests and slash commands both funnel into mutations of this one type, and the
+recommender/candidate generator read only from it.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class Objective(str, Enum):
+    """What to optimize for when strategies trade off against each other."""
+
+    QUALITY = "quality"
+    SPEED = "speed"
+    MEMORY = "memory"
+    BALANCED = "balanced"
+
+    def __str__(self) -> str:
+        return self.value
+
+
+class OptimizationGoal(BaseModel):
+    """Constraints and preferences that steer quantization strategy selection."""
+
+    max_vram_gb: float | None = Field(default=None, gt=0)
+    max_model_size_gb: float | None = Field(default=None, gt=0)
+    runtime: str | None = None
+    objective: Objective = Objective.BALANCED
+    #: Context length (tokens) to size the KV cache against when checking fit.
+    context_length: int = Field(default=4096, gt=0)
+    #: Maximum acceptable relative quality loss (e.g. 0.01 == 1%).
+    max_quality_loss: float | None = Field(default=None, ge=0)
+    #: Optional named GPU to plan *for* (may differ from the local machine).
+    target_gpu: str | None = None
+
+    def with_updates(self, **changes: Any) -> OptimizationGoal:
+        """Return a new goal with ``changes`` applied (originals are immutable-ish
+        snapshots for undo/history)."""
+        return self.model_copy(update={k: v for k, v in changes.items() if v is not None})
+
+    def is_empty(self) -> bool:
+        return (
+            self.max_vram_gb is None
+            and self.max_model_size_gb is None
+            and self.runtime is None
+            and self.max_quality_loss is None
+            and self.objective is Objective.BALANCED
+        )
+
+    def describe(self) -> list[tuple[str, str]]:
+        rows: list[tuple[str, str]] = [("Objective", str(self.objective))]
+        if self.max_vram_gb is not None:
+            rows.append(("Max VRAM", f"{self.max_vram_gb:g} GB"))
+        if self.max_model_size_gb is not None:
+            rows.append(("Max model size", f"{self.max_model_size_gb:g} GB"))
+        if self.runtime is not None:
+            rows.append(("Runtime", self.runtime))
+        rows.append(("Context length", f"{self.context_length} tokens"))
+        if self.max_quality_loss is not None:
+            rows.append(("Max quality loss", f"{self.max_quality_loss * 100:g}%"))
+        if self.target_gpu is not None:
+            rows.append(("Target GPU", self.target_gpu))
+        return rows

zeroquantz/optimization/fit.py

@@ -0,0 +1,203 @@
+"""``llmfit``-style memory-fit math.
+
+Models inference VRAM the way the ``llmfit`` tool and the "can I run this LLM"
+calculators (gpu_poor) do:
+
+    total = weights + kv_cache + overhead
+
+* **weights** — parameter count × effective bytes/weight for the format. For
+  whole-model formats (GGUF) the bytes/weight already averages in embeddings; for
+  linear-only formats (AWQ/GPTQ/bnb/TorchAO) embeddings stay at compute precision.
+* **kv_cache** — ``2 · n_layers · n_kv_heads · head_dim · ctx · batch · cache_bytes``
+  (GQA-aware, batch defaults to 1, as llmfit assumes for decode).
+* **overhead** — a flat 0.5 GB for CUDA/context/buffers (llmfit's constant); no
+  separate activation term for single-stream decode.
+
+From this we also derive the **maximum context length** that still fits a budget,
+and a rough QLoRA finetuning estimate (gpu_poor's training formula).
+
+All sizes are GiB (see :mod:`zeroquantz.utils.units`).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel
+
+from zeroquantz.utils import units
+
+if TYPE_CHECKING:
+    from zeroquantz.models.metadata import ModelProfile
+    from zeroquantz.quantization.catalog import QuantFormat
+
+#: Flat CUDA/context/buffer overhead (GiB) — llmfit's constant.
+OVERHEAD_GIB = 0.5
+#: KV-cache element size (fp16) unless a format uses a quantized cache.
+DEFAULT_CACHE_BYTES = 2.0
+#: Bytes/param for the parts a weight-only format keeps at compute precision.
+COMPUTE_BASE_BYTES = 2.0
+
+
+def weight_bytes(model: ModelProfile, fmt: QuantFormat) -> float:
+    """Estimated weight bytes for ``model`` stored in ``fmt``."""
+    bpw_bytes = fmt.bytes_per_weight
+    if fmt.whole_model:
+        return model.parameter_count * bpw_bytes
+    quant = model.quantizable_params * bpw_bytes
+    non_quant = model.non_quantizable_params * COMPUTE_BASE_BYTES
+    return quant + non_quant
+
+
+def kv_cache_bytes(
+    model: ModelProfile,
+    context_length: int,
+    *,
+    batch: int = 1,
+    cache_bytes: float = DEFAULT_CACHE_BYTES,
+) -> float:
+    """KV-cache bytes for a context length and batch (GQA-aware, llmfit formula)."""
+    layers = model.num_layers
+    kv_heads = model.num_key_value_heads or model.num_attention_heads
+    head_dim = model.head_dim
+    if not (layers and kv_heads and head_dim):
+        # llmfit fallback when layer geometry is unknown.
+        return 8e-6 * model.parameter_count * context_length
+    per_token = 2 * layers * kv_heads * head_dim * cache_bytes
+    return per_token * context_length * batch
+
+
+def kv_bytes_per_token(model: ModelProfile, *, cache_bytes: float = DEFAULT_CACHE_BYTES) -> float:
+    layers = model.num_layers
+    kv_heads = model.num_key_value_heads or model.num_attention_heads
+    head_dim = model.head_dim
+    if not (layers and kv_heads and head_dim):
+        return 8e-6 * model.parameter_count
+    return 2 * layers * kv_heads * head_dim * cache_bytes
+
+
+class FitReport(BaseModel):
+    """Result of a fit check for one model + format + budget + context."""
+
+    format_id: str
+    context_length: int
+    batch: int = 1
+    weight_gb: float
+    kv_cache_gb: float
+    overhead_gb: float
+    total_gb: float
+    budget_gb: float | None = None
+    fits: bool = True
+    max_context: int | None = None  # largest context that fits the budget
+
+    def breakdown_rows(self) -> list[tuple[str, str]]:
+        rows = [
+            ("Weights", f"{self.weight_gb:.2f} GB"),
+            (f"KV cache @ {self.context_length} tok", f"{self.kv_cache_gb:.2f} GB"),
+            ("Overhead", f"{self.overhead_gb:.2f} GB"),
+            ("Total", f"{self.total_gb:.2f} GB"),
+        ]
+        if self.budget_gb is not None:
+            rows.append(("Budget", f"{self.budget_gb:g} GB"))
+            rows.append(("Fits", "yes" if self.fits else "no"))
+        if self.max_context is not None:
+            rows.append(("Max context", f"{self.max_context:,} tokens"))
+        return rows
+
+
+def total_vram_gb(
+    model: ModelProfile,
+    fmt: QuantFormat,
+    *,
+    context_length: int = 4096,
+    batch: int = 1,
+) -> float:
+    weights = weight_bytes(model, fmt)
+    kv = kv_cache_bytes(model, context_length, batch=batch)
+    return round(units.bytes_to_gb(weights + kv) + OVERHEAD_GIB, 2)
+
+
+def weight_gb(model: ModelProfile, fmt: QuantFormat) -> float:
+    return round(units.bytes_to_gb(weight_bytes(model, fmt)), 2)
+
+
+def max_context_that_fits(
+    model: ModelProfile,
+    fmt: QuantFormat,
+    budget_gb: float,
+    *,
+    batch: int = 1,
+) -> int:
+    """Largest context length whose total VRAM still fits ``budget_gb`` (0 if none)."""
+    fixed = units.bytes_to_gb(weight_bytes(model, fmt)) + OVERHEAD_GIB
+    remaining_gb = budget_gb - fixed
+    if remaining_gb <= 0:
+        return 0
+    per_token_gb = units.bytes_to_gb(kv_bytes_per_token(model)) * batch
+    if per_token_gb <= 0:
+        return 0
+    return int(remaining_gb / per_token_gb)
+
+
+def fit(
+    model: ModelProfile,
+    fmt: QuantFormat,
+    *,
+    budget_gb: float | None = None,
+    context_length: int = 4096,
+    batch: int = 1,
+) -> FitReport:
+    """Full fit report for a model + format under a budget and context length."""
+    w_bytes = weight_bytes(model, fmt)
+    kv = kv_cache_bytes(model, context_length, batch=batch)
+    weight_g = units.bytes_to_gb(w_bytes)
+    kv_g = units.bytes_to_gb(kv)
+    total = round(weight_g + kv_g + OVERHEAD_GIB, 2)
+    fits = True if budget_gb is None else total <= budget_gb
+    max_ctx = (
+        max_context_that_fits(model, fmt, budget_gb, batch=batch)
+        if budget_gb is not None
+        else None
+    )
+    return FitReport(
+        format_id=fmt.id,
+        context_length=context_length,
+        batch=batch,
+        weight_gb=round(weight_g, 2),
+        kv_cache_gb=round(kv_g, 2),
+        overhead_gb=OVERHEAD_GIB,
+        total_gb=total,
+        budget_gb=budget_gb,
+        fits=fits,
+        max_context=max_ctx,
+    )
+
+
+def qlora_finetune_gb(
+    model: ModelProfile,
+    *,
+    context_length: int = 2048,
+    batch: int = 1,
+) -> float:
+    """Rough QLoRA (4-bit base) finetuning VRAM estimate (gpu_poor-style).
+
+    Approximate: 4-bit frozen weights + fp16 KV/activation working set + a small
+    LoRA/optimizer allowance + overhead. Clearly an estimate, not a guarantee.
+    """
+    from zeroquantz.quantization.catalog import get_format
+
+    nf4 = get_format("nf4")
+    base_bytes = weight_bytes(model, nf4) if nf4 else model.parameter_count * 0.53
+    # Training keeps activations for backprop: ~2x the decode KV working set here.
+    working = 2.0 * kv_cache_bytes(model, context_length, batch=batch)
+    activations = (
+        (model.hidden_size or 4096)
+        * context_length
+        * batch
+        * (model.num_layers or 32)
+        * COMPUTE_BASE_BYTES
+        * 0.25  # gradient-checkpointed approximation
+    )
+    lora_and_opt = units.gb_to_bytes(1.0)  # LoRA params + Adam states + grads (small)
+    total_bytes = base_bytes + working + activations + lora_and_opt
+    return round(units.bytes_to_gb(total_bytes) + OVERHEAD_GIB, 2)

zeroquantz/optimization/pareto.py

@@ -0,0 +1,66 @@
+"""Pareto-frontier computation over quantization candidates.
+
+Objectives, by default:
+
+* minimize ``estimated_vram_gb`` (memory)
+* minimize ``estimated_quality_risk`` (quality)
+* maximize ``estimated_speedup`` (speed)
+
+A candidate is Pareto-optimal if no other candidate is at least as good on every
+objective and strictly better on at least one.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from zeroquantz.optimization.candidate import QuantizationCandidate
+
+# (attribute, direction) where direction is "min" (lower better) or "max".
+Objective = tuple[str, str]
+
+DEFAULT_OBJECTIVES: tuple[Objective, ...] = (
+    ("estimated_vram_gb", "min"),
+    ("estimated_quality_risk", "min"),
+    ("estimated_speedup", "max"),
+)
+
+
+def _no_worse(a_val: float, b_val: float, direction: str) -> bool:
+    return a_val <= b_val if direction == "min" else a_val >= b_val
+
+
+def _strictly_better(a_val: float, b_val: float, direction: str) -> bool:
+    return a_val < b_val if direction == "min" else a_val > b_val
+
+
+def dominates(
+    a: QuantizationCandidate,
+    b: QuantizationCandidate,
+    objectives: tuple[Objective, ...] = DEFAULT_OBJECTIVES,
+) -> bool:
+    """True if ``a`` dominates ``b`` across ``objectives``."""
+    no_worse_all = True
+    strictly_better_any = False
+    for attr, direction in objectives:
+        av = float(getattr(a, attr))
+        bv = float(getattr(b, attr))
+        if not _no_worse(av, bv, direction):
+            no_worse_all = False
+            break
+        if _strictly_better(av, bv, direction):
+            strictly_better_any = True
+    return no_worse_all and strictly_better_any
+
+
+def pareto_frontier(
+    candidates: list[QuantizationCandidate],
+    objectives: tuple[Objective, ...] = DEFAULT_OBJECTIVES,
+) -> list[QuantizationCandidate]:
+    """Return the non-dominated subset of ``candidates`` (input order preserved)."""
+    frontier: list[QuantizationCandidate] = []
+    for cand in candidates:
+        if not any(dominates(other, cand, objectives) for other in candidates if other is not cand):
+            frontier.append(cand)
+    return frontier