quantcost 0.2.0__py3-none-any.whl

edgellm/__init__.py

@@ -0,0 +1,7 @@
+"""quantcost: measure what quantization actually costs on your own hardware."""
+
+from __future__ import annotations
+
+__version__ = "0.2.0"
+
+__all__ = ["__version__"]

edgellm/base.py

@@ -0,0 +1,44 @@
+"""Backend-agnostic interfaces, shared by the torch and the torch-free paths.
+
+These used to live in :mod:`edgellm.runners`, which imports ``torch`` at module
+level. The lightweight benchmark path (ONNX Runtime + NumPy, no torch, no
+transformers) needs the same interface without paying for that import, so the
+parts that carry no torch dependency live here and ``runners`` re-exports them.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass
+
+from edgellm.config import GenerationConfig
+
+
+@dataclass
+class GenerationResult:
+    """The text plus the timing/throughput numbers for one generation call."""
+
+    backend: str
+    prompt: str
+    text: str
+    prompt_tokens: int
+    generated_tokens: int
+    latency_s: float
+    tokens_per_second: float
+
+    def as_dict(self) -> dict:
+        return asdict(self)
+
+
+class InferenceRunner(ABC):
+    """Common interface for all inference backends."""
+
+    name: str = "base"
+
+    @abstractmethod
+    def generate(self, prompt: str, generation: GenerationConfig) -> GenerationResult:
+        """Generate a completion for ``prompt`` and report timing."""
+
+    def close(self) -> None:
+        """Release whatever the backend holds (sessions, memory). Optional."""
+        return None

edgellm/benchmark.py

@@ -0,0 +1,226 @@
+"""Benchmark harness: latency, throughput, peak RAM, on-disk size, perplexity.
+
+Every metric here is measured on the real machine. Nothing is estimated or
+scaled. The harness is backend-agnostic: it drives any :class:`InferenceRunner`
+and computes perplexity from any callable that returns logits, so PyTorch and
+ONNX Runtime (FP32 or quantized) are measured identically.
+
+Perplexity method (documented so results are reproducible): the eval text is a
+slice of WikiText-2, tokenized once and split into non-overlapping windows of
+``eval_max_length`` tokens. For each window we compute the summed
+next-token cross-entropy; perplexity is ``exp(total_nll / total_tokens)``.
+"""
+
+from __future__ import annotations
+
+import json
+import statistics
+import threading
+import time
+from collections.abc import Callable
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import psutil
+
+from edgellm.base import InferenceRunner
+from edgellm.config import BenchmarkConfig, GenerationConfig
+
+if TYPE_CHECKING:  # ``torch`` is only needed by PerplexityEvaluator (the heavy path).
+    import torch
+
+
+@dataclass
+class BenchmarkResult:
+    """All measured numbers for one (backend, precision) combination."""
+
+    backend: str
+    precision: str
+    device: str
+    model_id: str
+    size_mb: float | None
+    latency_s_mean: float
+    latency_s_std: float
+    tokens_per_second: float
+    peak_ram_mb: float
+    perplexity: float | None
+    generated_tokens: int
+    measured_runs: int
+
+    def as_dict(self) -> dict:
+        return asdict(self)
+
+
+class _PeakRSSSampler:
+    """Context manager sampling process RSS on a background thread to find its peak."""
+
+    def __init__(self, interval_s: float = 0.02) -> None:
+        self.interval_s = interval_s
+        self.peak_bytes = 0
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+        self._proc = psutil.Process()
+
+    def _run(self) -> None:
+        while not self._stop.is_set():
+            self.peak_bytes = max(self.peak_bytes, self._proc.memory_info().rss)
+            time.sleep(self.interval_s)
+
+    def __enter__(self) -> _PeakRSSSampler:
+        self.peak_bytes = self._proc.memory_info().rss
+        self._thread = threading.Thread(target=self._run, daemon=True)
+        self._thread.start()
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self._stop.set()
+        if self._thread is not None:
+            self._thread.join()
+
+    @property
+    def peak_mb(self) -> float:
+        return self.peak_bytes / (1024 * 1024)
+
+
+def measure_size_mb(path: Path, patterns: tuple[str, ...]) -> float:
+    """Sum the on-disk size (MB) of files under ``path`` matching any glob pattern."""
+    total = 0
+    for pattern in patterns:
+        for f in path.glob(pattern):
+            if f.is_file():
+                total += f.stat().st_size
+    return total / (1024 * 1024)
+
+
+class PerplexityEvaluator:
+    """Compute perplexity on a WikiText slice from any logits-producing callable."""
+
+    def __init__(self, config: BenchmarkConfig) -> None:
+        self.config = config
+
+    def _load_text(self) -> str:
+        from datasets import load_dataset
+
+        ds = load_dataset(self.config.eval_dataset, self.config.eval_config, split="test")
+        lines = [t for t in ds["text"] if t.strip()]
+        return "\n\n".join(lines[: self.config.eval_num_samples])
+
+    def evaluate(
+        self,
+        forward_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        tokenizer,
+    ) -> float:
+        """``forward_fn(input_ids, attention_mask) -> logits`` (all torch tensors)."""
+        import torch
+
+        text = self._load_text()
+        input_ids = tokenizer(text, return_tensors="pt").input_ids
+        max_len = self.config.eval_max_length
+
+        total_nll = torch.zeros((), dtype=torch.float64)
+        total_tokens = 0
+        for start in range(0, input_ids.shape[1], max_len):
+            window = input_ids[:, start : start + max_len]
+            if window.shape[1] < 2:
+                continue
+            attn = torch.ones_like(window)
+            logits = forward_fn(window, attn).float()
+            shift_logits = logits[:, :-1, :].reshape(-1, logits.size(-1))
+            shift_labels = window[:, 1:].reshape(-1).to(shift_logits.device)
+            nll = torch.nn.functional.cross_entropy(shift_logits, shift_labels, reduction="sum")
+            total_nll += nll.double().cpu()
+            total_tokens += int(shift_labels.numel())
+
+        return float(torch.exp(total_nll / total_tokens))
+
+
+class BenchmarkHarness:
+    """Drive a runner through warmup + measured runs and collect all metrics."""
+
+    def __init__(self, config: BenchmarkConfig) -> None:
+        self.config = config
+
+    def _bench_generation_config(self, base: GenerationConfig) -> GenerationConfig:
+        """Deterministic, fixed-length decoding so runs are comparable."""
+        n = self.config.gen_tokens
+        return GenerationConfig(
+            max_new_tokens=n,
+            min_new_tokens=n,  # force full length -> stable token count across backends
+            do_sample=False,
+            seed=base.seed,
+        )
+
+    def run(
+        self,
+        runner: InferenceRunner,
+        *,
+        precision: str,
+        device: str,
+        model_id: str,
+        generation: GenerationConfig,
+        size_mb: float | None = None,
+        perplexity: float | None = None,
+    ) -> BenchmarkResult:
+        gen = self._bench_generation_config(generation)
+        prompt = self.config.prompt
+
+        for _ in range(self.config.warmup_runs):
+            runner.generate(prompt, gen)
+
+        latencies: list[float] = []
+        throughputs: list[float] = []
+        generated_tokens = 0
+        with _PeakRSSSampler() as sampler:
+            for _ in range(self.config.measured_runs):
+                result = runner.generate(prompt, gen)
+                latencies.append(result.latency_s)
+                throughputs.append(result.tokens_per_second)
+                generated_tokens = result.generated_tokens
+
+        return BenchmarkResult(
+            backend=runner.name,
+            precision=precision,
+            device=device,
+            model_id=model_id,
+            size_mb=round(size_mb, 2) if size_mb is not None else None,
+            latency_s_mean=round(statistics.fmean(latencies), 4),
+            latency_s_std=round(statistics.pstdev(latencies), 4),
+            tokens_per_second=round(statistics.fmean(throughputs), 2),
+            peak_ram_mb=round(sampler.peak_mb, 1),
+            perplexity=round(perplexity, 3) if perplexity is not None else None,
+            generated_tokens=generated_tokens,
+            measured_runs=self.config.measured_runs,
+        )
+
+
+def save_results(results: list[BenchmarkResult], path: Path) -> None:
+    """Merge results into a JSON file keyed by (backend, precision)."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    existing: dict[str, dict] = {}
+    if path.exists():
+        for row in json.loads(path.read_text()):
+            existing[f"{row['backend']}|{row['precision']}"] = row
+    for r in results:
+        existing[f"{r.backend}|{r.precision}"] = r.as_dict()
+    path.write_text(json.dumps(list(existing.values()), indent=2) + "\n")
+
+
+def render_markdown(results_json: Path) -> str:
+    """Render the JSON results as a Markdown table (real numbers only)."""
+    rows = json.loads(results_json.read_text()) if results_json.exists() else []
+    header = (
+        "| Backend | Precision | Device | Size (MB) | Latency (s) | "
+        "Throughput (tok/s) | Peak RAM (MB) | Perplexity |\n"
+        "| --- | --- | --- | --- | --- | --- | --- | --- |\n"
+    )
+    lines = []
+    for r in rows:
+        size = f"{r['size_mb']}" if r.get("size_mb") is not None else "—"
+        ppl = f"{r['perplexity']}" if r.get("perplexity") is not None else "—"
+        lat = f"{r['latency_s_mean']} ± {r['latency_s_std']}"
+        lines.append(
+            f"| {r['backend']} | {r['precision']} | {r['device']} | {size} | "
+            f"{lat} | {r['tokens_per_second']} | {r['peak_ram_mb']} | {ppl} |"
+        )
+    return header + "\n".join(lines) + "\n"

edgellm/card.py

@@ -0,0 +1,234 @@
+"""The result card: what one machine measured, in a shape that can be compared.
+
+A card is the unit a contributor submits. It pairs the measured rows with enough
+hardware and software context to make them meaningful, and with the knobs that
+would otherwise silently change the numbers (thread count, token budget, eval
+window count, corpus hash).
+
+**On privacy.** Cards get committed to a public repository, so the fingerprint is
+deliberately narrow: CPU model string, architecture, core count, rounded RAM, OS
+name and release, and version strings. It never collects hostname, username,
+local paths, MAC address, IP, serial number or any machine ID. ``fingerprint()``
+is the only thing that reads the host, so that promise is auditable in one place.
+"""
+
+from __future__ import annotations
+
+import platform
+import re
+import subprocess
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+#: Bumped when the card layout changes in a way a validator must notice.
+CARD_SCHEMA_VERSION = 2
+
+
+def _cpu_model() -> str:
+    """A human-readable CPU model string, or 'unknown' if the OS will not say."""
+    system = platform.system()
+    try:
+        if system == "Darwin":
+            out = subprocess.run(
+                ["sysctl", "-n", "machdep.cpu.brand_string"],
+                capture_output=True,
+                text=True,
+                timeout=5,
+            )
+            if out.returncode == 0 and out.stdout.strip():
+                return out.stdout.strip()
+        elif system == "Linux":
+            with open("/proc/cpuinfo") as fh:
+                for line in fh:
+                    if line.startswith(("model name", "Model")):
+                        return line.split(":", 1)[1].strip()
+    except (OSError, subprocess.SubprocessError):
+        pass
+    return platform.processor() or platform.machine() or "unknown"
+
+
+def performance_cores() -> tuple[int, str]:
+    """Count the machine's *fast* cores, and say how that was determined.
+
+    This matters far more than it looks. On a heterogeneous CPU — Apple Silicon,
+    ARM big.LITTLE, Intel hybrid — an ONNX Runtime parallel region ends on a
+    barrier, so one thread scheduled onto an efficiency core gates the whole
+    region. Measured on an M4 (4 performance + 6 efficiency), pinning all ten
+    physical cores cost int8 **66% of its throughput** (24.1 vs 40.0 tok/s) and
+    doubled run-to-run spread, while barely moving fp32: the quantized kernel is
+    compute-bound and waits on the slow thread, fp32 is bandwidth-bound and does
+    not. Defaulting to every physical core therefore understates quantized
+    performance and makes results unreproducible.
+
+    Returns ``(count, policy)`` so the card records which rule produced it.
+    """
+    system = platform.system()
+
+    if system == "Darwin":
+        # perflevel0 is the fastest tier; the key is absent on uniform chips.
+        try:
+            out = subprocess.run(
+                ["sysctl", "-n", "hw.perflevel0.physicalcpu"],
+                capture_output=True,
+                text=True,
+                timeout=5,
+            )
+            if out.returncode == 0 and out.stdout.strip().isdigit():
+                count = int(out.stdout.strip())
+                if count > 0:
+                    return count, "macos-perflevel0"
+        except (OSError, subprocess.SubprocessError):
+            pass
+
+    elif system == "Linux":
+        # ARM big.LITTLE exposes a per-cpu capacity; the fast tier is the max.
+        try:
+            caps: dict[int, int] = {}
+            for path in Path("/sys/devices/system/cpu").glob("cpu[0-9]*/cpu_capacity"):
+                try:
+                    caps[int(path.parent.name[3:])] = int(path.read_text().strip())
+                except (OSError, ValueError):
+                    continue
+            if caps:
+                peak = max(caps.values())
+                count = sum(1 for value in caps.values() if value == peak)
+                if 0 < count < len(caps):
+                    return count, "linux-cpu-capacity"
+        except OSError:
+            pass
+
+    try:
+        import psutil
+
+        return (psutil.cpu_count(logical=False) or psutil.cpu_count() or 1), "physical-cores"
+    except Exception:
+        return 1, "fallback"
+
+
+def _total_ram_gb() -> float:
+    """Installed RAM, rounded to 1 GB — precise enough to compare, too coarse to identify."""
+    try:
+        import psutil
+
+        return round(psutil.virtual_memory().total / (1024**3))
+    except Exception:
+        return 0.0
+
+
+def slugify(text: str, *, max_length: int = 48) -> str:
+    """Lowercase, hyphenated, filesystem-safe — used to name the card file."""
+    slug = re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")
+    return slug[:max_length].rstrip("-") or "unknown"
+
+
+@dataclass
+class Fingerprint:
+    """Non-identifying description of the machine and software stack."""
+
+    cpu: str
+    arch: str
+    physical_cores: int
+    logical_cores: int
+    ram_gb: float
+    os: str
+    os_release: str
+    python: str
+    onnxruntime: str
+    provider: str
+    intra_op_threads: int
+    thread_policy: str = "unknown"
+
+    @property
+    def slug(self) -> str:
+        return f"{slugify(self.cpu, max_length=40)}-{slugify(self.os, max_length=10)}"
+
+
+def fingerprint(
+    *, provider: str, intra_op_threads: int, thread_policy: str = "unknown"
+) -> Fingerprint:
+    """Collect the host description. The only function in the project that reads the machine."""
+    import onnxruntime as ort
+    import psutil
+
+    return Fingerprint(
+        cpu=_cpu_model(),
+        arch=platform.machine(),
+        physical_cores=psutil.cpu_count(logical=False) or 0,
+        logical_cores=psutil.cpu_count(logical=True) or 0,
+        ram_gb=_total_ram_gb(),
+        os=platform.system(),
+        os_release=platform.release(),
+        python=platform.python_version(),
+        onnxruntime=ort.__version__,
+        provider=provider,
+        intra_op_threads=intra_op_threads,
+        thread_policy=thread_policy,
+    )
+
+
+@dataclass
+class PrecisionRow:
+    """One measured (model, precision) combination."""
+
+    precision: str
+    size_mb: float
+    latency_s_mean: float
+    latency_s_std: float
+    tokens_per_second: float
+    peak_ram_mb: float
+    perplexity: float | None
+    generated_tokens: int
+    measured_runs: int
+
+    def as_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class ResultCard:
+    """A full submission: one machine, one model, every precision it measured."""
+
+    schema_version: int
+    model_id: str
+    revision: str
+    created_utc: str
+    tool_version: str
+    machine: Fingerprint
+    rows: list[PrecisionRow]
+    eval_windows: int
+    eval_corpus_sha256: str
+    warmup_runs: int
+    gen_tokens: int
+    prompt_sha256: str
+    notes: str = ""
+    submitted_by: str = ""
+    extra: dict[str, Any] = field(default_factory=dict)
+
+    def as_dict(self) -> dict[str, Any]:
+        payload = asdict(self)
+        payload["machine"] = asdict(self.machine)
+        payload["rows"] = [r.as_dict() for r in self.rows]
+        return payload
+
+    @property
+    def filename(self) -> str:
+        """``<cpu>-<os>--<model>.json`` — stable, so a re-run updates its own card."""
+        return f"{self.machine.slug}--{slugify(self.model_id.split('/')[-1], max_length=32)}.json"
+
+    def speedup_vs(self, baseline: str = "fp32") -> dict[str, float]:
+        """Throughput of each precision relative to the baseline precision."""
+        by_precision = {r.precision: r for r in self.rows}
+        base = by_precision.get(baseline)
+        if base is None or base.tokens_per_second <= 0:
+            return {}
+        return {
+            r.precision: round(r.tokens_per_second / base.tokens_per_second, 3)
+            for r in self.rows
+            if r.precision != baseline
+        }
+
+
+def now_utc_iso() -> str:
+    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")