zeroquantz 0.1.0__py3-none-any.whl

zeroquantz/optimization/planner.py

@@ -0,0 +1,297 @@
+"""Deterministic mixed-precision planning.
+
+Given a memory budget, per-layer sensitivity, and the precisions the hardware
+supports, assign a precision to each module. The v0.1 algorithm is an explicit,
+replaceable greedy strategy (spec §12):
+
+1. Start every transformer sub-module at the lowest requested precision.
+2. Keep critical modules (embeddings, LM head) at base precision.
+3. Rank sub-modules by sensitivity (from a profiler, or a sensible prior).
+4. Promote the most sensitive sub-modules to the higher precision while the
+   memory budget still allows.
+
+The result is both a human-readable plan and a :class:`QuantizationConfig` the
+mixed backend can consume.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field
+
+from zeroquantz.models import architecture as arch
+from zeroquantz.profiling.sensitivity import layer_sensitivity_prior
+from zeroquantz.quantization import memory
+from zeroquantz.quantization.config import QuantizationConfig
+from zeroquantz.utils import units
+
+if TYPE_CHECKING:
+    from zeroquantz.hardware.capabilities import HardwareProfile
+    from zeroquantz.models.metadata import ModelProfile
+    from zeroquantz.optimization.constraints import OptimizationGoal
+    from zeroquantz.profiling.sensitivity import SensitivityProfile
+
+_PRECISION_BITS = {"bf16": 16, "fp16": 16, "int8": 8, "int4": 4, "nf4": 4}
+
+
+class PrecisionAssignment(BaseModel):
+    module: str
+    precision: str
+    params: int = 0
+    sensitivity: float | None = None
+    critical: bool = False
+
+    @property
+    def weight_bytes(self) -> float:
+        return units.param_bytes(self.params, self.precision)
+
+
+class MixedPrecisionPlan(BaseModel):
+    """A per-module precision assignment plus its estimated cost."""
+
+    model_id: str
+    base_precision: str
+    low_precision: str
+    high_precision: str
+    assignments: list[PrecisionAssignment]
+    estimated_size_gb: float
+    estimated_vram_gb: float
+    budget_gb: float | None = None
+    fits_budget: bool = True
+    notes: list[str] = Field(default_factory=list)
+
+    def to_config(self, backend: str = "mixed") -> QuantizationConfig:
+        overrides = {
+            a.module: _PRECISION_BITS.get(a.precision, 4)
+            for a in self.assignments
+            if not a.critical
+        }
+        return QuantizationConfig(
+            backend=backend,
+            method="mixed",
+            weight_bits=_PRECISION_BITS.get(self.low_precision, 4),
+            activation_bits=16,
+            compute_dtype="bfloat16" if self.base_precision == "bf16" else "float16",
+            skip_modules=[a.module for a in self.assignments if a.critical],
+            module_overrides=overrides,
+            extra={
+                "low_precision": self.low_precision,
+                "high_precision": self.high_precision,
+                "high_fraction": self._high_fraction(),
+            },
+        )
+
+    def _high_fraction(self) -> float:
+        quant = [a for a in self.assignments if not a.critical]
+        if not quant:
+            return 0.0
+        promoted = sum(a.params for a in quant if a.precision == self.high_precision)
+        total = sum(a.params for a in quant) or 1
+        return round(promoted / total, 3)
+
+    def grouped_rows(self) -> list[tuple[str, str]]:
+        """Collapse consecutive same-precision layers into ranges for display,
+        e.g. ``("layers.1-33", "INT4")`` (spec §12 style)."""
+        rows: list[tuple[str, str]] = []
+        run_start: int | None = None
+        run_prec: str | None = None
+        run_end = -1
+
+        def flush() -> None:
+            if run_start is None or run_prec is None:
+                return
+            name = f"layers.{run_start}" if run_start == run_end else f"layers.{run_start}-{run_end}"
+            rows.append((name, run_prec.upper()))
+
+        layer_assigns = _per_layer_precision(self.assignments)
+        for idx, prec in layer_assigns:
+            if prec == run_prec:
+                run_end = idx
+            else:
+                flush()
+                run_start = run_end = idx
+                run_prec = prec
+        flush()
+
+        # Prepend/append critical modules.
+        critical = [(a.module, a.precision.upper()) for a in self.assignments if a.critical]
+        return critical[:1] + rows + critical[1:]
+
+
+class MixedPrecisionPlanner:
+    """Greedy budget-aware precision assignment (spec §12)."""
+
+    @staticmethod
+    def plan(
+        model: ModelProfile,
+        hardware: HardwareProfile | None = None,
+        goal: OptimizationGoal | None = None,
+        sensitivity: SensitivityProfile | None = None,
+        *,
+        low_precision: str = "int4",
+        high_precision: str = "int8",
+        base_precision: str | None = None,
+    ) -> MixedPrecisionPlan:
+        base_precision = base_precision or (
+            "bf16" if (hardware and hardware.supports_bf16) else "fp16"
+        )
+        assignments = MixedPrecisionPlanner._initial_assignments(
+            model, low_precision, base_precision, sensitivity
+        )
+
+        budget_gb = MixedPrecisionPlanner._resolve_budget(goal, hardware)
+        budget_bytes = units.gb_to_bytes(budget_gb) if budget_gb is not None else None
+        overhead = memory.estimate_runtime_overhead_bytes(
+            model, compute_dtype=base_precision if base_precision != "bf16" else "bfloat16"
+        )
+
+        promotable = sorted(
+            (a for a in assignments if not a.critical and a.params > 0),
+            key=lambda a: (a.sensitivity or 0.0),
+            reverse=True,
+        )
+
+        current = _weight_bytes(assignments)
+        notes: list[str] = []
+
+        if budget_bytes is not None and current + overhead > budget_bytes:
+            notes.append(
+                "Budget cannot be met even with every layer at the lowest precision; "
+                "consider a smaller model or quantizing embeddings/LM head."
+            )
+        else:
+            for assign in promotable:
+                delta = units.param_bytes(assign.params, high_precision) - units.param_bytes(
+                    assign.params, low_precision
+                )
+                if budget_bytes is None:
+                    # No hard budget: promote the most sensitive ~25% by param mass.
+                    promoted_mass = sum(
+                        a.params for a in assignments if a.precision == high_precision
+                    )
+                    quant_mass = sum(a.params for a in assignments if not a.critical) or 1
+                    if promoted_mass / quant_mass >= 0.25:
+                        break
+                    assign.precision = high_precision
+                    current += delta
+                elif current + overhead + delta <= budget_bytes:
+                    assign.precision = high_precision
+                    current += delta
+
+        size_gb = round(units.bytes_to_gb(current), 2)
+        vram_gb = round(units.bytes_to_gb(current + overhead), 2)
+        fits = budget_gb is None or vram_gb <= budget_gb
+
+        return MixedPrecisionPlan(
+            model_id=model.model_id,
+            base_precision=base_precision,
+            low_precision=low_precision,
+            high_precision=high_precision,
+            assignments=assignments,
+            estimated_size_gb=size_gb,
+            estimated_vram_gb=vram_gb,
+            budget_gb=budget_gb,
+            fits_budget=fits,
+            notes=notes,
+        )
+
+    # ---- internals ----------------------------------------------------------
+
+    @staticmethod
+    def _resolve_budget(
+        goal: OptimizationGoal | None, hardware: HardwareProfile | None
+    ) -> float | None:
+        if goal and goal.max_vram_gb is not None:
+            return goal.max_vram_gb
+        if hardware and hardware.total_vram_gb > 0:
+            return round(hardware.total_vram_gb * 0.9, 2)
+        return None
+
+    @staticmethod
+    def _initial_assignments(
+        model: ModelProfile,
+        low_precision: str,
+        base_precision: str,
+        sensitivity: SensitivityProfile | None,
+    ) -> list[PrecisionAssignment]:
+        layers = model.num_layers or 0
+        hidden = model.hidden_size or 0
+        heads = model.num_attention_heads or 0
+        kv_heads = model.num_key_value_heads or heads
+        head_dim = model.head_dim or (hidden // heads if heads else 0)
+        inter = model.intermediate_size or (4 * hidden)
+
+        attn_params = (
+            hidden * (heads * head_dim)
+            + (heads * head_dim) * hidden
+            + 2 * hidden * (kv_heads * head_dim)
+        )
+        gated = arch.uses_gated_mlp(model.model_type)
+        mlp_params = (3 if gated else 2) * hidden * inter
+
+        sens_lookup = sensitivity.as_lookup() if sensitivity else {}
+
+        assignments: list[PrecisionAssignment] = [
+            PrecisionAssignment(
+                module="model.embed_tokens",
+                precision=base_precision,
+                params=model.embedding_params,
+                critical=True,
+            )
+        ]
+        for i in range(layers):
+            attn_mod = f"model.layers.{i}.self_attn"
+            mlp_mod = f"model.layers.{i}.mlp"
+            assignments.append(
+                PrecisionAssignment(
+                    module=attn_mod,
+                    precision=low_precision,
+                    params=attn_params,
+                    sensitivity=sens_lookup.get(attn_mod, layer_sensitivity_prior(i, layers, "attn")),
+                )
+            )
+            assignments.append(
+                PrecisionAssignment(
+                    module=mlp_mod,
+                    precision=low_precision,
+                    params=mlp_params,
+                    sensitivity=sens_lookup.get(mlp_mod, layer_sensitivity_prior(i, layers, "mlp")),
+                )
+            )
+        assignments.append(
+            PrecisionAssignment(
+                module="lm_head",
+                precision=base_precision,
+                params=model.lm_head_params,
+                critical=True,
+            )
+        )
+        return assignments
+
+
+def _weight_bytes(assignments: list[PrecisionAssignment]) -> float:
+    return sum(a.weight_bytes for a in assignments)
+
+
+def _per_layer_precision(
+    assignments: list[PrecisionAssignment],
+) -> list[tuple[int, str]]:
+    """Reduce attn/mlp assignments to one precision per layer index (the lower of
+    the two, since that dominates the visual grouping)."""
+    per_layer: dict[int, list[str]] = {}
+    for a in assignments:
+        if ".layers." not in a.module:
+            continue
+        try:
+            idx = int(a.module.split(".layers.")[1].split(".")[0])
+        except (IndexError, ValueError):
+            continue
+        per_layer.setdefault(idx, []).append(a.precision)
+    result: list[tuple[int, str]] = []
+    for idx in sorted(per_layer):
+        precs = per_layer[idx]
+        # If mixed within a layer, show the lower precision as the layer summary.
+        chosen = min(precs, key=lambda p: _PRECISION_BITS.get(p, 4))
+        result.append((idx, chosen))
+    return result

zeroquantz/optimization/recommender.py

@@ -0,0 +1,149 @@
+"""Rank quantization candidates against an :class:`OptimizationGoal`.
+
+The ranking is a normalized multi-objective score whose weights depend on the
+goal's objective, plus penalties for infeasibility (over budget, quality bound
+exceeded, poor runtime fit, not-yet-implemented). Different objectives genuinely
+reorder the candidates, which is the whole point of the recommender.
+"""
+
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+from zeroquantz.optimization.candidate import QuantizationCandidate
+from zeroquantz.optimization.constraints import Objective, OptimizationGoal
+from zeroquantz.runtimes.compatibility import RuntimeCompat
+
+# objective -> (quality_weight, memory_weight, speed_weight)
+_WEIGHTS: dict[Objective, tuple[float, float, float]] = {
+    Objective.QUALITY: (0.60, 0.25, 0.15),
+    Objective.MEMORY: (0.25, 0.60, 0.15),
+    Objective.SPEED: (0.25, 0.20, 0.55),
+    Objective.BALANCED: (0.34, 0.33, 0.33),
+}
+
+
+class Recommendation(BaseModel):
+    """A scored candidate with a human-readable justification."""
+
+    candidate: QuantizationCandidate
+    score: float
+    reason: str
+    rank: int = 0
+
+
+def _normalize(value: float, lo: float, hi: float, direction: str) -> float:
+    if hi <= lo:
+        return 0.5
+    t = (value - lo) / (hi - lo)
+    return (1.0 - t) if direction == "min" else t
+
+
+class Recommender:
+    """Score and order candidates for a goal."""
+
+    @staticmethod
+    def rank(
+        candidates: list[QuantizationCandidate],
+        goal: OptimizationGoal | None = None,
+    ) -> list[Recommendation]:
+        goal = goal or OptimizationGoal()
+        if not candidates:
+            return []
+
+        vrams = [c.estimated_vram_gb for c in candidates]
+        risks = [c.estimated_quality_risk for c in candidates]
+        speeds = [c.estimated_speedup for c in candidates]
+        v_lo, v_hi = min(vrams), max(vrams)
+        r_lo, r_hi = min(risks), max(risks)
+        s_lo, s_hi = min(speeds), max(speeds)
+
+        qw, mw, sw = _WEIGHTS[goal.objective]
+
+        scored: list[Recommendation] = []
+        for cand in candidates:
+            quality = _normalize(cand.estimated_quality_risk, r_lo, r_hi, "min")
+            mem = _normalize(cand.estimated_vram_gb, v_lo, v_hi, "min")
+            speed = _normalize(cand.estimated_speedup, s_lo, s_hi, "max")
+            score = qw * quality + mw * mem + sw * speed
+            score *= Recommender._penalty(cand)
+            scored.append(
+                Recommendation(
+                    candidate=cand,
+                    score=round(score, 4),
+                    reason=Recommender._reason(cand, goal),
+                )
+            )
+
+        scored.sort(
+            key=lambda r: (r.candidate.feasible, r.score),
+            reverse=True,
+        )
+        for i, rec in enumerate(scored, start=1):
+            rec.rank = i
+        return scored
+
+    @staticmethod
+    def recommend(
+        candidates: list[QuantizationCandidate],
+        goal: OptimizationGoal | None = None,
+    ) -> Recommendation | None:
+        """The single best pick — preferring a feasible, runnable candidate."""
+        ranked = Recommender.rank(candidates, goal)
+        if not ranked:
+            return None
+        for rec in ranked:
+            if rec.candidate.feasible:
+                return rec
+        return ranked[0]
+
+    # ---- internals ----------------------------------------------------------
+
+    @staticmethod
+    def _penalty(cand: QuantizationCandidate) -> float:
+        penalty = 1.0
+        if not cand.fits_memory:
+            penalty *= 0.15
+        if not cand.meets_quality:
+            penalty *= 0.30
+        if not cand.implemented:
+            penalty *= 0.50
+        elif not cand.backend_available:
+            penalty *= 0.85
+        compat_mult = {
+            RuntimeCompat.COMPATIBLE: 1.0,
+            RuntimeCompat.UNKNOWN: 1.0,
+            RuntimeCompat.LIMITED: 0.9,
+            RuntimeCompat.NOT_RECOMMENDED: 0.6,
+            RuntimeCompat.UNSUPPORTED: 0.1,
+        }[cand.runtime_compat]
+        return penalty * compat_mult
+
+    @staticmethod
+    def _reason(cand: QuantizationCandidate, goal: OptimizationGoal) -> str:
+        parts: list[str] = []
+
+        if cand.method == "mixed":
+            parts.append(
+                "Keeps sensitive layers at higher precision while pushing the rest to 4-bit"
+            )
+        elif cand.weight_bits <= 4:
+            parts.append(f"4-bit weights give the smallest footprint (~{cand.estimated_size_gb:.1f} GB)")
+        elif cand.weight_bits == 8:
+            parts.append(f"8-bit weights preserve quality at ~{cand.estimated_size_gb:.1f} GB")
+
+        if goal.max_vram_gb is not None and cand.fits_memory:
+            parts.append(f"fits the {goal.max_vram_gb:g} GB budget (est. {cand.estimated_vram_gb:.1f} GB)")
+        elif goal.max_vram_gb is not None and not cand.fits_memory:
+            parts.append(f"exceeds the {goal.max_vram_gb:g} GB budget (est. {cand.estimated_vram_gb:.1f} GB)")
+
+        parts.append(f"quality risk {cand.quality_risk_label}")
+
+        if goal.runtime:
+            parts.append(f"{cand.runtime_compat.label} with {goal.runtime}")
+
+        if not cand.implemented:
+            parts.append("planned for a future release")
+
+        reason = "; ".join(parts)
+        return reason[0].upper() + reason[1:] + "." if reason else ""

zeroquantz/profiling/__init__.py

@@ -0,0 +1,18 @@
+"""Layer-sensitivity profiling and calibration data."""
+
+from __future__ import annotations
+
+from zeroquantz.profiling.calibration import CalibrationDataset, default_calibration
+from zeroquantz.profiling.sensitivity import (
+    LayerSensitivity,
+    SensitivityProfile,
+    SensitivityProfiler,
+)
+
+__all__ = [
+    "CalibrationDataset",
+    "LayerSensitivity",
+    "SensitivityProfile",
+    "SensitivityProfiler",
+    "default_calibration",
+]

zeroquantz/profiling/calibration.py

@@ -0,0 +1,74 @@
+"""Calibration data for sensitivity profiling and calibrated quantization.
+
+Ships a tiny built-in prompt set so profiling works with zero external data, and
+can load user prompts from ``.jsonl`` (one ``{"text": ...}`` or raw string per
+line) or plain ``.txt`` (one prompt per line).
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+# A small, domain-diverse default calibration set (kept intentionally short).
+_DEFAULT_PROMPTS: tuple[str, ...] = (
+    "The theory of relativity fundamentally changed our understanding of space and time.",
+    "def fibonacci(n):\n    return n if n < 2 else fibonacci(n - 1) + fibonacci(n - 2)",
+    "In 1969, the Apollo 11 mission successfully landed the first humans on the Moon.",
+    "The mitochondria is often described as the powerhouse of the cell.",
+    "To be, or not to be, that is the question posed in Shakespeare's Hamlet.",
+    "A balanced diet includes proteins, carbohydrates, fats, vitamins, and minerals.",
+    "Machine learning models learn patterns from data to make predictions.",
+    "The stock market reacted sharply to the central bank's interest rate decision.",
+    "Photosynthesis converts sunlight, water, and carbon dioxide into glucose and oxygen.",
+    "She walked along the quiet beach as the sun dipped below the horizon.",
+    "The French Revolution began in 1789 and reshaped modern European politics.",
+    "Quantum computers exploit superposition and entanglement to process information.",
+    "A good API is easy to use correctly and hard to use incorrectly.",
+    "The recipe calls for two cups of flour, one egg, and a pinch of salt.",
+    "Climate change is driven largely by greenhouse gas emissions from human activity.",
+    "The ancient library of Alexandria was one of the largest of the ancient world.",
+)
+
+
+class CalibrationDataset(BaseModel):
+    """A named set of calibration prompts."""
+
+    name: str = "default"
+    prompts: list[str] = Field(default_factory=list)
+
+    def __len__(self) -> int:
+        return len(self.prompts)
+
+    def sample(self, n: int) -> list[str]:
+        """First ``n`` prompts (deterministic — reproducible profiling)."""
+        return self.prompts[:n]
+
+    @classmethod
+    def from_file(cls, path: str | Path) -> CalibrationDataset:
+        p = Path(path)
+        if not p.exists():
+            raise FileNotFoundError(f"calibration file not found: {p}")
+        prompts: list[str] = []
+        if p.suffix == ".jsonl":
+            for line in p.read_text(encoding="utf-8").splitlines():
+                line = line.strip()
+                if not line:
+                    continue
+                obj = json.loads(line)
+                if isinstance(obj, str):
+                    prompts.append(obj)
+                elif isinstance(obj, dict):
+                    prompts.append(obj.get("text") or obj.get("prompt") or "")
+        else:
+            prompts = [ln for ln in p.read_text(encoding="utf-8").splitlines() if ln.strip()]
+        return cls(name=p.stem, prompts=[pr for pr in prompts if pr])
+
+
+def default_calibration(n: int | None = None) -> CalibrationDataset:
+    prompts = list(_DEFAULT_PROMPTS)
+    if n is not None:
+        prompts = prompts[:n]
+    return CalibrationDataset(name="default", prompts=prompts)