zeroquantz 0.1.0__py3-none-any.whl

zeroquantz/deploy/targets.py

@@ -0,0 +1,154 @@
+"""Deployment targets: the runtimes ZeroQuantz can generate serve/deploy assets for.
+
+Each target maps a quantization format to the right serve command, the correct
+``--quantization`` flag, a container image, and (for OpenAI-compatible servers) a
+client snippet. Templates are best-effort against current runtime CLIs.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from zeroquantz.quantization.catalog import QuantFormat
+
+
+@dataclass(frozen=True)
+class DeployTarget:
+    name: str                       # matches catalog runtime names
+    label: str
+    kind: str                       # "server" | "library" | "gguf"
+    default_port: int
+    openai_compatible: bool
+    container_image: str | None
+    install_hint: str
+    #: format.family -> serving --quantization flag (None = auto-detected / not needed)
+    quant_flags: dict[str, str | None] = field(default_factory=dict)
+    serve_template: Callable[[str, QuantFormat, int, str | None], str] | None = None
+    note: str = ""
+    # ---- isolated sub-environment (vLLM / SGLang / TensorRT-LLM) ----
+    #: When True, this runtime is provisioned in its own venv under
+    #: ~/.zeroquantz/envs/serve/<name>/ (its stack conflicts with the others).
+    isolated_env: bool = False
+    #: pip packages that install the runtime into that env.
+    pip: tuple[str, ...] = ()
+    #: extra package index (e.g. NVIDIA's PyPI for TensorRT-LLM).
+    extra_index_url: str | None = None
+    #: Python version for the venv.
+    env_python: str = "3.11"
+
+    def quant_flag(self, fmt: QuantFormat | None) -> str | None:
+        if fmt is not None and fmt.family in self.quant_flags:
+            return self.quant_flags[fmt.family]
+        return None
+
+    def serve_command(
+        self, model_path: str, fmt: QuantFormat | None = None, *, port: int | None = None
+    ) -> str:
+        port = port or self.default_port
+        flag = self.quant_flag(fmt)
+        return self.serve_template(model_path, fmt, port, flag) if self.serve_template else ""
+
+
+# ---- serve-command templates ------------------------------------------------
+
+
+def _vllm(path: str, fmt: QuantFormat, port: int, flag: str | None) -> str:
+    cmd = f"vllm serve {path}"
+    if flag:
+        cmd += f" --quantization {flag}"
+    return cmd + f" --max-model-len 4096 --port {port}"
+
+
+def _sglang(path: str, fmt: QuantFormat, port: int, flag: str | None) -> str:
+    cmd = f"python -m sglang.launch_server --model-path {path}"
+    if flag:
+        cmd += f" --quantization {flag}"
+    return cmd + f" --host 0.0.0.0 --port {port}"
+
+
+def _trtllm(path: str, fmt: QuantFormat, port: int, flag: str | None) -> str:
+    return f"trtllm-serve {path} --backend pytorch --host 0.0.0.0 --port {port}"
+
+
+def _transformers(path: str, fmt: QuantFormat, port: int, flag: str | None) -> str:
+    return (
+        "python -c \"from transformers import AutoModelForCausalLM, AutoTokenizer; "
+        f"m=AutoModelForCausalLM.from_pretrained('{path}', device_map='auto'); "
+        f"t=AutoTokenizer.from_pretrained('{path}'); print('loaded', m.config.model_type)\""
+    )
+
+
+def _llamacpp(path: str, fmt: QuantFormat, port: int, flag: str | None) -> str:
+    return f"llama-server -m {path} --host 0.0.0.0 --port {port} -c 4096 -ngl 999"
+
+
+def _ollama(path: str, fmt: QuantFormat, port: int, flag: str | None) -> str:
+    return "ollama create zeroquantz-model -f Modelfile && ollama run zeroquantz-model"
+
+
+# ---- the targets ------------------------------------------------------------
+
+_VLLM_FLAGS = {
+    "awq": "awq_marlin", "gptq": "gptq_marlin", "autoround": "gptq_marlin",
+    "fp8": "fp8", "compressed_tensors": None, "torchao": "torchao",
+    "bitsandbytes": "bitsandbytes", "mxfp4": "mxfp4", "nvfp4": "modelopt_fp4",
+    "hqq": "hqq", "aqlm": "aqlm",
+}
+_SGLANG_FLAGS = {
+    "awq": "awq_marlin", "gptq": "gptq_marlin", "autoround": "gptq_marlin",
+    "fp8": "fp8", "compressed_tensors": "w8a8_int8", "nvfp4": "modelopt_fp4",
+    "mxfp4": "mxfp4",
+}
+_TRT_FLAGS = {"fp8": "fp8", "nvfp4": "nvfp4", "awq": "int4_awq", "gptq": "int4_gptq"}
+
+_TARGETS: dict[str, DeployTarget] = {
+    "vllm": DeployTarget(
+        "vllm", "vLLM", "server", 8000, True, "vllm/vllm-openai:latest",
+        "pip install vllm  (Linux + CUDA)", _VLLM_FLAGS, _vllm,
+        note="High-throughput OpenAI-compatible server.",
+        isolated_env=True, pip=("vllm",), env_python="3.12",
+    ),
+    "sglang": DeployTarget(
+        "sglang", "SGLang", "server", 30000, True, "lmsysorg/sglang:latest",
+        "pip install \"sglang[all]\"  (Linux + CUDA)", _SGLANG_FLAGS, _sglang,
+        note="High-throughput server with RadixAttention.",
+        isolated_env=True, pip=("sglang[all]",), env_python="3.12",
+    ),
+    "tensorrt_llm": DeployTarget(
+        "tensorrt_llm", "TensorRT-LLM", "server", 8000, True,
+        "nvcr.io/nvidia/tritonserver:24.10-trtllm-python-py3",
+        "pip install tensorrt-llm  (Linux; FP4 needs Blackwell)", _TRT_FLAGS, _trtllm,
+        note="NVIDIA's fastest inference; FP8/FP4 on recent GPUs.",
+        isolated_env=True, pip=("tensorrt-llm",),
+        extra_index_url="https://pypi.nvidia.com", env_python="3.12",
+    ),
+    "transformers": DeployTarget(
+        "transformers", "Transformers", "library", 0, False, "pytorch/pytorch:latest",
+        "pip install transformers accelerate", {}, _transformers,
+        note="Library load (not a server); good for scripts / TGI.",
+    ),
+    "llamacpp": DeployTarget(
+        "llamacpp", "llama.cpp", "gguf", 8080, True, "ghcr.io/ggml-org/llama.cpp:server",
+        "download a llama.cpp release (prebuilt binaries)", {}, _llamacpp,
+        note="Runs GGUF on CPU + many GPUs; OpenAI-compatible /v1.",
+    ),
+    "ollama": DeployTarget(
+        "ollama", "Ollama", "gguf", 11434, True, "ollama/ollama:latest",
+        "install Ollama from ollama.com", {}, _ollama,
+        note="Simplest local GGUF serving.",
+    ),
+}
+
+
+def get_target(name: str) -> DeployTarget | None:
+    return _TARGETS.get(name)
+
+
+def deploy_targets_for(fmt: QuantFormat) -> list[DeployTarget]:
+    """Targets that can serve ``fmt`` (intersection of the format's runtimes and
+    the targets we can generate assets for)."""
+    targets = [_TARGETS[r] for r in fmt.runtimes if r in _TARGETS]
+    return targets or [_TARGETS["transformers"]]

zeroquantz/export/__init__.py

@@ -0,0 +1,8 @@
+"""Export quantized models and generate ZeroQuantz reports."""
+
+from __future__ import annotations
+
+from zeroquantz.export.exporter import Exporter
+from zeroquantz.export.report import QuantizationReport, build_report
+
+__all__ = ["Exporter", "QuantizationReport", "build_report"]

zeroquantz/export/exporter.py

@@ -0,0 +1,68 @@
+"""Export orchestration: write a quantized model + report, safely.
+
+A partially-written export is marked with a ``.zeroquantz_incomplete`` sentinel
+that is only removed once everything has been written, so an interrupted export
+never masquerades as a finished model (spec §22).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from zeroquantz.core.exceptions import ZeroQuantzError
+from zeroquantz.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from zeroquantz.export.report import QuantizationReport
+    from zeroquantz.quantization.base import QuantizationBackend, QuantizationResult
+
+log = get_logger(__name__)
+
+_INCOMPLETE_MARKER = ".zeroquantz_incomplete"
+
+
+class Exporter:
+    """Write models and reports to disk with completeness guarantees."""
+
+    @staticmethod
+    def write_report(report: QuantizationReport, output_dir: str | Path) -> tuple[Path, Path]:
+        """Write just the JSON + Markdown report (no model weights)."""
+        return report.write(output_dir)
+
+    @staticmethod
+    def export_model(
+        result: QuantizationResult,
+        output_dir: str | Path,
+        backend: QuantizationBackend,
+        *,
+        report: QuantizationReport | None = None,
+    ) -> Path:
+        """Write the quantized model (and optional report) to ``output_dir``."""
+        out = Path(output_dir)
+        out.mkdir(parents=True, exist_ok=True)
+        marker = out / _INCOMPLETE_MARKER
+        marker.write_text("export in progress\n", encoding="utf-8")
+
+        try:
+            backend.export(result, str(out))
+            if report is not None:
+                report.model_size_after_gb = (
+                    result.quantized_size_gb or report.model_size_after_gb
+                )
+                report.write(out)
+        except ZeroQuantzError:
+            raise
+        except Exception as exc:
+            raise ZeroQuantzError(
+                f"Failed to export the quantized model to '{out}'.",
+                detail=str(exc),
+                suggestions=["ensure the target directory is writable"],
+            ) from exc
+        else:
+            marker.unlink(missing_ok=True)
+        return out
+
+    @staticmethod
+    def is_incomplete(directory: str | Path) -> bool:
+        return (Path(directory) / _INCOMPLETE_MARKER).exists()

zeroquantz/export/report.py

@@ -0,0 +1,203 @@
+"""The ZeroQuantz report: ``zeroquantz_report.json`` (machine-readable) and
+``zeroquantz_report.md`` (human-readable, HF-card friendly)."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field
+
+from zeroquantz.benchmark.runner import BenchmarkResult
+from zeroquantz.quantization.config import QuantizationConfig
+from zeroquantz.utils import units
+from zeroquantz.utils.env import capture_environment, utc_timestamp
+from zeroquantz.verification.report import VerificationReport
+
+if TYPE_CHECKING:
+    from zeroquantz.hardware.capabilities import HardwareProfile
+    from zeroquantz.models.metadata import ModelProfile
+    from zeroquantz.optimization.constraints import OptimizationGoal
+    from zeroquantz.optimization.planner import MixedPrecisionPlan
+
+
+class QuantizationReport(BaseModel):
+    """Everything worth recording about a quantization run."""
+
+    base_model: str
+    architecture: str | None = None
+    model_type: str | None = None
+    parameter_count: int = 0
+    original_dtype: str | None = None
+
+    backend: str | None = None
+    method: str | None = None
+    quantization_config: dict = Field(default_factory=dict)
+    calibration: dict = Field(default_factory=dict)
+
+    model_size_before_gb: float | None = None
+    model_size_after_gb: float | None = None
+    peak_vram_gb: float | None = None
+
+    runtime_target: str | None = None
+    goal: dict = Field(default_factory=dict)
+    plan: dict | None = None
+
+    benchmark: BenchmarkResult | None = None
+    verification: VerificationReport | None = None
+
+    zeroquantz_version: str = ""
+    timestamp: str = ""
+    environment: dict = Field(default_factory=dict)
+
+    # ---- serialization ------------------------------------------------------
+
+    def to_json(self, *, indent: int = 2) -> str:
+        return json.dumps(self.model_dump(mode="json"), indent=indent)
+
+    def to_markdown(self) -> str:
+        lines: list[str] = []
+        a = lines.append
+        a(f"# ZeroQuantz Report — {self.base_model}")
+        a("")
+        a(f"*Generated {self.timestamp} by ZeroQuantz {self.zeroquantz_version}*")
+        a("")
+
+        a("## Model")
+        a("")
+        a("| Field | Value |")
+        a("| --- | --- |")
+        a(f"| Base model | `{self.base_model}` |")
+        if self.architecture:
+            a(f"| Architecture | {self.architecture} |")
+        if self.parameter_count:
+            a(f"| Parameters | {units.humanize_params(self.parameter_count)} |")
+        if self.original_dtype:
+            a(f"| Original dtype | {self.original_dtype.upper()} |")
+        a("")
+
+        a("## Quantization")
+        a("")
+        a("| Field | Value |")
+        a("| --- | --- |")
+        if self.backend:
+            a(f"| Backend | {self.backend} |")
+        if self.method:
+            a(f"| Method | {self.method} |")
+        if self.runtime_target:
+            a(f"| Runtime target | {self.runtime_target} |")
+        if self.model_size_before_gb is not None:
+            a(f"| Size before | {self.model_size_before_gb:.2f} GB |")
+        if self.model_size_after_gb is not None:
+            a(f"| Size after | {self.model_size_after_gb:.2f} GB |")
+        if self.peak_vram_gb is not None:
+            a(f"| Peak VRAM | {self.peak_vram_gb:.2f} GB |")
+        a("")
+        if self.quantization_config:
+            a("<details><summary>Full quantization config</summary>")
+            a("")
+            a("```json")
+            a(json.dumps(self.quantization_config, indent=2))
+            a("```")
+            a("</details>")
+            a("")
+
+        if self.plan:
+            a("## Mixed-precision plan")
+            a("")
+            a("| Modules | Precision |")
+            a("| --- | --- |")
+            for name, prec in self.plan.get("groups", []):
+                a(f"| {name} | {prec} |")
+            a("")
+
+        if self.benchmark is not None:
+            a("## Benchmark (measured)")
+            a("")
+            a("| Metric | Value |")
+            a("| --- | --- |")
+            for label, value in self.benchmark.summary_rows():
+                a(f"| {label} | {value} |")
+            a("")
+
+        if self.verification is not None:
+            a("## Verification")
+            a("")
+            a("| Metric | Value |")
+            a("| --- | --- |")
+            for label, value in self.verification.summary_rows():
+                a(f"| {label} | {value} |")
+            a("")
+
+        a("## Environment")
+        a("")
+        a("```json")
+        a(json.dumps(self.environment, indent=2))
+        a("```")
+        a("")
+        return "\n".join(lines)
+
+    def write(self, output_dir: str | Path) -> tuple[Path, Path]:
+        out = Path(output_dir)
+        out.mkdir(parents=True, exist_ok=True)
+        json_path = out / "zeroquantz_report.json"
+        md_path = out / "zeroquantz_report.md"
+        json_path.write_text(self.to_json(), encoding="utf-8")
+        md_path.write_text(self.to_markdown(), encoding="utf-8")
+        return json_path, md_path
+
+
+def build_report(
+    model: ModelProfile,
+    *,
+    config: QuantizationConfig | None = None,
+    hardware: HardwareProfile | None = None,
+    goal: OptimizationGoal | None = None,
+    plan: MixedPrecisionPlan | None = None,
+    benchmark: BenchmarkResult | None = None,
+    verification: VerificationReport | None = None,
+    size_after_gb: float | None = None,
+    peak_vram_gb: float | None = None,
+    calibration: dict | None = None,
+) -> QuantizationReport:
+    """Assemble a :class:`QuantizationReport` from the pieces ZeroQuantz has."""
+    from zeroquantz import __version__
+
+    plan_dict = None
+    if plan is not None:
+        plan_dict = {
+            "base_precision": plan.base_precision,
+            "low_precision": plan.low_precision,
+            "high_precision": plan.high_precision,
+            "estimated_size_gb": plan.estimated_size_gb,
+            "estimated_vram_gb": plan.estimated_vram_gb,
+            "groups": plan.grouped_rows(),
+        }
+
+    return QuantizationReport(
+        base_model=model.model_id,
+        architecture=model.architecture,
+        model_type=model.model_type,
+        parameter_count=model.parameter_count,
+        original_dtype=model.dtype,
+        backend=config.backend if config else None,
+        method=config.method if config else None,
+        quantization_config=config.model_dump(mode="json") if config else {},
+        calibration=calibration or {},
+        model_size_before_gb=model.estimated_weight_size_gb,
+        model_size_after_gb=size_after_gb
+        if size_after_gb is not None
+        else (benchmark.model_size_gb if benchmark else None),
+        peak_vram_gb=peak_vram_gb
+        if peak_vram_gb is not None
+        else (benchmark.peak_vram_gb if benchmark else None),
+        runtime_target=goal.runtime if goal else None,
+        goal=goal.model_dump(mode="json") if goal else {},
+        plan=plan_dict,
+        benchmark=benchmark,
+        verification=verification,
+        zeroquantz_version=__version__,
+        timestamp=utc_timestamp(),
+        environment=capture_environment(hardware),
+    )

zeroquantz/hardware/__init__.py

@@ -0,0 +1,15 @@
+"""Hardware detection and capability modelling."""
+
+from __future__ import annotations
+
+from zeroquantz.hardware.capabilities import HardwareProfile, PrecisionStatus, derive_precision_support
+from zeroquantz.hardware.detector import HardwareDetector
+from zeroquantz.hardware.gpu import GpuInfo
+
+__all__ = [
+    "GpuInfo",
+    "HardwareDetector",
+    "HardwareProfile",
+    "PrecisionStatus",
+    "derive_precision_support",
+]

zeroquantz/hardware/capabilities.py

@@ -0,0 +1,152 @@
+"""Hardware capability modelling: which numeric precisions a device supports.
+
+This is the single home for "does this GPU support precision X" reasoning. Both
+the ``/hardware`` view and the candidate generator consult it, so support rules
+live here rather than being scattered across the codebase.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+from zeroquantz.hardware.gpu import GpuInfo
+
+# The set of precisions ZeroQuantz reasons about, in display order.
+PRECISIONS: tuple[str, ...] = ("fp32", "fp16", "bf16", "int8", "int4", "fp8")
+
+
+class PrecisionStatus(str, Enum):
+    """How well a precision is supported on a given device."""
+
+    SUPPORTED = "supported"
+    LIMITED = "limited"
+    UNSUPPORTED = "unsupported"
+
+    @property
+    def label(self) -> str:
+        return {"supported": "yes", "limited": "limited", "unsupported": "no"}[self.value]
+
+    def __bool__(self) -> bool:  # truthy if usable at all
+        return self is not PrecisionStatus.UNSUPPORTED
+
+
+def derive_precision_support(
+    compute_capability: tuple[int, int] | None,
+    *,
+    cuda_available: bool,
+) -> dict[str, PrecisionStatus]:
+    """Map compute capability -> per-precision :class:`PrecisionStatus`.
+
+    Rules (CUDA):
+      * fp32/fp16 — supported on every supported CUDA GPU.
+      * bf16      — native on Ampere+ (SM >= 8.0); limited (emulated) below.
+      * int8      — supported from Turing/late-Pascal (SM >= 6.1).
+      * int4      — supported from Turing (SM >= 7.5) via weight-only kernels.
+      * fp8       — full on Hopper+ (SM >= 9.0), limited on Ada (SM 8.9),
+                    unsupported below.
+
+    Without CUDA everything but fp32 is unsupported (CPU int8 exists but is out of
+    scope for v0.1's CUDA-first target).
+    """
+    S, L, U = (
+        PrecisionStatus.SUPPORTED,
+        PrecisionStatus.LIMITED,
+        PrecisionStatus.UNSUPPORTED,
+    )
+
+    if not cuda_available or compute_capability is None:
+        return {"fp32": S, "fp16": L, "bf16": U, "int8": U, "int4": U, "fp8": U}
+
+    major, minor = compute_capability
+    cc = major * 10 + minor  # e.g. (8, 9) -> 89
+
+    return {
+        "fp32": S,
+        "fp16": S if cc >= 60 else L,
+        "bf16": S if cc >= 80 else L,
+        "int8": S if cc >= 61 else L,
+        "int4": S if cc >= 75 else (L if cc >= 61 else U),
+        "fp8": S if cc >= 90 else (L if cc >= 89 else U),
+    }
+
+
+class HardwareProfile(BaseModel):
+    """Detected hardware plus its precision capabilities.
+
+    The spec-shaped scalar fields (``gpu_name``, ``total_vram_gb``,
+    ``compute_capability``, ``cuda_version``, ``supports_*``) are exposed as
+    properties derived from the primary GPU, while ``gpus`` and
+    ``precision_support`` carry the full detail.
+    """
+
+    gpus: list[GpuInfo] = Field(default_factory=list)
+    cuda_available: bool = False
+    cuda_version: str | None = None
+    driver_version: str | None = None
+    torch_version: str | None = None
+    cpu_cores: int | None = None
+    system_ram_gb: float | None = None
+    precision_support: dict[str, PrecisionStatus] = Field(default_factory=dict)
+    notes: list[str] = Field(default_factory=list)
+
+    # ---- primary-GPU conveniences (match the spec's flat shape) -------------
+
+    @property
+    def primary_gpu(self) -> GpuInfo | None:
+        return self.gpus[0] if self.gpus else None
+
+    @property
+    def gpu_name(self) -> str:
+        gpu = self.primary_gpu
+        return gpu.name if gpu else "CPU (no CUDA device)"
+
+    @property
+    def gpu_count(self) -> int:
+        return len(self.gpus)
+
+    @property
+    def total_vram_gb(self) -> float:
+        return self.primary_gpu.total_vram_gb if self.primary_gpu else 0.0
+
+    @property
+    def available_vram_gb(self) -> float | None:
+        return self.primary_gpu.available_vram_gb if self.primary_gpu else None
+
+    @property
+    def compute_capability(self) -> tuple[int, int] | None:
+        return self.primary_gpu.compute_capability if self.primary_gpu else None
+
+    @property
+    def sm(self) -> str | None:
+        return self.primary_gpu.sm if self.primary_gpu else None
+
+    # ---- precision helpers --------------------------------------------------
+
+    def status(self, precision: str) -> PrecisionStatus:
+        return self.precision_support.get(precision, PrecisionStatus.UNSUPPORTED)
+
+    def supports(self, precision: str) -> bool:
+        """True if ``precision`` is usable at all (supported or limited)."""
+        return bool(self.status(precision))
+
+    @property
+    def supports_fp16(self) -> bool:
+        return self.supports("fp16")
+
+    @property
+    def supports_bf16(self) -> bool:
+        return self.supports("bf16")
+
+    @property
+    def supports_int8(self) -> bool:
+        return self.supports("int8")
+
+    @property
+    def supports_int4(self) -> bool:
+        return self.supports("int4")
+
+    @property
+    def supports_fp8(self) -> bool:
+        return self.supports("fp8")