quantcost 0.2.0__py3-none-any.whl

edgellm/quantize.py

@@ -0,0 +1,134 @@
+"""Quantization: PyTorch INT8 and ONNX Runtime INT8 / INT4.
+
+Plain-English summary of what happens here (the README expands on this):
+
+* **INT8 (8-bit)** stores each weight as an 8-bit integer plus a floating-point
+  *scale*, cutting size ~4x vs FP32. **Dynamic** quantization keeps activations
+  in float and quantizes weights only; per-token activation scales are computed
+  on the fly at inference (no calibration data needed) — the robust default for
+  transformers.
+* **INT4 (4-bit)** packs weights into 4 bits using *block-wise* quantization:
+  each contiguous block of ``block_size`` weights shares one scale (and, when
+  asymmetric, a zero-point). Smaller blocks = better accuracy, more overhead.
+
+On CUDA hardware, GPTQ/AWQ give higher-quality INT4; they are unavailable on
+this Apple-Silicon machine (both require CUDA), so the INT4 path here uses ONNX
+Runtime's block-wise weight-only quantizer, which runs on CPU. See
+:meth:`Quantizer.int4_availability` for the honest state of each backend.
+"""
+
+from __future__ import annotations
+
+import logging
+import shutil
+from pathlib import Path
+
+from edgellm.config import QuantizeConfig
+
+logger = logging.getLogger(__name__)
+
+# Non-weight files copied alongside a quantized model.onnx so an ORT runner can
+# load the directory standalone (tokenizer, config, chat template, ...).
+_SIDE_CAR_GLOBS = (
+    "*.json",
+    "*.txt",
+    "*.jinja",
+    "*.model",
+)
+
+
+def _copy_sidecars(src_dir: Path, dst_dir: Path) -> None:
+    dst_dir.mkdir(parents=True, exist_ok=True)
+    for pattern in _SIDE_CAR_GLOBS:
+        for f in src_dir.glob(pattern):
+            if f.is_file():
+                shutil.copy2(f, dst_dir / f.name)
+
+
+class Quantizer:
+    """Produce quantized model artifacts from an exported FP32 ONNX directory."""
+
+    def __init__(self, config: QuantizeConfig) -> None:
+        self.config = config
+
+    # --- ONNX Runtime INT8 (dynamic, weight-only) ---
+    def ort_dynamic_int8(self, fp32_dir: Path, out_dir: Path) -> Path:
+        """Dynamically quantize ONNX weights to INT8. Returns the output dir."""
+        from onnxruntime.quantization import QuantType, quantize_dynamic
+
+        out_dir.mkdir(parents=True, exist_ok=True)
+        model_out = out_dir / "model.onnx"
+        logger.info("ORT dynamic INT8: %s -> %s", fp32_dir, model_out)
+        quantize_dynamic(
+            model_input=str(fp32_dir / "model.onnx"),
+            model_output=str(model_out),
+            weight_type=QuantType.QInt8,
+            per_channel=self.config.int8.per_channel,
+        )
+        _copy_sidecars(fp32_dir, out_dir)
+        return out_dir
+
+    # --- ONNX Runtime INT4 (block-wise, weight-only) ---
+    def ort_int4(self, fp32_dir: Path, out_dir: Path) -> Path:
+        """Block-wise weight-only INT4 quantization (native ORT quantizer).
+
+        The model is passed *by path* so ONNX Runtime streams the >2GB external
+        data instead of requiring it all in a single in-memory ModelProto.
+        """
+        from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer
+
+        out_dir.mkdir(parents=True, exist_ok=True)
+        model_out = out_dir / "model.onnx"
+        logger.info("ORT INT4 (block_size=%d): %s", self.config.int4.group_size, model_out)
+
+        quantizer = MatMulNBitsQuantizer(
+            str(fp32_dir / "model.onnx"),
+            block_size=self.config.int4.group_size,
+            is_symmetric=True,
+        )
+        quantizer.process()
+        quantizer.model.save_model_to_file(str(model_out), use_external_data_format=False)
+        _copy_sidecars(fp32_dir, out_dir)
+        return out_dir
+
+    # --- PyTorch INT8 (dynamic) ---
+    @staticmethod
+    def pytorch_dynamic_int8(model):
+        """Dynamically quantize ``nn.Linear`` layers of a torch model to INT8 (CPU).
+
+        Selects the ``qnnpack`` engine, which is the INT8 kernel backend available
+        on ARM (Apple Silicon); x86 would use ``fbgemm``.
+        """
+        import torch
+        from torch.ao.quantization import quantize_dynamic
+
+        engines = torch.backends.quantized.supported_engines
+        if "qnnpack" in engines:
+            torch.backends.quantized.engine = "qnnpack"
+        elif "fbgemm" in engines:
+            torch.backends.quantized.engine = "fbgemm"
+
+        model = model.to("cpu").eval()
+        return quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
+
+    @staticmethod
+    def int4_availability() -> dict[str, str]:
+        """Report which INT4 backends can actually run on this machine."""
+        import importlib.util
+
+        import torch
+
+        cuda = torch.cuda.is_available()
+        return {
+            "gptq (auto-gptq)": (
+                "available"
+                if (cuda and importlib.util.find_spec("auto_gptq"))
+                else "unavailable (needs CUDA)"
+            ),
+            "awq (autoawq)": (
+                "available"
+                if (cuda and importlib.util.find_spec("awq"))
+                else "unavailable (needs CUDA)"
+            ),
+            "onnxruntime block-wise INT4": "available (CPU)",
+        }

edgellm/render.py

@@ -0,0 +1,164 @@
+"""Turn a result card into the thing a person actually wanted: a tradeoff table.
+
+Raw per-precision numbers are not the answer to "should I quantize this?". The
+answer is *relative*: how much faster, how much smaller, how much worse. So every
+view here is anchored to a baseline precision and reports deltas, with the
+absolute numbers kept alongside so nothing is hidden.
+
+The verdict line exists because the honest answer is frequently "don't" — INT4 on
+a CPU without native 4-bit kernels routinely dequantizes to fp32 inside the GEMM
+and comes out slower than the baseline it was supposed to beat. A tool that
+prints numbers but will not say that is making the reader do the work.
+"""
+
+from __future__ import annotations
+
+from edgellm.card import PrecisionRow, ResultCard
+
+BASELINE = "fp32"
+
+
+def _fmt(value: float | None, spec: str = ".1f", dash: str = "—") -> str:
+    return dash if value is None else format(value, spec)
+
+
+def _signed(value: float, spec: str = ".1f") -> str:
+    """Format with an explicit sign, so a quality *improvement* reads as negative."""
+    return format(value, f"+{spec}")
+
+
+def tradeoff_table(card: ResultCard, *, baseline: str = BASELINE) -> str:
+    """A Markdown table of absolute numbers plus deltas against the baseline."""
+    by_precision = {r.precision: r for r in card.rows}
+    base = by_precision.get(baseline)
+
+    header = (
+        "| Precision | Size (MB) | Throughput (tok/s) | Speedup | "
+        "Peak RAM (MB) | Perplexity | PPL change |\n"
+        "| --- | --- | --- | --- | --- | --- | --- |\n"
+    )
+    lines = []
+    for row in card.rows:
+        speedup = "baseline" if row.precision == baseline else "—"
+        ppl_delta = "baseline" if row.precision == baseline else "—"
+
+        if base and row.precision != baseline:
+            if base.tokens_per_second > 0:
+                ratio = row.tokens_per_second / base.tokens_per_second
+                speedup = f"{ratio:.2f}x"
+            if base.perplexity and row.perplexity:
+                pct = (row.perplexity / base.perplexity - 1) * 100
+                ppl_delta = f"{_signed(pct)}%"
+
+        lines.append(
+            f"| {row.precision} | {_fmt(row.size_mb, '.0f')} | "
+            f"{_fmt(row.tokens_per_second, '.2f')} | {speedup} | "
+            f"{_fmt(row.peak_ram_mb, '.0f')} | {_fmt(row.perplexity, '.2f')} | {ppl_delta} |"
+        )
+    return header + "\n".join(lines) + "\n"
+
+
+def unstable_rows(card: ResultCard, *, ratio: float = 0.15) -> list[str]:
+    """Precisions whose run-to-run spread is too wide for the number to be trusted."""
+    return [
+        row.precision
+        for row in card.rows
+        if row.latency_s_mean > 0 and row.latency_s_std / row.latency_s_mean > ratio
+    ]
+
+
+def verdict(card: ResultCard, *, baseline: str = BASELINE) -> list[str]:
+    """Plain-language findings. Says 'not worth it' when the numbers say so."""
+    by_precision = {r.precision: r for r in card.rows}
+    base = by_precision.get(baseline)
+    if base is None:
+        return [f"No {baseline} baseline was measured, so nothing can be compared against it."]
+
+    findings: list[str] = []
+    candidates = [r for r in card.rows if r.precision != baseline]
+
+    def shrink(row: PrecisionRow) -> float:
+        return base.size_mb / row.size_mb if row.size_mb else 0.0
+
+    for row in candidates:
+        ratio = row.tokens_per_second / base.tokens_per_second if base.tokens_per_second else 0.0
+        size_x = shrink(row)
+        bits = [f"{size_x:.1f}x smaller"] if size_x > 1.05 else []
+
+        if ratio >= 1.15:
+            bits.append(f"{ratio:.2f}x faster")
+        elif ratio <= 0.9:
+            bits.append(f"**{1 / ratio:.2f}x slower**")
+        else:
+            bits.append("about the same speed")
+
+        if base.perplexity and row.perplexity:
+            pct = (row.perplexity / base.perplexity - 1) * 100
+            if pct > 10:
+                bits.append(f"and {pct:.0f}% worse perplexity")
+            elif pct > 2:
+                bits.append(f"for {pct:.0f}% worse perplexity")
+            else:
+                bits.append("at essentially no quality cost")
+
+        findings.append(f"- **{row.precision}**: " + ", ".join(bits) + ".")
+
+    base_tps = base.tokens_per_second
+    faster = [r for r in candidates if base_tps and r.tokens_per_second > base_tps * 1.15]
+    slower = [r for r in candidates if base_tps and r.tokens_per_second < base_tps * 0.9]
+
+    if slower:
+        names = ", ".join(r.precision for r in slower)
+        verb = "is" if len(slower) == 1 else "are"
+        findings.append(
+            f"- On this machine **{names} {verb} slower than {baseline}** despite being smaller. "
+            "This is a property of *autoregressive decode*, not of quantized arithmetic: "
+            "ONNX Runtime unpacks the weights back to float inside each step, and a single "
+            "token cannot amortise unpacking a whole weight matrix. With a long prompt to "
+            "amortise over, the same artifacts often come out faster than the baseline. "
+            "What that means in practice: quantize here to fit in memory, not to speed up "
+            "generation."
+        )
+    if faster:
+        best = max(faster, key=lambda r: r.tokens_per_second)
+        findings.append(
+            f"- Best throughput on this machine is **{best.precision}** at "
+            f"{best.tokens_per_second:.1f} tok/s."
+        )
+    return findings
+
+
+def console_report(card: ResultCard, errors: dict[str, str] | None = None) -> str:
+    """The full human-facing report printed after a run."""
+    machine = card.machine
+    measured = card.rows[0].measured_runs if card.rows else 0
+    out = [
+        "",
+        f"  model    {card.model_id} @ {card.revision}",
+        f"  machine  {machine.cpu} ({machine.arch}), {machine.physical_cores} physical cores, "
+        f"{machine.ram_gb:.0f} GB RAM",
+        f"  software {machine.os} {machine.os_release} · Python {machine.python} · "
+        f"onnxruntime {machine.onnxruntime} · {machine.provider}",
+        f"  settings {machine.intra_op_threads} threads ({machine.thread_policy}) · "
+        f"{card.gen_tokens} generated tokens · "
+        f"{card.warmup_runs} warmup + {measured} measured runs"
+        + (f" · {card.eval_windows} perplexity windows" if card.eval_windows else ""),
+        "",
+        tradeoff_table(card),
+    ]
+    findings = verdict(card)
+    if findings:
+        out += ["What this means on this machine:", ""] + findings + [""]
+
+    if unstable := unstable_rows(card):
+        out += [
+            "Measurement warning:",
+            "",
+            f"- {', '.join(unstable)} varied by more than 15% between runs, so these "
+            "numbers are noisy. Close other applications and re-run before submitting — "
+            "a busy machine makes the comparison between precisions meaningless.",
+            "",
+        ]
+    if errors:
+        out += ["Skipped:"] + [f"- {p}: {why}" for p, why in errors.items()] + [""]
+    return "\n".join(out)

edgellm/report.py

@@ -0,0 +1,53 @@
+"""Turn the benchmark JSON into charts (and, in Phase 6, the full report)."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+
+def _label(row: dict) -> str:
+    return f"{row['backend']}\n{row['precision']}/{row['device']}"
+
+
+def render_charts(results_json: Path, out_path: Path) -> Path:
+    """Render size / throughput / perplexity bar charts to a single PNG."""
+    import matplotlib
+
+    matplotlib.use("Agg")  # headless: no display needed
+    import matplotlib.pyplot as plt
+
+    rows = json.loads(results_json.read_text())
+    labels = [_label(r) for r in rows]
+
+    panels = [
+        ("On-disk size (MB)", [r.get("size_mb") for r in rows], "#4C72B0"),
+        ("Throughput (tok/s)", [r.get("tokens_per_second") for r in rows], "#55A868"),
+        ("Perplexity (lower=better)", [r.get("perplexity") for r in rows], "#C44E52"),
+    ]
+
+    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+    for ax, (title, values, color) in zip(axes, panels, strict=False):
+        xs = list(range(len(labels)))
+        vals = [v if v is not None else 0 for v in values]
+        bars = ax.bar(xs, vals, color=color)
+        ax.set_title(title)
+        ax.set_xticks(xs)
+        ax.set_xticklabels(labels, fontsize=8)
+        for bar, original in zip(bars, values, strict=False):
+            text = "n/a" if original is None else f"{original:g}"
+            ax.text(
+                bar.get_x() + bar.get_width() / 2,
+                bar.get_height(),
+                text,
+                ha="center",
+                va="bottom",
+                fontsize=8,
+            )
+
+    fig.suptitle("EdgeLLM benchmark (real runs)", fontsize=13)
+    fig.tight_layout()
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(out_path, dpi=120)
+    plt.close(fig)
+    return out_path

edgellm/runners.py

@@ -0,0 +1,170 @@
+"""Inference backends.
+
+Every backend implements the same :class:`InferenceRunner` interface so the
+benchmark harness can treat PyTorch, ONNX Runtime (CPU/GPU) and the Qualcomm QNN
+NPU path identically. Phase 1 ships the PyTorch backend; later phases add the
+ONNX Runtime and QNN runners.
+"""
+
+from __future__ import annotations
+
+import time
+
+import torch
+
+from edgellm.base import GenerationResult, InferenceRunner
+from edgellm.config import GenerationConfig
+from edgellm.models import LoadedModel
+
+__all__ = [
+    "GenerationResult",
+    "InferenceRunner",
+    "PyTorchRunner",
+    "ORTRunner",
+    "QNNRunner",
+    "encode_prompt",
+]
+
+
+def encode_prompt(tokenizer, prompt: str) -> dict[str, torch.Tensor]:
+    """Turn a prompt into model inputs (CPU tensors, not yet placed on a device).
+
+    Uses the chat template for instruct models; falls back to plain tokenization
+    for base models. Shared by every backend so all runners see identical inputs.
+    """
+    if tokenizer.chat_template:
+        messages = [{"role": "user", "content": prompt}]
+        return tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            return_tensors="pt",
+            return_dict=True,
+        )
+    return dict(tokenizer(prompt, return_tensors="pt"))
+
+
+class PyTorchRunner(InferenceRunner):
+    """Eager PyTorch generation — the correctness/quality reference backend."""
+
+    name = "pytorch"
+
+    def __init__(self, loaded: LoadedModel) -> None:
+        self.loaded = loaded
+
+    def _build_inputs(self, prompt: str) -> dict[str, torch.Tensor]:
+        """Encode ``prompt`` and place the tensors on the model's device."""
+        inputs = encode_prompt(self.loaded.tokenizer, prompt)
+        return {k: v.to(self.loaded.device) for k, v in inputs.items()}
+
+    @torch.inference_mode()
+    def generate(self, prompt: str, generation: GenerationConfig) -> GenerationResult:
+        torch.manual_seed(generation.seed)
+        tokenizer = self.loaded.tokenizer
+        inputs = self._build_inputs(prompt)
+        prompt_tokens = int(inputs["input_ids"].shape[-1])
+
+        start = time.perf_counter()
+        output_ids = self.loaded.model.generate(
+            **inputs,
+            max_new_tokens=generation.max_new_tokens,
+            min_new_tokens=generation.min_new_tokens,
+            do_sample=generation.do_sample,
+            temperature=generation.temperature,
+            top_p=generation.top_p,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+        if self.loaded.device == "mps":
+            torch.mps.synchronize()
+        elif self.loaded.device == "cuda":
+            torch.cuda.synchronize()
+        latency_s = time.perf_counter() - start
+
+        new_ids = output_ids[0, prompt_tokens:]
+        generated_tokens = int(new_ids.shape[-1])
+        text = tokenizer.decode(new_ids, skip_special_tokens=True)
+        tokens_per_second = generated_tokens / latency_s if latency_s > 0 else 0.0
+
+        return GenerationResult(
+            backend=self.name,
+            prompt=prompt,
+            text=text,
+            prompt_tokens=prompt_tokens,
+            generated_tokens=generated_tokens,
+            latency_s=latency_s,
+            tokens_per_second=tokens_per_second,
+        )
+
+
+class ORTRunner(InferenceRunner):
+    """ONNX Runtime backend (loads an exported/quantized ONNX model directory).
+
+    Used for the FP32 ONNX baseline (Phase 2) and every quantized ONNX model
+    (Phase 3). The execution provider selects the hardware: ``CPUExecutionProvider``
+    for CPU, ``CoreMLExecutionProvider`` on Apple, ``QNNExecutionProvider`` for the
+    Qualcomm NPU (Phase 5).
+    """
+
+    def __init__(
+        self,
+        model_dir: str,
+        tokenizer,
+        provider: str = "CPUExecutionProvider",
+        name: str = "ort-cpu",
+    ) -> None:
+        import onnxruntime as ort
+        from optimum.onnxruntime import ORTModelForCausalLM
+
+        available = ort.get_available_providers()
+        if provider not in available:
+            raise RuntimeError(
+                f"Execution provider '{provider}' is not available in this ONNX "
+                f"Runtime build (have: {available}). For the Qualcomm NPU, install "
+                "onnxruntime-qnn on a Snapdragon host, or use aihub/run_on_snapdragon.py."
+            )
+        self.name = name
+        self.tokenizer = tokenizer
+        self.model = ORTModelForCausalLM.from_pretrained(model_dir, provider=provider)
+
+
+class QNNRunner(ORTRunner):
+    """ONNX Runtime QNN Execution Provider runner (Qualcomm NPU).
+
+    Requires an ``onnxruntime-qnn`` build on a Snapdragon host; on other platforms
+    construction raises a clear error. For laptop-side profiling of real Snapdragon
+    hardware without a device in hand, use ``aihub/run_on_snapdragon.py`` instead.
+    """
+
+    def __init__(self, model_dir: str, tokenizer, name: str = "ort-qnn") -> None:
+        super().__init__(model_dir, tokenizer, provider="QNNExecutionProvider", name=name)
+
+    def generate(self, prompt: str, generation: GenerationConfig) -> GenerationResult:
+        torch.manual_seed(generation.seed)
+        inputs = encode_prompt(self.tokenizer, prompt)
+        prompt_tokens = int(inputs["input_ids"].shape[-1])
+
+        start = time.perf_counter()
+        output_ids = self.model.generate(
+            **inputs,
+            max_new_tokens=generation.max_new_tokens,
+            min_new_tokens=generation.min_new_tokens,
+            do_sample=generation.do_sample,
+            temperature=generation.temperature,
+            top_p=generation.top_p,
+            pad_token_id=self.tokenizer.pad_token_id,
+        )
+        latency_s = time.perf_counter() - start
+
+        new_ids = output_ids[0, prompt_tokens:]
+        generated_tokens = int(new_ids.shape[-1])
+        text = self.tokenizer.decode(new_ids, skip_special_tokens=True)
+        tokens_per_second = generated_tokens / latency_s if latency_s > 0 else 0.0
+
+        return GenerationResult(
+            backend=self.name,
+            prompt=prompt,
+            text=text,
+            prompt_tokens=prompt_tokens,
+            generated_tokens=generated_tokens,
+            latency_s=latency_s,
+            tokens_per_second=tokens_per_second,
+        )