zeroquantz 0.1.0__py3-none-any.whl

zeroquantz/hardware/detector.py

@@ -0,0 +1,200 @@
+"""Detect the local hardware and turn it into a :class:`HardwareProfile`.
+
+Uses PyTorch's CUDA APIs as the primary source and NVML (``pynvml`` /
+``nvidia-ml-py``) for richer telemetry (free memory, driver version). Every probe
+is wrapped so that a missing library or an unavailable subsystem degrades
+gracefully to a CPU profile instead of raising.
+"""
+
+from __future__ import annotations
+
+from zeroquantz.hardware.capabilities import HardwareProfile, derive_precision_support
+from zeroquantz.hardware.gpu import GpuInfo
+from zeroquantz.utils import units
+from zeroquantz.utils.logging import get_logger
+
+log = get_logger(__name__)
+
+# Presets used only by :meth:`HardwareDetector.simulate`, so that recommendations
+# can be produced for a target card you are not currently running on.
+_GPU_PRESETS: dict[str, tuple[float, tuple[int, int]]] = {
+    # name substring -> (vram GiB, compute capability)
+    "rtx 4090": (24.0, (8, 9)),
+    "rtx 4080": (16.0, (8, 9)),
+    "rtx 4070": (12.0, (8, 9)),
+    "rtx 4060": (8.0, (8, 9)),
+    "rtx 3090": (24.0, (8, 6)),
+    "rtx 3080": (10.0, (8, 6)),
+    "a100": (80.0, (8, 0)),
+    "h100": (80.0, (9, 0)),
+    "l40s": (48.0, (8, 9)),
+    "l4": (24.0, (8, 9)),
+    "t4": (16.0, (7, 5)),
+    "v100": (32.0, (7, 0)),
+}
+
+
+class HardwareDetector:
+    """Probe the machine for GPUs, VRAM, compute capability, and CUDA versions."""
+
+    @staticmethod
+    def detect() -> HardwareProfile:
+        gpus: list[GpuInfo] = []
+        cuda_available = False
+        cuda_version: str | None = None
+        torch_version: str | None = None
+        driver_version: str | None = None
+        notes: list[str] = []
+
+        # ---- PyTorch probe --------------------------------------------------
+        try:
+            import torch
+
+            torch_version = torch.__version__
+            cuda_version = getattr(torch.version, "cuda", None)
+            cuda_available = bool(torch.cuda.is_available())
+            if cuda_available:
+                for i in range(torch.cuda.device_count()):
+                    props = torch.cuda.get_device_properties(i)
+                    try:
+                        free_b, _total_b = torch.cuda.mem_get_info(i)
+                        available_gb = units.bytes_to_gb(free_b)
+                    except Exception:  # pragma: no cover - driver dependent
+                        available_gb = None
+                    gpus.append(
+                        GpuInfo(
+                            index=i,
+                            name=props.name,
+                            total_vram_gb=round(units.bytes_to_gb(props.total_memory), 2),
+                            available_vram_gb=(
+                                round(available_gb, 2) if available_gb is not None else None
+                            ),
+                            compute_capability=(props.major, props.minor),
+                            multiprocessor_count=getattr(props, "multi_processor_count", None),
+                        )
+                    )
+        except ImportError:
+            notes.append("PyTorch is not installed; GPU detection is limited.")
+        except Exception as exc:  # pragma: no cover - defensive
+            log.warning("torch hardware probe failed: %s", exc)
+            notes.append(f"PyTorch hardware probe failed: {exc}")
+
+        # ---- NVML probe (driver version, free memory, CPU-only fallback) ----
+        nvml_gpus = HardwareDetector._probe_nvml()
+        if nvml_gpus is not None:
+            driver_version = nvml_gpus.get("driver")
+            if not gpus and nvml_gpus.get("gpus"):
+                # torch unavailable but NVML present -> still report the devices.
+                gpus = nvml_gpus["gpus"]
+                cuda_available = bool(gpus)
+            elif gpus and nvml_gpus.get("gpus"):
+                # Enrich torch data with NVML free memory where torch could not.
+                for gpu, nv in zip(gpus, nvml_gpus["gpus"], strict=False):
+                    if gpu.available_vram_gb is None and nv.available_vram_gb is not None:
+                        gpu.available_vram_gb = nv.available_vram_gb
+
+        # ---- CPU / RAM ------------------------------------------------------
+        cpu_cores, system_ram_gb = HardwareDetector._probe_cpu()
+
+        primary_cc = gpus[0].compute_capability if gpus else None
+        precision = derive_precision_support(primary_cc, cuda_available=cuda_available)
+
+        if not gpus:
+            notes.append("No CUDA GPU detected — running in CPU/planning-only mode.")
+
+        return HardwareProfile(
+            gpus=gpus,
+            cuda_available=cuda_available,
+            cuda_version=cuda_version,
+            driver_version=driver_version,
+            torch_version=torch_version,
+            cpu_cores=cpu_cores,
+            system_ram_gb=system_ram_gb,
+            precision_support=precision,
+            notes=notes,
+        )
+
+    @staticmethod
+    def _probe_nvml() -> dict | None:
+        try:
+            import pynvml  # provided by nvidia-ml-py or the legacy pynvml package
+        except ImportError:
+            return None
+        try:
+            pynvml.nvmlInit()
+        except Exception:  # pragma: no cover - no driver
+            return None
+        try:
+            driver = _as_str(pynvml.nvmlSystemGetDriverVersion())
+            gpus: list[GpuInfo] = []
+            for i in range(pynvml.nvmlDeviceGetCount()):
+                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+                mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
+                try:
+                    major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
+                    cc: tuple[int, int] | None = (major, minor)
+                except Exception:
+                    cc = None
+                gpus.append(
+                    GpuInfo(
+                        index=i,
+                        name=_as_str(pynvml.nvmlDeviceGetName(handle)),
+                        total_vram_gb=round(units.bytes_to_gb(mem.total), 2),
+                        available_vram_gb=round(units.bytes_to_gb(mem.free), 2),
+                        compute_capability=cc,
+                    )
+                )
+            return {"driver": driver, "gpus": gpus}
+        except Exception as exc:  # pragma: no cover - defensive
+            log.debug("NVML probe failed: %s", exc)
+            return None
+        finally:
+            try:
+                pynvml.nvmlShutdown()
+            except Exception:
+                pass
+
+    @staticmethod
+    def _probe_cpu() -> tuple[int | None, float | None]:
+        cores: int | None = None
+        ram_gb: float | None = None
+        try:
+            import os
+
+            cores = os.cpu_count()
+        except Exception:
+            pass
+        try:
+            import psutil
+
+            ram_gb = round(units.bytes_to_gb(psutil.virtual_memory().total), 1)
+        except Exception:
+            pass
+        return cores, ram_gb
+
+    @staticmethod
+    def simulate(gpu_name: str) -> HardwareProfile:
+        """Build a profile for a *named* GPU without needing that hardware present.
+
+        Powers ``recommend --gpu <name>``: plan for the card you will deploy on.
+        Raises :class:`ValueError` if the name is not in the preset table.
+        """
+        key = gpu_name.strip().lower()
+        match = next((preset for name, preset in _GPU_PRESETS.items() if name in key), None)
+        if match is None:
+            known = ", ".join(sorted(_GPU_PRESETS))
+            raise ValueError(f"unknown GPU preset {gpu_name!r}; known presets: {known}")
+        vram, cc = match
+        gpu = GpuInfo(index=0, name=gpu_name, total_vram_gb=vram, compute_capability=cc)
+        return HardwareProfile(
+            gpus=[gpu],
+            cuda_available=True,
+            precision_support=derive_precision_support(cc, cuda_available=True),
+            notes=[f"Simulated hardware profile for '{gpu_name}'."],
+        )
+
+
+def _as_str(value: object) -> str:
+    if isinstance(value, bytes):
+        return value.decode("utf-8", errors="replace")
+    return str(value)

zeroquantz/hardware/gpu.py

@@ -0,0 +1,31 @@
+"""Per-GPU information model."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+from zeroquantz.utils import units
+
+
+class GpuInfo(BaseModel):
+    """A single detected CUDA device."""
+
+    index: int
+    name: str
+    total_vram_gb: float = Field(ge=0)
+    available_vram_gb: float | None = Field(default=None, ge=0)
+    compute_capability: tuple[int, int] | None = None
+    multiprocessor_count: int | None = None
+    uuid: str | None = None
+
+    @property
+    def sm(self) -> str | None:
+        """Compute capability rendered like ``"SM89"``."""
+        if self.compute_capability is None:
+            return None
+        major, minor = self.compute_capability
+        return f"SM{major}{minor}"
+
+    @property
+    def total_vram_bytes(self) -> float:
+        return units.gb_to_bytes(self.total_vram_gb)

zeroquantz/models/__init__.py

@@ -0,0 +1,8 @@
+"""Model inspection, metadata, and (lazy) loading."""
+
+from __future__ import annotations
+
+from zeroquantz.models.inspector import ModelInspector
+from zeroquantz.models.metadata import ModelProfile
+
+__all__ = ["ModelInspector", "ModelProfile"]

zeroquantz/models/architecture.py

@@ -0,0 +1,168 @@
+"""Read a Hugging Face ``config.json`` into normalized dimensions, and estimate a
+parameter count from those dimensions when exact tensor headers are unavailable.
+
+The estimate targets modern decoder-only transformers (LLaMA/Qwen/Mistral-style
+gated MLP with grouped-query attention), which is v0.1's scope. It is clearly
+labelled as an *estimate* everywhere it surfaces.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+# Families that use a 3-matrix gated MLP (gate/up/down). Anything else is assumed
+# to use a 2-matrix MLP. Only affects the *estimate* path; the safetensors path is
+# exact regardless.
+_GATED_MLP_FAMILIES = {
+    "llama",
+    "mistral",
+    "mixtral",
+    "qwen2",
+    "qwen3",
+    "qwen2_moe",
+    "gemma",
+    "gemma2",
+    "gemma3",
+    "phi3",
+    "phi",
+    "stablelm",
+    "starcoder2",
+    "cohere",
+    "olmo",
+    "granite",
+    "internlm2",
+    "yi",
+}
+
+_NON_GATED_MLP_FAMILIES = {"gpt2", "gpt_neox", "gptj", "opt", "bloom", "falcon", "mpt"}
+
+
+def _first(config: dict, *keys: str, default: Any = None) -> Any:
+    for key in keys:
+        if key in config and config[key] is not None:
+            return config[key]
+    return default
+
+
+def extract_dimensions(config: dict) -> dict[str, Any]:
+    """Pull normalized architecture dimensions out of a raw HF config dict."""
+    architectures = config.get("architectures") or []
+    architecture = architectures[0] if architectures else None
+
+    num_heads = _first(config, "num_attention_heads", "n_head", "num_heads")
+    hidden = _first(config, "hidden_size", "n_embd", "d_model", "hidden_dim")
+    head_dim = config.get("head_dim")
+    if head_dim is None and hidden and num_heads:
+        head_dim = hidden // num_heads
+
+    return {
+        "architecture": architecture,
+        "model_type": config.get("model_type"),
+        "num_layers": _first(config, "num_hidden_layers", "n_layer", "num_layers"),
+        "hidden_size": hidden,
+        "intermediate_size": _first(
+            config, "intermediate_size", "n_inner", "ffn_dim", "d_ff"
+        ),
+        "num_attention_heads": num_heads,
+        "num_key_value_heads": _first(config, "num_key_value_heads", default=num_heads),
+        "head_dim": head_dim,
+        "vocab_size": config.get("vocab_size"),
+        "max_position_embeddings": _first(
+            config, "max_position_embeddings", "n_positions", "max_seq_len", "seq_length"
+        ),
+        # HF's PretrainedConfig default is True; large models usually override to False.
+        "tie_word_embeddings": bool(config.get("tie_word_embeddings", True)),
+    }
+
+
+def uses_gated_mlp(model_type: str | None) -> bool:
+    """Whether a model family uses a 3-matrix gated MLP (gate/up/down)."""
+    if model_type is None:
+        return True  # modern default
+    mt = model_type.lower()
+    if mt in _NON_GATED_MLP_FAMILIES:
+        return False
+    if mt in _GATED_MLP_FAMILIES:
+        return True
+    return True
+
+
+# Backwards-compatible private alias.
+_uses_gated_mlp = uses_gated_mlp
+
+
+def estimate_parameter_count(config: dict) -> int | None:
+    """Estimate total parameters from config dimensions, or ``None`` if the config
+    lacks the dimensions needed."""
+    dims = extract_dimensions(config)
+    layers = dims["num_layers"]
+    hidden = dims["hidden_size"]
+    inter = dims["intermediate_size"]
+    heads = dims["num_attention_heads"]
+    kv_heads = dims["num_key_value_heads"] or heads
+    head_dim = dims["head_dim"]
+    vocab = dims["vocab_size"]
+
+    if not all((layers, hidden, heads, vocab)):
+        return None
+    if head_dim is None:
+        head_dim = hidden // heads
+    if inter is None:
+        inter = 4 * hidden  # common fallback ratio
+
+    q_proj = hidden * (heads * head_dim)
+    o_proj = (heads * head_dim) * hidden
+    kv_proj = hidden * (kv_heads * head_dim)
+    attn = q_proj + o_proj + 2 * kv_proj
+
+    if _uses_gated_mlp(dims["model_type"]):
+        mlp = 3 * hidden * inter
+    else:
+        mlp = 2 * hidden * inter
+
+    norms = 2 * hidden  # two RMS/LayerNorms per block
+    per_layer = attn + mlp + norms
+
+    embed = vocab * hidden
+    lm_head = 0 if dims["tie_word_embeddings"] else vocab * hidden
+    final_norm = hidden
+
+    return int(layers * per_layer + embed + lm_head + final_norm)
+
+
+# safetensors dtype codes -> canonical names
+_ST_DTYPES = {
+    "F64": "float64",
+    "F32": "float32",
+    "F16": "float16",
+    "BF16": "bfloat16",
+    "I64": "int64",
+    "I32": "int32",
+    "I16": "int16",
+    "I8": "int8",
+    "U8": "uint8",
+    "F8_E4M3": "fp8",
+    "F8_E5M2": "fp8",
+}
+
+# dtypes that represent actual model weights (not integer buffers / rotary caches)
+_WEIGHT_DTYPES = {"float32", "float16", "bfloat16", "float64", "fp8"}
+
+
+def normalize_st_dtype(code: str) -> str:
+    return _ST_DTYPES.get(str(code).upper(), str(code).lower())
+
+
+def dominant_weight_dtype(param_count_by_dtype: dict[str, int]) -> str:
+    """Given ``{dtype_code: n_params}`` pick the dtype that best represents the
+    model's weights (largest float bucket, ignoring integer buffers)."""
+    normalized: dict[str, int] = {}
+    for code, count in param_count_by_dtype.items():
+        name = normalize_st_dtype(code)
+        normalized[name] = normalized.get(name, 0) + int(count)
+
+    weight_only = {d: c for d, c in normalized.items() if d in _WEIGHT_DTYPES}
+    pool = weight_only or normalized
+    if not pool:
+        return "float16"
+    return max(pool.items(), key=lambda kv: kv[1])[0]

zeroquantz/models/downloader.py

@@ -0,0 +1,161 @@
+"""Download a model's weights from the Hub with real-time byte progress.
+
+Files are fetched one at a time (so a download can be cancelled between files and
+overall byte progress is exact), hooking each file's tqdm bar to report byte
+deltas. The essential files are selected — safetensors weights (or .bin when
+there are no safetensors) plus config/tokenizer — skipping other formats
+(.gguf/.onnx/.h5/…) and repo cruft.
+"""
+
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from zeroquantz.core.exceptions import ZeroQuantzError
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class DownloadCancelled(ZeroQuantzError):
+    """Raised when a download is cancelled by the user."""
+
+
+@dataclass(frozen=True)
+class DownloadFile:
+    name: str
+    size: int
+
+
+_KEEP_SUFFIXES = (".safetensors", ".json", ".model", ".txt", ".jinja", ".tiktoken")
+_WEIGHT_FALLBACK = (".bin", ".pth")
+_SKIP_SUFFIXES = (".gguf", ".onnx", ".h5", ".msgpack", ".pt", ".ckpt", ".gitattributes", ".md")
+# Alternative-format subdirectories that hold duplicate configs/tokenizers/weights
+# we never need for a transformers/vLLM load — skip them so we don't download, e.g.,
+# an `onnx/` copy of the tokenizer alongside the real one.
+_SKIP_DIRS = frozenset({"onnx", "coreml", "tf", "flax", "tflite", "openvino"})
+
+
+def plan_download(model_id: str, *, revision: str | None = None, token: str | None = None) -> tuple[list[DownloadFile], int]:
+    """Return the essential files to download and their total size in bytes."""
+    from huggingface_hub import HfApi
+
+    api = HfApi(token=token)
+    try:
+        info = api.model_info(model_id, revision=revision, files_metadata=True)
+    except Exception as exc:
+        raise ZeroQuantzError(
+            f"Could not read the file list for '{model_id}'.", detail=str(exc)
+        ) from exc
+
+    siblings = info.siblings or []
+    names = [s.rfilename for s in siblings]
+    has_safetensors = any(n.endswith(".safetensors") for n in names)
+
+    files: list[DownloadFile] = []
+    for sib in siblings:
+        name = sib.rfilename
+        low = name.lower()
+        if low.endswith(_SKIP_SUFFIXES):
+            continue
+        if "/" in low and low.split("/", 1)[0] in _SKIP_DIRS:
+            continue
+        keep = (
+            low.endswith(_KEEP_SUFFIXES)
+            or "tokenizer" in low
+            or (low.endswith(_WEIGHT_FALLBACK) and not has_safetensors)
+        )
+        if keep:
+            files.append(DownloadFile(name=name, size=int(sib.size or 0)))
+    total = sum(f.size for f in files)
+    return files, total
+
+
+def _make_tqdm(on_delta: Callable[[int], None]):
+    from tqdm.auto import tqdm as _base
+
+    class _ProgressTqdm(_base):
+        def __init__(self, *args, **kwargs) -> None:
+            # Silence the child bar (would corrupt the TUI); we report via on_delta.
+            kwargs["disable"] = True
+            super().__init__(*args, **kwargs)
+
+        def update(self, n: float | None = 1) -> bool | None:
+            # hf_hub only uses this tqdm_class for a file's byte download, so every
+            # positive delta is bytes. (We must NOT gate on ``self.unit``: tqdm's
+            # __init__ returns early when ``disable=True`` and never assigns it.)
+            try:
+                if n:
+                    on_delta(int(n))
+            except Exception:
+                pass
+            return super().update(n)
+
+    return _ProgressTqdm
+
+
+def download_model(
+    model_id: str,
+    files: list[DownloadFile],
+    *,
+    revision: str | None = None,
+    token: str | None = None,
+    cancel: threading.Event | None = None,
+    on_progress: Callable[[int], None] | None = None,
+) -> str:
+    """Download ``files`` for ``model_id``; return the local snapshot directory.
+
+    Progress is reported two ways so it is always correct and always reaches
+    ``total``: intra-file byte deltas (from each file's tqdm) give a smooth bar,
+    and a snap to the exact byte boundary after every completed file guarantees
+    the reported total lands on 100% even when deltas are unavailable — e.g. tiny
+    files with no progress bar, or the ``hf_transfer`` accelerator, which doesn't
+    drive our tqdm.
+    """
+    from huggingface_hub import hf_hub_download, snapshot_download
+
+    total = sum(f.size for f in files)
+    completed = 0  # bytes from fully-finished files
+
+    def report(done: int) -> None:
+        if on_progress:
+            on_progress(min(done, total))
+
+    for f in files:
+        if cancel is not None and cancel.is_set():
+            raise DownloadCancelled("Download cancelled.")
+        file_seen = 0
+
+        # ``_base``/``_size`` are bound at closure creation (per file); ``completed``
+        # is not mutated during this file's download, so the pre-file base is correct.
+        def bump(delta: int, _base: int = completed, _size: int = f.size) -> None:
+            nonlocal file_seen
+            file_seen = min(file_seen + delta, _size)
+            report(_base + file_seen)
+
+        try:
+            hf_hub_download(
+                model_id, f.name, revision=revision, token=token,
+                tqdm_class=_make_tqdm(bump),
+            )
+        except Exception as exc:
+            raise ZeroQuantzError(
+                f"Failed to download '{f.name}' from '{model_id}'.", detail=str(exc)
+            ) from exc
+        completed += f.size
+        report(completed)  # snap to the file boundary (covers any missed deltas)
+
+    report(total)  # guarantee the bar reaches 100%
+
+    # Resolve the snapshot root from the now-cached files (no re-download).
+    try:
+        root = snapshot_download(
+            model_id, revision=revision, token=token,
+            allow_patterns=[f.name for f in files], local_files_only=True,
+        )
+    except Exception:
+        root = str(Path.home() / ".cache" / "huggingface")
+    return root

zeroquantz/models/hf_auth.py

@@ -0,0 +1,105 @@
+"""Hugging Face token management.
+
+Validates and persists a token using the standard huggingface_hub token store, so
+every subsequent Hub call (search, inspect, download) authenticates automatically
+without threading a token through the code. The token itself is never logged or
+displayed — only a masked form or the resolved username.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import io
+import os
+
+from zeroquantz.core.exceptions import ZeroQuantzError
+from zeroquantz.utils.logging import get_logger
+
+log = get_logger(__name__)
+
+_ENV_VARS = ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HUGGINGFACE_HUB_TOKEN")
+
+
+def current_token() -> str | None:
+    """The effective token (env var or saved file), or None. No network."""
+    try:
+        from huggingface_hub import get_token
+
+        return get_token()
+    except Exception:
+        for var in _ENV_VARS:
+            if os.environ.get(var):
+                return os.environ[var]
+        return None
+
+
+def masked(token: str | None) -> str | None:
+    if not token:
+        return None
+    t = token.strip()
+    return f"{t[:3]}…{t[-4:]}" if len(t) > 10 else "set"
+
+
+def whoami_name(token: str | None = None) -> str | None:
+    """Resolve the username for a token (network). Raises if invalid/unreachable."""
+    from huggingface_hub import whoami
+
+    info = whoami(token=token) if token else whoami()
+    if isinstance(info, dict):
+        return info.get("name") or info.get("fullname")
+    return None
+
+
+def save_token(token: str) -> str | None:
+    """Validate ``token`` (via whoami) and persist it. Returns the username.
+
+    Raises :class:`ZeroQuantzError` with an actionable message if the token is
+    invalid or the Hub is unreachable.
+    """
+    token = token.strip()
+    if not token:
+        raise ZeroQuantzError("No token entered.")
+
+    try:
+        name = whoami_name(token)
+    except Exception as exc:
+        raise ZeroQuantzError(
+            "That Hugging Face token was rejected.",
+            detail=str(exc),
+            suggestions=["create a token at https://huggingface.co/settings/tokens"],
+        ) from exc
+
+    # Effective immediately for this process...
+    os.environ["HF_TOKEN"] = token
+    # ...and persisted for future sessions (quietly — no stray prints in the TUI).
+    with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
+        try:
+            from huggingface_hub import HfFolder
+
+            HfFolder.save_token(token)
+        except Exception:
+            try:
+                from huggingface_hub import login
+
+                login(token=token, add_to_git_credential=False)
+            except Exception as exc:  # pragma: no cover - persistence best-effort
+                log.warning("token set for this session but not persisted: %s", exc)
+    return name
+
+
+def clear_token() -> None:
+    """Remove the token from the environment and the saved store."""
+    for var in _ENV_VARS:
+        os.environ.pop(var, None)
+    with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
+        try:
+            from huggingface_hub import logout
+
+            logout()
+        except Exception:  # pragma: no cover
+            try:
+                from huggingface_hub import HfFolder
+
+                HfFolder.delete_token()
+            except Exception:
+                pass