From c221281ccd4077ad6d066d92ad4676216edde063 Mon Sep 17 00:00:00 2001 From: nathanhubens Date: Mon, 6 Jul 2026 13:26:37 +0200 Subject: [PATCH 1/2] fix(benchmark): don't segfault benchmarking quantized models on CUDA Quantized models (INT8 fbgemm/qnnpack) have CPU-only kernels. benchmark() defaults to devices [cpu, cuda] on a CUDA host; moving a quantized model to CUDA dispatches e.g. quantized::conv2d_relu on a QuantizedCUDA tensor with no registered kernel -> PyTorch's dispatcher SEGFAULTS (uncatchable). Hits any workflow that INT8-quantizes a model and then benchmarks it. Fix: - _is_quantized() + _device_supported()/_ensure_device_supported() in core. - _run_on_devices skips unsupported devices (CUDA for quantized) -> NaN + warning: one choke point for the speed/memory/energy multi-device profilers. - Belt-and-suspenders guard at each single-device forward entry raises a catchable RuntimeError instead of segfaulting. Non-quantized models unaffected (ResNet-18 still benches cpu+cuda). Regression tests added; nbdev-test green on all 5 touched notebooks. --- fasterbench/__init__.py | 4 +- fasterbench/_modidx.py | 4 + fasterbench/core.py | 45 +++++++++ fasterbench/energy.py | 3 +- fasterbench/memory.py | 3 +- fasterbench/speed.py | 3 +- nbs/analysis/benchmark.ipynb | 8 ++ nbs/core/core.ipynb | 4 +- nbs/metrics/energy.ipynb | 127 +---------------------- nbs/metrics/memory.ipynb | 4 +- nbs/metrics/speed.ipynb | 189 +---------------------------------- 11 files changed, 74 insertions(+), 320 deletions(-) diff --git a/fasterbench/__init__.py b/fasterbench/__init__.py index ec82712..560dcee 100644 --- a/fasterbench/__init__.py +++ b/fasterbench/__init__.py @@ -1,4 +1,6 @@ -"""Comprehensive benchmarking toolkit for deep learning models""" +"""Comprehensive benchmarking toolkit for deep learning models + +Docs: https://FasterAI-Labs.github.io/fasterbench/index.html.md""" # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/index.ipynb. diff --git a/fasterbench/_modidx.py b/fasterbench/_modidx.py index b5647f9..f7ce56f 100644 --- a/fasterbench/_modidx.py +++ b/fasterbench/_modidx.py @@ -40,10 +40,14 @@ 'fasterbench.core': { 'fasterbench.core._bytes_to_mib': ('core/core.html#_bytes_to_mib', 'fasterbench/core.py'), 'fasterbench.core._default_devices': ('core/core.html#_default_devices', 'fasterbench/core.py'), 'fasterbench.core._device_ctx': ('core/core.html#_device_ctx', 'fasterbench/core.py'), + 'fasterbench.core._device_supported': ('core/core.html#_device_supported', 'fasterbench/core.py'), + 'fasterbench.core._ensure_device_supported': ( 'core/core.html#_ensure_device_supported', + 'fasterbench/core.py'), 'fasterbench.core._fmt_float': ('core/core.html#_fmt_float', 'fasterbench/core.py'), 'fasterbench.core._fmt_human': ('core/core.html#_fmt_human', 'fasterbench/core.py'), 'fasterbench.core._fmt_macs': ('core/core.html#_fmt_macs', 'fasterbench/core.py'), 'fasterbench.core._fmt_table': ('core/core.html#_fmt_table', 'fasterbench/core.py'), + 'fasterbench.core._is_quantized': ('core/core.html#_is_quantized', 'fasterbench/core.py'), 'fasterbench.core._run_on_devices': ('core/core.html#_run_on_devices', 'fasterbench/core.py'), 'fasterbench.core._section': ('core/core.html#_section', 'fasterbench/core.py'), 'fasterbench.core._sync': ('core/core.html#_sync', 'fasterbench/core.py')}, diff --git a/fasterbench/core.py b/fasterbench/core.py index cb6af19..7975c16 100644 --- a/fasterbench/core.py +++ b/fasterbench/core.py @@ -66,6 +66,41 @@ def _default_devices() -> list[str]: return devices +def _is_quantized(model: torch.nn.Module) -> bool: # model to inspect + """True if the model contains PyTorch quantized modules. + + PyTorch's eager/FX quantized ops (fbgemm/qnnpack) are **CPU-only**. Running + them on a CUDA tensor has no registered kernel and *segfaults* the process + rather than raising, so callers must avoid dispatching them off-CPU. + """ + return any("quantized" in type(m).__module__ for m in model.modules()) + + +def _device_supported( + model: torch.nn.Module, # model to run + device: str | torch.device, # target device +) -> bool: + """False if `model` cannot execute on `device` (quantized models are CPU-only).""" + return not (torch.device(device).type != "cpu" and _is_quantized(model)) + + +def _ensure_device_supported( + model: torch.nn.Module, # model to run + device: str | torch.device, # target device +) -> None: + """Raise a *catchable* error instead of letting PyTorch segfault. + + Guards the forward-executing profilers: a quantized model dispatched on a + non-CPU device would crash the interpreter with SIGSEGV, so we fail loudly + (and recoverably) first. + """ + if not _device_supported(model, device): + raise RuntimeError( + f"Quantized models run on CPU only; cannot benchmark on '{device}'. " + "PyTorch quantized ops (fbgemm/qnnpack) have no CUDA kernels." + ) + + def _run_on_devices( compute_fn, # single-device compute function model: torch.nn.Module, # model to benchmark @@ -78,6 +113,9 @@ def _run_on_devices( """Run compute_fn on multiple devices with unified error handling. Returns dict mapping device string to metrics (or NaN metrics on failure). + Devices a model cannot run on (e.g. CUDA for a quantized model) are skipped + with NaN metrics — this avoids an uncatchable SIGSEGV from dispatching + CPU-only quantized ops on a CUDA tensor. """ if devices is None: devices = _default_devices() @@ -85,6 +123,13 @@ def _run_on_devices( out = {} for d in devices: d_str = str(d) + if not _device_supported(model, d): + warnings.warn( + f"{metric_name} benchmark skipped on {d}: quantized models run on " + "CPU only (PyTorch quantized ops have no CUDA kernels)." + ) + out[d_str] = nan_factory(d_str) + continue try: out[d_str] = compute_fn(model, sample, device=d, **kwargs) except (RuntimeError, torch.cuda.OutOfMemoryError) as e: diff --git a/fasterbench/energy.py b/fasterbench/energy.py index e4c5b80..9c5f425 100644 --- a/fasterbench/energy.py +++ b/fasterbench/energy.py @@ -8,7 +8,7 @@ import torch -from .core import _device_ctx, _sync, _run_on_devices +from .core import _device_ctx, _sync, _run_on_devices, _ensure_device_supported try: from codecarbon import EmissionsTracker, OfflineEmissionsTracker @@ -94,6 +94,7 @@ def compute_energy( ) with _device_ctx(device) as dev: + _ensure_device_supported(model, dev) # quantized ops are CPU-only; avoid SIGSEGV model = model.eval().to(dev) sample = sample.to(dev, non_blocking=True) diff --git a/fasterbench/memory.py b/fasterbench/memory.py index edfa4d3..70fb76f 100644 --- a/fasterbench/memory.py +++ b/fasterbench/memory.py @@ -6,7 +6,7 @@ import warnings from dataclasses import dataclass, asdict from typing import Sequence -from .core import _bytes_to_mib, _run_on_devices +from .core import _bytes_to_mib, _run_on_devices, _ensure_device_supported import numpy as np import torch @@ -49,6 +49,7 @@ def _gpu_metrics( ) -> MemoryMetrics: """Measure GPU memory usage.""" dev = torch.device("cuda") + _ensure_device_supported(model, dev) # quantized ops are CPU-only; avoid SIGSEGV model = model.eval().to(dev) sample = sample.to(dev) diff --git a/fasterbench/speed.py b/fasterbench/speed.py index 0b8625b..9af3d31 100644 --- a/fasterbench/speed.py +++ b/fasterbench/speed.py @@ -11,7 +11,7 @@ import torch.nn as nn from torch.utils.benchmark import Timer -from .core import _device_ctx, _sync, _run_on_devices +from .core import _device_ctx, _sync, _run_on_devices, _ensure_device_supported # %% auto #0 __all__ = ['SpeedMetrics', 'compute_speed', 'compute_speed_multi', 'sweep_threads', 'sweep_batch_sizes', 'sweep_latency'] @@ -68,6 +68,7 @@ def _forward_latencies( use_torch_timer = torch.device(device).type == "cpu" with _device_ctx(device) as dev: + _ensure_device_supported(model, dev) # quantized ops are CPU-only; avoid SIGSEGV model.eval().to(dev) sample = sample.to(dev, non_blocking=True) diff --git a/nbs/analysis/benchmark.ipynb b/nbs/analysis/benchmark.ipynb index a354ea6..a0434eb 100644 --- a/nbs/analysis/benchmark.ipynb +++ b/nbs/analysis/benchmark.ipynb @@ -143,6 +143,14 @@ "#| hide\nfrom fastcore.test import *\n\nimport torch, torch.nn as nn\n_m = nn.Linear(10, 5)\n_x = torch.randn(1, 10)\n_r = benchmark(_m, _x, metrics=[\"size\"])\nassert isinstance(_r, BenchmarkResult)\nassert \"size_num_params\" in _r" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "075b7b5f", + "metadata": {}, + "outputs": [], + "source": "#| hide\n# Regression: quantized models are CPU-only. PyTorch's quantized ops (fbgemm/qnnpack)\n# have no CUDA kernels, so dispatching them on a CUDA tensor *segfaults* the process.\n# The forward-executing profilers (speed/memory/energy) default to cpu+cuda, so\n# benchmarking a quantized model on a CUDA host must NOT crash — the CUDA leg is\n# skipped (NaN), the CPU leg is real. See: fasterbench/core._is_quantized.\nimport math\nimport torch.ao.quantization as tq\nfrom fasterbench.core import _is_quantized, _ensure_device_supported\n\nclass _QNet(nn.Module):\n \"Tiny static-quantizable net → QuantizedConvReLU2d + QuantizedLinear (pure torch).\"\n def __init__(self):\n super().__init__()\n self.quant, self.dequant = tq.QuantStub(), tq.DeQuantStub()\n self.conv, self.relu = nn.Conv2d(3, 4, 3, padding=1), nn.ReLU()\n self.pool, self.flat, self.fc = nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(4, 2)\n def forward(self, x):\n x = self.quant(x); x = self.relu(self.conv(x))\n x = self.fc(self.flat(self.pool(x))); return self.dequant(x)\n\n_qm = _QNet().eval()\n_qm.qconfig = tq.get_default_qconfig(\"fbgemm\")\ntq.fuse_modules(_qm, [[\"conv\", \"relu\"]], inplace=True)\ntq.prepare(_qm, inplace=True)\n_qm(torch.randn(4, 3, 8, 8)) # calibrate\ntq.convert(_qm, inplace=True) # -> quantized (CPU-only) modules\nassert _is_quantized(_qm)\n\n_qx = torch.randn(1, 3, 8, 8)\n# Must complete WITHOUT SIGSEGV, even on a CUDA machine.\n_qr = benchmark(_qm, _qx, metrics=[\"size\", \"speed\", \"compute\", \"memory\"], warmup=2, steps=3)\nassert isinstance(_qr, BenchmarkResult)\nassert math.isfinite(_qr.speed[\"cpu\"].mean_ms) # CPU result is real\nassert all(math.isnan(m.mean_ms) for d, m in _qr.speed.items() if d != \"cpu\") # CUDA skipped -> NaN\n\n# Low-level guard: a direct quantized-on-CUDA call raises (catchable), never segfaults.\nwith ExceptionExpected(RuntimeError):\n _ensure_device_supported(_qm, \"cuda\")" + }, { "cell_type": "markdown", "execution_count": null, diff --git a/nbs/core/core.ipynb b/nbs/core/core.ipynb index 5840c12..c53a5e9 100644 --- a/nbs/core/core.ipynb +++ b/nbs/core/core.ipynb @@ -48,7 +48,7 @@ "id": "ed4d37d2", "metadata": {}, "outputs": [], - "source": "#| export\ndef _bytes_to_mib(n: int) -> float: # number of bytes to convert\n \"\"\"Convert bytes to MiB (base-2).\"\"\"\n return n / 1024 ** 2\n\ndef _fmt_human(n: int | float) -> str: # number to format\n \"\"\"Format large number with K/M/B suffix for human-readable output.\"\"\"\n if n >= 1e9: return f\"{n/1e9:.2f}B\"\n if n >= 1e6: return f\"{n/1e6:.2f}M\"\n if n >= 1e3: return f\"{n/1e3:.2f}K\"\n return str(int(n))\n\n\ndef _fmt_table(n: int | float) -> str: # number to format\n \"\"\"Format number with commas for tabular output.\"\"\"\n return f\"{int(n):>12,}\"\n\n\ndef _fmt_float(v: float, width: int = 8, decimals: int = 3) -> str: # value, width, decimal places\n \"\"\"Format float with specified width and decimal places.\"\"\"\n return f\"{v:{width}.{decimals}f}\"\n\n\ndef _fmt_macs(v: float) -> str: # MACs value to format\n \"\"\"Format MACs with appropriate M/K suffix.\"\"\"\n if v >= 1e6: return f\"{v/1e6:8.2f} M\"\n if v >= 1e3: return f\"{v/1e3:8.2f} K\"\n return f\"{int(v):>10}\"\n\n\ndef _section(title: str, width: int = 40) -> str: # section title, total width\n \"\"\"Create a section header with box-drawing characters.\"\"\"\n return f\"═══ {title} \" + \"═\" * (width - len(title))\n\n@contextlib.contextmanager\ndef _device_ctx(dev: str | torch.device): # device string or torch.device\n \"\"\"Context manager that validates device availability and yields resolved device.\"\"\"\n dev = torch.device(dev)\n if dev.type == \"cuda\" and not torch.cuda.is_available():\n warnings.warn(\"CUDA requested but not available – falling back to CPU\")\n dev = torch.device(\"cpu\")\n yield dev\n\ndef _sync(dev: torch.device) -> None: # device to synchronize\n \"\"\"Synchronize CUDA device if applicable.\"\"\"\n if dev.type == \"cuda\":\n torch.cuda.synchronize(dev)\n\n\ndef _default_devices() -> list[str]:\n \"\"\"Return default device list: ['cpu'] + ['cuda'] if available.\"\"\"\n devices = [\"cpu\"]\n if torch.cuda.is_available():\n devices.append(\"cuda\")\n return devices\n\n\ndef _run_on_devices(\n compute_fn, # single-device compute function\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor\n devices: list[str | torch.device] | None, # devices to run on (None = default)\n nan_factory, # callable(device_str) -> metrics with NaN values\n metric_name: str, # for warning messages (e.g., \"Speed\", \"Memory\")\n **kwargs,\n) -> dict:\n \"\"\"Run compute_fn on multiple devices with unified error handling.\n \n Returns dict mapping device string to metrics (or NaN metrics on failure).\n \"\"\"\n if devices is None:\n devices = _default_devices()\n\n out = {}\n for d in devices:\n d_str = str(d)\n try:\n out[d_str] = compute_fn(model, sample, device=d, **kwargs)\n except (RuntimeError, torch.cuda.OutOfMemoryError) as e:\n warnings.warn(f\"{metric_name} benchmark failed on {d}: {e}\")\n out[d_str] = nan_factory(d_str)\n except Exception as e:\n warnings.warn(f\"Unexpected error during {metric_name} benchmark on {d}: {e}\")\n out[d_str] = nan_factory(d_str)\n return out" + "source": "#| export\ndef _bytes_to_mib(n: int) -> float: # number of bytes to convert\n \"\"\"Convert bytes to MiB (base-2).\"\"\"\n return n / 1024 ** 2\n\ndef _fmt_human(n: int | float) -> str: # number to format\n \"\"\"Format large number with K/M/B suffix for human-readable output.\"\"\"\n if n >= 1e9: return f\"{n/1e9:.2f}B\"\n if n >= 1e6: return f\"{n/1e6:.2f}M\"\n if n >= 1e3: return f\"{n/1e3:.2f}K\"\n return str(int(n))\n\n\ndef _fmt_table(n: int | float) -> str: # number to format\n \"\"\"Format number with commas for tabular output.\"\"\"\n return f\"{int(n):>12,}\"\n\n\ndef _fmt_float(v: float, width: int = 8, decimals: int = 3) -> str: # value, width, decimal places\n \"\"\"Format float with specified width and decimal places.\"\"\"\n return f\"{v:{width}.{decimals}f}\"\n\n\ndef _fmt_macs(v: float) -> str: # MACs value to format\n \"\"\"Format MACs with appropriate M/K suffix.\"\"\"\n if v >= 1e6: return f\"{v/1e6:8.2f} M\"\n if v >= 1e3: return f\"{v/1e3:8.2f} K\"\n return f\"{int(v):>10}\"\n\n\ndef _section(title: str, width: int = 40) -> str: # section title, total width\n \"\"\"Create a section header with box-drawing characters.\"\"\"\n return f\"═══ {title} \" + \"═\" * (width - len(title))\n\n@contextlib.contextmanager\ndef _device_ctx(dev: str | torch.device): # device string or torch.device\n \"\"\"Context manager that validates device availability and yields resolved device.\"\"\"\n dev = torch.device(dev)\n if dev.type == \"cuda\" and not torch.cuda.is_available():\n warnings.warn(\"CUDA requested but not available – falling back to CPU\")\n dev = torch.device(\"cpu\")\n yield dev\n\ndef _sync(dev: torch.device) -> None: # device to synchronize\n \"\"\"Synchronize CUDA device if applicable.\"\"\"\n if dev.type == \"cuda\":\n torch.cuda.synchronize(dev)\n\n\ndef _default_devices() -> list[str]:\n \"\"\"Return default device list: ['cpu'] + ['cuda'] if available.\"\"\"\n devices = [\"cpu\"]\n if torch.cuda.is_available():\n devices.append(\"cuda\")\n return devices\n\n\ndef _is_quantized(model: torch.nn.Module) -> bool: # model to inspect\n \"\"\"True if the model contains PyTorch quantized modules.\n\n PyTorch's eager/FX quantized ops (fbgemm/qnnpack) are **CPU-only**. Running\n them on a CUDA tensor has no registered kernel and *segfaults* the process\n rather than raising, so callers must avoid dispatching them off-CPU.\n \"\"\"\n return any(\"quantized\" in type(m).__module__ for m in model.modules())\n\n\ndef _device_supported(\n model: torch.nn.Module, # model to run\n device: str | torch.device, # target device\n) -> bool:\n \"\"\"False if `model` cannot execute on `device` (quantized models are CPU-only).\"\"\"\n return not (torch.device(device).type != \"cpu\" and _is_quantized(model))\n\n\ndef _ensure_device_supported(\n model: torch.nn.Module, # model to run\n device: str | torch.device, # target device\n) -> None:\n \"\"\"Raise a *catchable* error instead of letting PyTorch segfault.\n\n Guards the forward-executing profilers: a quantized model dispatched on a\n non-CPU device would crash the interpreter with SIGSEGV, so we fail loudly\n (and recoverably) first.\n \"\"\"\n if not _device_supported(model, device):\n raise RuntimeError(\n f\"Quantized models run on CPU only; cannot benchmark on '{device}'. \"\n \"PyTorch quantized ops (fbgemm/qnnpack) have no CUDA kernels.\"\n )\n\n\ndef _run_on_devices(\n compute_fn, # single-device compute function\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor\n devices: list[str | torch.device] | None, # devices to run on (None = default)\n nan_factory, # callable(device_str) -> metrics with NaN values\n metric_name: str, # for warning messages (e.g., \"Speed\", \"Memory\")\n **kwargs,\n) -> dict:\n \"\"\"Run compute_fn on multiple devices with unified error handling.\n \n Returns dict mapping device string to metrics (or NaN metrics on failure).\n Devices a model cannot run on (e.g. CUDA for a quantized model) are skipped\n with NaN metrics — this avoids an uncatchable SIGSEGV from dispatching\n CPU-only quantized ops on a CUDA tensor.\n \"\"\"\n if devices is None:\n devices = _default_devices()\n\n out = {}\n for d in devices:\n d_str = str(d)\n if not _device_supported(model, d):\n warnings.warn(\n f\"{metric_name} benchmark skipped on {d}: quantized models run on \"\n \"CPU only (PyTorch quantized ops have no CUDA kernels).\"\n )\n out[d_str] = nan_factory(d_str)\n continue\n try:\n out[d_str] = compute_fn(model, sample, device=d, **kwargs)\n except (RuntimeError, torch.cuda.OutOfMemoryError) as e:\n warnings.warn(f\"{metric_name} benchmark failed on {d}: {e}\")\n out[d_str] = nan_factory(d_str)\n except Exception as e:\n warnings.warn(f\"Unexpected error during {metric_name} benchmark on {d}: {e}\")\n out[d_str] = nan_factory(d_str)\n return out" }, { "cell_type": "code", @@ -56,7 +56,7 @@ "id": "1d0944ae", "metadata": {}, "outputs": [], - "source": [] + "source": "#| hide\nfrom fastcore.test import *\nimport torch.nn as nn\n\n# _is_quantized: plain fp modules are not quantized\n_fp = nn.Sequential(nn.Conv2d(3, 4, 3), nn.ReLU())\ntest_eq(_is_quantized(_fp), False)\n# fp models can run anywhere; quantized models are CPU-only\ntest_eq(_device_supported(_fp, \"cuda\"), True)\ntest_eq(_device_supported(_fp, \"cpu\"), True)\n_ensure_device_supported(_fp, \"cuda\") # no raise for fp model\n\n# A module whose class lives under a *.quantized.* path is detected as quantized\nimport torch.ao.nn.quantized as nnq\n_q = nn.Sequential(nnq.Linear(4, 2))\ntest_eq(_is_quantized(_q), True)\ntest_eq(_device_supported(_q, \"cpu\"), True)\ntest_eq(_device_supported(_q, \"cuda\"), False)\nwith ExceptionExpected(RuntimeError):\n _ensure_device_supported(_q, \"cuda\")" }, { "cell_type": "markdown", diff --git a/nbs/metrics/energy.ipynb b/nbs/metrics/energy.ipynb index 1bef0f3..1366ee8 100644 --- a/nbs/metrics/energy.ipynb +++ b/nbs/metrics/energy.ipynb @@ -34,7 +34,7 @@ "id": "d27f26a4", "metadata": {}, "outputs": [], - "source": "#| export\nfrom __future__ import annotations\nimport logging, os, time, warnings\nfrom dataclasses import dataclass, asdict\nfrom typing import Sequence\n\nimport torch\n\nfrom fasterbench.core import _device_ctx, _sync, _run_on_devices\n\ntry:\n from codecarbon import EmissionsTracker, OfflineEmissionsTracker\n # Suppress codecarbon's verbose logging BEFORE any tracker is constructed,\n # because codecarbon logs messages in __init__ before applying log_level.\n logging.getLogger(\"codecarbon\").setLevel(logging.CRITICAL)\nexcept ImportError:\n EmissionsTracker = OfflineEmissionsTracker = None" + "source": "#| export\nfrom __future__ import annotations\nimport logging, os, time, warnings\nfrom dataclasses import dataclass, asdict\nfrom typing import Sequence\n\nimport torch\n\nfrom fasterbench.core import _device_ctx, _sync, _run_on_devices, _ensure_device_supported\n\ntry:\n from codecarbon import EmissionsTracker, OfflineEmissionsTracker\n # Suppress codecarbon's verbose logging BEFORE any tracker is constructed,\n # because codecarbon logs messages in __init__ before applying log_level.\n logging.getLogger(\"codecarbon\").setLevel(logging.CRITICAL)\nexcept ImportError:\n EmissionsTracker = OfflineEmissionsTracker = None" }, { "cell_type": "code", @@ -42,130 +42,7 @@ "id": "58dcf143", "metadata": {}, "outputs": [], - "source": [ - "#| export\n", - "@dataclass(slots=True)\n", - "class EnergyMetrics:\n", - " \"\"\"Energy consumption and carbon footprint metrics.\"\"\"\n", - " mean_watts: float # average power during measurement\n", - " energy_wh: float # Wh per inference\n", - " co2_eq_g: float # g CO₂-eq per inference\n", - "\n", - " def as_dict(self) -> dict[str, float]:\n", - " return asdict(self)\n", - "\n", - "\n", - "#| export\n", - "def _nan_energy_metrics(device: str) -> EnergyMetrics: # device string (unused, for consistent signature)\n", - " \"\"\"Create EnergyMetrics with NaN values for failed benchmarks.\"\"\"\n", - " nan = float(\"nan\")\n", - " return EnergyMetrics(nan, nan, nan)\n", - "\n", - "\n", - "#| export\n", - "def _clear_stale_codecarbon_lock() -> None:\n", - " \"\"\"Remove stale codecarbon lock file if the owning process no longer exists.\"\"\"\n", - " import tempfile\n", - " lock_path = os.path.join(tempfile.gettempdir(), \".codecarbon.lock\")\n", - " if not os.path.exists(lock_path):\n", - " return\n", - " try:\n", - " # Read the PID from the lock file (codecarbon writes its PID there)\n", - " with open(lock_path) as f:\n", - " content = f.read().strip()\n", - " if content:\n", - " pid = int(content)\n", - " os.kill(pid, 0) # Check if process exists (signal 0 = no-op)\n", - " # Process exists — lock is valid, don't remove\n", - " return\n", - " except (ValueError, ProcessLookupError, PermissionError, OSError):\n", - " pass # PID invalid or process dead — lock is stale\n", - " try:\n", - " os.remove(lock_path)\n", - " except OSError:\n", - " pass\n", - "\n", - "\n", - "#| export\n", - "@torch.no_grad()\n", - "def compute_energy(\n", - " model: torch.nn.Module, # model to benchmark\n", - " sample: torch.Tensor, # input tensor (with batch dimension)\n", - " *,\n", - " device: str | torch.device = \"cpu\", # device to run on\n", - " warmup: int = 20, # warmup iterations\n", - " steps: int = 100, # measurement iterations\n", - " offline: bool = True, # use offline emissions tracker\n", - " country_iso: str | None = None, # country ISO code for carbon intensity\n", - " measure_secs: int = 1, # power sampling interval\n", - ") -> EnergyMetrics:\n", - " \"\"\"Measure power consumption and carbon footprint using codecarbon.\"\"\"\n", - " if EmissionsTracker is None:\n", - " warnings.warn(\"codecarbon not installed – returning NaNs\")\n", - " return _nan_energy_metrics(str(device))\n", - "\n", - " _clear_stale_codecarbon_lock()\n", - "\n", - " Tracker = OfflineEmissionsTracker if offline else EmissionsTracker\n", - " tracker = Tracker(\n", - " project_name=\"fasterbench\",\n", - " country_iso_code=(country_iso or os.getenv(\"NNBENCH_ISO\", \"USA\")),\n", - " measure_power_secs=measure_secs,\n", - " save_to_file=False,\n", - " log_level=\"critical\",\n", - " )\n", - "\n", - " with _device_ctx(device) as dev:\n", - " model = model.eval().to(dev)\n", - " sample = sample.to(dev, non_blocking=True)\n", - "\n", - " for _ in range(warmup):\n", - " model(sample)\n", - " _sync(dev)\n", - "\n", - " tracker.start()\n", - " try:\n", - " t0 = time.perf_counter()\n", - " for _ in range(steps):\n", - " model(sample)\n", - " _sync(dev)\n", - " finally:\n", - " tracker.stop()\n", - " dur_s = time.perf_counter() - t0\n", - "\n", - " # codecarbon silently fails if another instance is running,\n", - " # leaving final_emissions_data as None\n", - " if tracker.final_emissions_data is None:\n", - " warnings.warn(\"codecarbon tracker did not collect data (another instance may be running)\")\n", - " return _nan_energy_metrics(str(device))\n", - "\n", - " ene_kwh = tracker.final_emissions_data.energy_consumed\n", - " co2_kg = tracker.final_emissions\n", - " mean_w = (ene_kwh * 3600_000) / dur_s\n", - "\n", - " return EnergyMetrics(\n", - " mean_watts=mean_w,\n", - " energy_wh=(ene_kwh * 1_000) / steps,\n", - " co2_eq_g=(co2_kg * 1_000) / steps,\n", - " )\n", - "\n", - "\n", - "#| export\n", - "def compute_energy_multi(\n", - " model: torch.nn.Module, # model to benchmark\n", - " sample: torch.Tensor, # input tensor (with batch dimension)\n", - " *,\n", - " devices: Sequence[str | torch.device] | None = None, # devices to benchmark (default: cpu + cuda)\n", - " **kwargs,\n", - ") -> dict[str, EnergyMetrics]:\n", - " \"\"\"Measure energy on multiple devices.\"\"\"\n", - " return _run_on_devices(\n", - " compute_energy, model, sample, devices,\n", - " nan_factory=_nan_energy_metrics,\n", - " metric_name=\"Energy\",\n", - " **kwargs\n", - " )" - ] + "source": "#| export\n@dataclass(slots=True)\nclass EnergyMetrics:\n \"\"\"Energy consumption and carbon footprint metrics.\"\"\"\n mean_watts: float # average power during measurement\n energy_wh: float # Wh per inference\n co2_eq_g: float # g CO₂-eq per inference\n\n def as_dict(self) -> dict[str, float]:\n return asdict(self)\n\n\n#| export\ndef _nan_energy_metrics(device: str) -> EnergyMetrics: # device string (unused, for consistent signature)\n \"\"\"Create EnergyMetrics with NaN values for failed benchmarks.\"\"\"\n nan = float(\"nan\")\n return EnergyMetrics(nan, nan, nan)\n\n\n#| export\ndef _clear_stale_codecarbon_lock() -> None:\n \"\"\"Remove stale codecarbon lock file if the owning process no longer exists.\"\"\"\n import tempfile\n lock_path = os.path.join(tempfile.gettempdir(), \".codecarbon.lock\")\n if not os.path.exists(lock_path):\n return\n try:\n # Read the PID from the lock file (codecarbon writes its PID there)\n with open(lock_path) as f:\n content = f.read().strip()\n if content:\n pid = int(content)\n os.kill(pid, 0) # Check if process exists (signal 0 = no-op)\n # Process exists — lock is valid, don't remove\n return\n except (ValueError, ProcessLookupError, PermissionError, OSError):\n pass # PID invalid or process dead — lock is stale\n try:\n os.remove(lock_path)\n except OSError:\n pass\n\n\n#| export\n@torch.no_grad()\ndef compute_energy(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n device: str | torch.device = \"cpu\", # device to run on\n warmup: int = 20, # warmup iterations\n steps: int = 100, # measurement iterations\n offline: bool = True, # use offline emissions tracker\n country_iso: str | None = None, # country ISO code for carbon intensity\n measure_secs: int = 1, # power sampling interval\n) -> EnergyMetrics:\n \"\"\"Measure power consumption and carbon footprint using codecarbon.\"\"\"\n if EmissionsTracker is None:\n warnings.warn(\"codecarbon not installed – returning NaNs\")\n return _nan_energy_metrics(str(device))\n\n _clear_stale_codecarbon_lock()\n\n Tracker = OfflineEmissionsTracker if offline else EmissionsTracker\n tracker = Tracker(\n project_name=\"fasterbench\",\n country_iso_code=(country_iso or os.getenv(\"NNBENCH_ISO\", \"USA\")),\n measure_power_secs=measure_secs,\n save_to_file=False,\n log_level=\"critical\",\n )\n\n with _device_ctx(device) as dev:\n _ensure_device_supported(model, dev) # quantized ops are CPU-only; avoid SIGSEGV\n model = model.eval().to(dev)\n sample = sample.to(dev, non_blocking=True)\n\n for _ in range(warmup):\n model(sample)\n _sync(dev)\n\n tracker.start()\n try:\n t0 = time.perf_counter()\n for _ in range(steps):\n model(sample)\n _sync(dev)\n finally:\n tracker.stop()\n dur_s = time.perf_counter() - t0\n\n # codecarbon silently fails if another instance is running,\n # leaving final_emissions_data as None\n if tracker.final_emissions_data is None:\n warnings.warn(\"codecarbon tracker did not collect data (another instance may be running)\")\n return _nan_energy_metrics(str(device))\n\n ene_kwh = tracker.final_emissions_data.energy_consumed\n co2_kg = tracker.final_emissions\n mean_w = (ene_kwh * 3600_000) / dur_s\n\n return EnergyMetrics(\n mean_watts=mean_w,\n energy_wh=(ene_kwh * 1_000) / steps,\n co2_eq_g=(co2_kg * 1_000) / steps,\n )\n\n\n#| export\ndef compute_energy_multi(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n devices: Sequence[str | torch.device] | None = None, # devices to benchmark (default: cpu + cuda)\n **kwargs,\n) -> dict[str, EnergyMetrics]:\n \"\"\"Measure energy on multiple devices.\"\"\"\n return _run_on_devices(\n compute_energy, model, sample, devices,\n nan_factory=_nan_energy_metrics,\n metric_name=\"Energy\",\n **kwargs\n )" }, { "cell_type": "code", diff --git a/nbs/metrics/memory.ipynb b/nbs/metrics/memory.ipynb index 5e92945..01f6be5 100644 --- a/nbs/metrics/memory.ipynb +++ b/nbs/metrics/memory.ipynb @@ -34,7 +34,7 @@ "id": "16cd91b6", "metadata": {}, "outputs": [], - "source": "#| export\nfrom __future__ import annotations\n\nimport warnings\nfrom dataclasses import dataclass, asdict\nfrom typing import Sequence\nfrom fasterbench.core import _bytes_to_mib, _run_on_devices\n\nimport numpy as np\nimport torch\n\ntry:\n import psutil\nexcept ImportError:\n psutil = None" + "source": "#| export\nfrom __future__ import annotations\n\nimport warnings\nfrom dataclasses import dataclass, asdict\nfrom typing import Sequence\nfrom fasterbench.core import _bytes_to_mib, _run_on_devices, _ensure_device_supported\n\nimport numpy as np\nimport torch\n\ntry:\n import psutil\nexcept ImportError:\n psutil = None" }, { "cell_type": "code", @@ -42,7 +42,7 @@ "id": "08a0264a", "metadata": {}, "outputs": [], - "source": "#| export\n@dataclass(slots=True)\nclass MemoryMetrics:\n \"\"\"Memory usage metrics for a single device.\"\"\"\n avg_mib: float\n peak_mib: float\n reserved_mib: float # GPU-only (NaN for CPU)\n device: str\n\n def as_dict(self) -> dict[str, float | str]:\n return asdict(self)\n\n\n#| export\ndef _nan_memory_metrics(device: str) -> MemoryMetrics: # device string\n \"\"\"Create MemoryMetrics with NaN values for failed benchmarks.\"\"\"\n nan = float(\"nan\")\n return MemoryMetrics(nan, nan, nan, device)\n\n\n#| export\ndef _gpu_metrics(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n warmup: int, # warmup iterations\n steps: int, # measurement iterations\n) -> MemoryMetrics:\n \"\"\"Measure GPU memory usage.\"\"\"\n dev = torch.device(\"cuda\")\n model = model.eval().to(dev)\n sample = sample.to(dev)\n\n for _ in range(warmup):\n model(sample)\n torch.cuda.synchronize(dev)\n\n alloc, alloc_peak, reserv_peak = [], [], []\n for _ in range(steps):\n torch.cuda.reset_peak_memory_stats(dev)\n model(sample)\n torch.cuda.synchronize(dev)\n alloc.append(torch.cuda.memory_allocated(dev))\n alloc_peak.append(torch.cuda.max_memory_allocated(dev))\n reserv_peak.append(torch.cuda.max_memory_reserved(dev))\n\n return MemoryMetrics(\n avg_mib=_bytes_to_mib(float(np.mean(alloc))),\n peak_mib=_bytes_to_mib(float(np.mean(alloc_peak))),\n reserved_mib=_bytes_to_mib(float(np.mean(reserv_peak))),\n device=\"cuda\",\n )\n\n\n#| export\ndef _cpu_metrics(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n warmup: int, # warmup iterations\n steps: int, # measurement iterations\n) -> MemoryMetrics:\n \"\"\"Measure CPU memory usage via psutil.\"\"\"\n if psutil is None:\n warnings.warn(\"psutil not available – returning NaNs for CPU memory\")\n return _nan_memory_metrics(\"cpu\")\n\n proc = psutil.Process()\n model = model.eval().cpu()\n sample = sample.cpu()\n\n rss0 = proc.memory_info().rss\n\n for _ in range(warmup):\n model(sample)\n\n diffs: list[int] = []\n peaks: list[int] = []\n for _ in range(steps):\n rss_before = proc.memory_info().rss\n model(sample)\n rss_after = proc.memory_info().rss\n diffs.append(max(0, rss_after - rss_before))\n peaks.append(max(0, rss_after - rss0))\n\n return MemoryMetrics(\n avg_mib=_bytes_to_mib(float(np.mean(diffs))),\n peak_mib=_bytes_to_mib(float(np.max(peaks))),\n reserved_mib=float(\"nan\"),\n device=\"cpu\",\n )\n\n\n#| export\ndef compute_memory(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n device: str | torch.device = \"cpu\", # device to run on\n warmup: int = 10, # warmup iterations\n steps: int = 100, # measurement iterations\n) -> MemoryMetrics:\n \"\"\"Measure memory usage on specified device.\"\"\"\n device_str = str(device)\n if device_str == \"cuda\" or (device_str == \"auto\" and torch.cuda.is_available()):\n if not torch.cuda.is_available():\n warnings.warn(\"CUDA requested but not available – falling back to CPU\")\n return _cpu_metrics(model, sample, warmup=warmup, steps=steps)\n return _gpu_metrics(model, sample, warmup=warmup, steps=steps)\n return _cpu_metrics(model, sample, warmup=warmup, steps=steps)\n\n\n#| export\ndef compute_memory_multi(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n devices: Sequence[str | torch.device] | None = None, # devices to benchmark (default: cpu + cuda)\n warmup: int = 10, # warmup iterations\n steps: int = 100, # measurement iterations\n) -> dict[str, MemoryMetrics]:\n \"\"\"Measure memory on multiple devices.\"\"\"\n return _run_on_devices(\n compute_memory, model, sample, devices,\n nan_factory=_nan_memory_metrics,\n metric_name=\"Memory\",\n warmup=warmup,\n steps=steps,\n )" + "source": "#| export\n@dataclass(slots=True)\nclass MemoryMetrics:\n \"\"\"Memory usage metrics for a single device.\"\"\"\n avg_mib: float\n peak_mib: float\n reserved_mib: float # GPU-only (NaN for CPU)\n device: str\n\n def as_dict(self) -> dict[str, float | str]:\n return asdict(self)\n\n\n#| export\ndef _nan_memory_metrics(device: str) -> MemoryMetrics: # device string\n \"\"\"Create MemoryMetrics with NaN values for failed benchmarks.\"\"\"\n nan = float(\"nan\")\n return MemoryMetrics(nan, nan, nan, device)\n\n\n#| export\ndef _gpu_metrics(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n warmup: int, # warmup iterations\n steps: int, # measurement iterations\n) -> MemoryMetrics:\n \"\"\"Measure GPU memory usage.\"\"\"\n dev = torch.device(\"cuda\")\n _ensure_device_supported(model, dev) # quantized ops are CPU-only; avoid SIGSEGV\n model = model.eval().to(dev)\n sample = sample.to(dev)\n\n for _ in range(warmup):\n model(sample)\n torch.cuda.synchronize(dev)\n\n alloc, alloc_peak, reserv_peak = [], [], []\n for _ in range(steps):\n torch.cuda.reset_peak_memory_stats(dev)\n model(sample)\n torch.cuda.synchronize(dev)\n alloc.append(torch.cuda.memory_allocated(dev))\n alloc_peak.append(torch.cuda.max_memory_allocated(dev))\n reserv_peak.append(torch.cuda.max_memory_reserved(dev))\n\n return MemoryMetrics(\n avg_mib=_bytes_to_mib(float(np.mean(alloc))),\n peak_mib=_bytes_to_mib(float(np.mean(alloc_peak))),\n reserved_mib=_bytes_to_mib(float(np.mean(reserv_peak))),\n device=\"cuda\",\n )\n\n\n#| export\ndef _cpu_metrics(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n warmup: int, # warmup iterations\n steps: int, # measurement iterations\n) -> MemoryMetrics:\n \"\"\"Measure CPU memory usage via psutil.\"\"\"\n if psutil is None:\n warnings.warn(\"psutil not available – returning NaNs for CPU memory\")\n return _nan_memory_metrics(\"cpu\")\n\n proc = psutil.Process()\n model = model.eval().cpu()\n sample = sample.cpu()\n\n rss0 = proc.memory_info().rss\n\n for _ in range(warmup):\n model(sample)\n\n diffs: list[int] = []\n peaks: list[int] = []\n for _ in range(steps):\n rss_before = proc.memory_info().rss\n model(sample)\n rss_after = proc.memory_info().rss\n diffs.append(max(0, rss_after - rss_before))\n peaks.append(max(0, rss_after - rss0))\n\n return MemoryMetrics(\n avg_mib=_bytes_to_mib(float(np.mean(diffs))),\n peak_mib=_bytes_to_mib(float(np.max(peaks))),\n reserved_mib=float(\"nan\"),\n device=\"cpu\",\n )\n\n\n#| export\ndef compute_memory(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n device: str | torch.device = \"cpu\", # device to run on\n warmup: int = 10, # warmup iterations\n steps: int = 100, # measurement iterations\n) -> MemoryMetrics:\n \"\"\"Measure memory usage on specified device.\"\"\"\n device_str = str(device)\n if device_str == \"cuda\" or (device_str == \"auto\" and torch.cuda.is_available()):\n if not torch.cuda.is_available():\n warnings.warn(\"CUDA requested but not available – falling back to CPU\")\n return _cpu_metrics(model, sample, warmup=warmup, steps=steps)\n return _gpu_metrics(model, sample, warmup=warmup, steps=steps)\n return _cpu_metrics(model, sample, warmup=warmup, steps=steps)\n\n\n#| export\ndef compute_memory_multi(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n devices: Sequence[str | torch.device] | None = None, # devices to benchmark (default: cpu + cuda)\n warmup: int = 10, # warmup iterations\n steps: int = 100, # measurement iterations\n) -> dict[str, MemoryMetrics]:\n \"\"\"Measure memory on multiple devices.\"\"\"\n return _run_on_devices(\n compute_memory, model, sample, devices,\n nan_factory=_nan_memory_metrics,\n metric_name=\"Memory\",\n warmup=warmup,\n steps=steps,\n )" }, { "cell_type": "code", diff --git a/nbs/metrics/speed.ipynb b/nbs/metrics/speed.ipynb index f8b876d..3102546 100644 --- a/nbs/metrics/speed.ipynb +++ b/nbs/metrics/speed.ipynb @@ -34,20 +34,7 @@ "id": "e6b40d9e", "metadata": {}, "outputs": [], - "source": [ - "#| export\n", - "from __future__ import annotations\n", - "import math, time, warnings\n", - "from dataclasses import dataclass, asdict\n", - "from typing import Mapping, Sequence\n", - "\n", - "import numpy as np\n", - "import torch\n", - "import torch.nn as nn\n", - "from torch.utils.benchmark import Timer\n", - "\n", - "from fasterbench.core import _device_ctx, _sync, _run_on_devices" - ] + "source": "#| export\nfrom __future__ import annotations\nimport math, time, warnings\nfrom dataclasses import dataclass, asdict\nfrom typing import Mapping, Sequence\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom torch.utils.benchmark import Timer\n\nfrom fasterbench.core import _device_ctx, _sync, _run_on_devices, _ensure_device_supported" }, { "cell_type": "code", @@ -55,179 +42,7 @@ "id": "eb5e3063", "metadata": {}, "outputs": [], - "source": [ - "#| export\n", - "@dataclass(slots=True)\n", - "class SpeedMetrics:\n", - " \"\"\"Latency and throughput metrics for a single device.\"\"\"\n", - " p50_ms: float\n", - " p90_ms: float\n", - " p99_ms: float\n", - " mean_ms: float\n", - " std_ms: float\n", - " throughput_s: float\n", - "\n", - " def as_dict(self) -> dict[str, float]:\n", - " return asdict(self)\n", - "\n", - "\n", - "#| export\n", - "def _nan_speed_metrics(device: str) -> SpeedMetrics: # device string (unused, for consistent signature)\n", - " \"\"\"Create SpeedMetrics with NaN values for failed benchmarks.\"\"\"\n", - " nan = float(\"nan\")\n", - " return SpeedMetrics(nan, nan, nan, nan, nan, nan)\n", - "\n", - "\n", - "#| export\n", - "def _stats(lat_ms: np.ndarray, batch: int) -> Mapping[str, float]:\n", - " \"\"\"Compute latency statistics from raw measurements.\"\"\"\n", - " p50, p90, p99 = np.percentile(lat_ms, [50, 90, 99])\n", - " mean = float(lat_ms.mean())\n", - " return {\n", - " \"p50_ms\": float(p50),\n", - " \"p90_ms\": float(p90),\n", - " \"p99_ms\": float(p99),\n", - " \"mean_ms\": mean,\n", - " \"std_ms\": float(lat_ms.std(ddof=1)) if lat_ms.size > 1 else 0.0,\n", - " \"throughput_s\": batch * 1000.0 / mean,\n", - " }\n", - "\n", - "\n", - "#| export\n", - "def _forward_latencies(\n", - " model: nn.Module, # model to benchmark\n", - " sample: torch.Tensor, # input tensor (with batch dimension)\n", - " *,\n", - " device: str | torch.device = \"cpu\", # device to run on\n", - " warmup: int = 20, # warmup iterations\n", - " steps: int = 100, # measurement iterations\n", - " use_torch_timer: bool | None = None, # force torch.utils.benchmark.Timer (default: auto)\n", - ") -> np.ndarray:\n", - " \"\"\"Collect forward-pass latencies (ms) using optimal timing for each device.\"\"\"\n", - " if use_torch_timer is None:\n", - " use_torch_timer = torch.device(device).type == \"cpu\"\n", - "\n", - " with _device_ctx(device) as dev:\n", - " model.eval().to(dev)\n", - " sample = sample.to(dev, non_blocking=True)\n", - "\n", - " for _ in range(warmup):\n", - " model(sample)\n", - " _sync(dev)\n", - "\n", - " lat: list[float] = []\n", - " if use_torch_timer and dev.type == \"cpu\":\n", - " t = Timer(stmt=\"model(x)\", globals={\"model\": model, \"x\": sample})\n", - " m = t.blocked_autorange(min_run_time=0.3)\n", - " per_iter = (np.asarray(m.raw_times) / m.number_per_run) * 1e3\n", - " lat = per_iter.tolist()\n", - " elif dev.type == \"cuda\":\n", - " start_evt, end_evt = torch.cuda.Event(True), torch.cuda.Event(True)\n", - " for _ in range(steps):\n", - " start_evt.record()\n", - " model(sample)\n", - " end_evt.record()\n", - " _sync(dev)\n", - " lat.append(start_evt.elapsed_time(end_evt))\n", - " else:\n", - " for _ in range(steps):\n", - " t0 = time.perf_counter()\n", - " model(sample)\n", - " lat.append((time.perf_counter() - t0) * 1e3)\n", - " return np.asarray(lat, dtype=np.float32)\n", - "\n", - "\n", - "#| export\n", - "@torch.no_grad()\n", - "def compute_speed(\n", - " model: nn.Module, # model to benchmark\n", - " sample: torch.Tensor, # input tensor (with batch dimension)\n", - " *,\n", - " device: str | torch.device = \"cpu\", # device to run on\n", - " warmup: int = 20, # warmup iterations\n", - " steps: int = 100, # measurement iterations\n", - ") -> SpeedMetrics:\n", - " \"\"\"Measure latency and throughput on a single device.\"\"\"\n", - " lat = _forward_latencies(model, sample, device=device, warmup=warmup, steps=steps)\n", - " return SpeedMetrics(**_stats(lat, sample.size(0)))\n", - "\n", - "\n", - "#| export\n", - "def compute_speed_multi(\n", - " model: nn.Module, # model to benchmark\n", - " sample: torch.Tensor, # input tensor (with batch dimension)\n", - " *,\n", - " devices: Sequence[str | torch.device] | None = None, # devices to benchmark (default: cpu + cuda if available)\n", - " **kwargs,\n", - ") -> dict[str, SpeedMetrics]:\n", - " \"\"\"Measure latency/throughput on multiple devices.\"\"\"\n", - " return _run_on_devices(\n", - " compute_speed, model, sample, devices,\n", - " nan_factory=_nan_speed_metrics,\n", - " metric_name=\"Speed\",\n", - " **kwargs\n", - " )\n", - "\n", - "\n", - "#| export\n", - "def sweep_threads(\n", - " model: nn.Module, # model to benchmark\n", - " sample: torch.Tensor, # input tensor (with batch dimension)\n", - " thread_counts: Sequence[int] = (1, 2, 4, 8), # thread counts to test\n", - " *,\n", - " warmup: int = 20, # warmup iterations per thread count\n", - " steps: int = 100, # measurement iterations per thread count\n", - ") -> list[dict]:\n", - " \"\"\"Sweep CPU thread counts to find optimal parallelism.\"\"\"\n", - " rows = []\n", - " for n in thread_counts:\n", - " torch.set_num_threads(n)\n", - " lat = _forward_latencies(model, sample, device=\"cpu\", warmup=warmup, steps=steps, use_torch_timer=True)\n", - " rows.append({\"threads\": n, **_stats(lat, sample.size(0))})\n", - " return rows\n", - "\n", - "\n", - "#| export\n", - "def sweep_batch_sizes(\n", - " model: nn.Module, # model to benchmark\n", - " input_shape: Sequence[int], # input shape WITHOUT batch dim, e.g. (3, 224, 224)\n", - " batch_sizes: Sequence[int] = (1, 2, 4, 8, 16, 32), # batch sizes to test\n", - " *,\n", - " device: str | torch.device = \"cuda\", # device to run on\n", - " warmup: int = 20, # warmup iterations per batch size\n", - " steps: int = 100, # measurement iterations per batch size\n", - ") -> list[dict]:\n", - " \"\"\"Sweep batch sizes to find optimal throughput.\"\"\"\n", - " rows = []\n", - " for bs in batch_sizes:\n", - " try:\n", - " dummy = torch.randn(bs, *input_shape)\n", - " lat = _forward_latencies(model, dummy, device=device, warmup=warmup, steps=steps)\n", - " stats = _stats(lat, bs)\n", - " rows.append({\"batch_size\": bs, \"latency_per_sample_ms\": stats[\"mean_ms\"] / bs, **stats})\n", - " except (RuntimeError, torch.cuda.OutOfMemoryError) as e:\n", - " warnings.warn(f\"Batch size {bs} failed (likely OOM): {e}\")\n", - " rows.append({\"batch_size\": bs, \"mean_ms\": float(\"nan\"), \"throughput_s\": float(\"nan\")})\n", - " return rows\n", - "\n", - "\n", - "#| export\n", - "def sweep_latency(\n", - " model: nn.Module, # model to benchmark\n", - " shapes: Sequence[Sequence[int]], # input shapes to test, e.g. [(1,3,224,224), (1,3,384,384)]\n", - " *,\n", - " device: str | torch.device = \"cuda\", # device to run on\n", - " warmup: int = 20, # warmup iterations per shape\n", - " steps: int = 100, # measurement iterations per shape\n", - ") -> list[dict]:\n", - " \"\"\"Sweep input shapes to analyze latency vs resolution.\"\"\"\n", - " rows = []\n", - " for shape in shapes:\n", - " dummy = torch.empty(*shape)\n", - " lat = _forward_latencies(model, dummy, device=device, warmup=warmup, steps=steps)\n", - " rows.append({\"shape\": \"×\".join(map(str, shape)), **_stats(lat, shape[0])})\n", - " return rows" - ] + "source": "#| export\n@dataclass(slots=True)\nclass SpeedMetrics:\n \"\"\"Latency and throughput metrics for a single device.\"\"\"\n p50_ms: float\n p90_ms: float\n p99_ms: float\n mean_ms: float\n std_ms: float\n throughput_s: float\n\n def as_dict(self) -> dict[str, float]:\n return asdict(self)\n\n\n#| export\ndef _nan_speed_metrics(device: str) -> SpeedMetrics: # device string (unused, for consistent signature)\n \"\"\"Create SpeedMetrics with NaN values for failed benchmarks.\"\"\"\n nan = float(\"nan\")\n return SpeedMetrics(nan, nan, nan, nan, nan, nan)\n\n\n#| export\ndef _stats(lat_ms: np.ndarray, batch: int) -> Mapping[str, float]:\n \"\"\"Compute latency statistics from raw measurements.\"\"\"\n p50, p90, p99 = np.percentile(lat_ms, [50, 90, 99])\n mean = float(lat_ms.mean())\n return {\n \"p50_ms\": float(p50),\n \"p90_ms\": float(p90),\n \"p99_ms\": float(p99),\n \"mean_ms\": mean,\n \"std_ms\": float(lat_ms.std(ddof=1)) if lat_ms.size > 1 else 0.0,\n \"throughput_s\": batch * 1000.0 / mean,\n }\n\n\n#| export\ndef _forward_latencies(\n model: nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n device: str | torch.device = \"cpu\", # device to run on\n warmup: int = 20, # warmup iterations\n steps: int = 100, # measurement iterations\n use_torch_timer: bool | None = None, # force torch.utils.benchmark.Timer (default: auto)\n) -> np.ndarray:\n \"\"\"Collect forward-pass latencies (ms) using optimal timing for each device.\"\"\"\n if use_torch_timer is None:\n use_torch_timer = torch.device(device).type == \"cpu\"\n\n with _device_ctx(device) as dev:\n _ensure_device_supported(model, dev) # quantized ops are CPU-only; avoid SIGSEGV\n model.eval().to(dev)\n sample = sample.to(dev, non_blocking=True)\n\n for _ in range(warmup):\n model(sample)\n _sync(dev)\n\n lat: list[float] = []\n if use_torch_timer and dev.type == \"cpu\":\n t = Timer(stmt=\"model(x)\", globals={\"model\": model, \"x\": sample})\n m = t.blocked_autorange(min_run_time=0.3)\n per_iter = (np.asarray(m.raw_times) / m.number_per_run) * 1e3\n lat = per_iter.tolist()\n elif dev.type == \"cuda\":\n start_evt, end_evt = torch.cuda.Event(True), torch.cuda.Event(True)\n for _ in range(steps):\n start_evt.record()\n model(sample)\n end_evt.record()\n _sync(dev)\n lat.append(start_evt.elapsed_time(end_evt))\n else:\n for _ in range(steps):\n t0 = time.perf_counter()\n model(sample)\n lat.append((time.perf_counter() - t0) * 1e3)\n return np.asarray(lat, dtype=np.float32)\n\n\n#| export\n@torch.no_grad()\ndef compute_speed(\n model: nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n device: str | torch.device = \"cpu\", # device to run on\n warmup: int = 20, # warmup iterations\n steps: int = 100, # measurement iterations\n) -> SpeedMetrics:\n \"\"\"Measure latency and throughput on a single device.\"\"\"\n lat = _forward_latencies(model, sample, device=device, warmup=warmup, steps=steps)\n return SpeedMetrics(**_stats(lat, sample.size(0)))\n\n\n#| export\ndef compute_speed_multi(\n model: nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n devices: Sequence[str | torch.device] | None = None, # devices to benchmark (default: cpu + cuda if available)\n **kwargs,\n) -> dict[str, SpeedMetrics]:\n \"\"\"Measure latency/throughput on multiple devices.\"\"\"\n return _run_on_devices(\n compute_speed, model, sample, devices,\n nan_factory=_nan_speed_metrics,\n metric_name=\"Speed\",\n **kwargs\n )\n\n\n#| export\ndef sweep_threads(\n model: nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n thread_counts: Sequence[int] = (1, 2, 4, 8), # thread counts to test\n *,\n warmup: int = 20, # warmup iterations per thread count\n steps: int = 100, # measurement iterations per thread count\n) -> list[dict]:\n \"\"\"Sweep CPU thread counts to find optimal parallelism.\"\"\"\n rows = []\n for n in thread_counts:\n torch.set_num_threads(n)\n lat = _forward_latencies(model, sample, device=\"cpu\", warmup=warmup, steps=steps, use_torch_timer=True)\n rows.append({\"threads\": n, **_stats(lat, sample.size(0))})\n return rows\n\n\n#| export\ndef sweep_batch_sizes(\n model: nn.Module, # model to benchmark\n input_shape: Sequence[int], # input shape WITHOUT batch dim, e.g. (3, 224, 224)\n batch_sizes: Sequence[int] = (1, 2, 4, 8, 16, 32), # batch sizes to test\n *,\n device: str | torch.device = \"cuda\", # device to run on\n warmup: int = 20, # warmup iterations per batch size\n steps: int = 100, # measurement iterations per batch size\n) -> list[dict]:\n \"\"\"Sweep batch sizes to find optimal throughput.\"\"\"\n rows = []\n for bs in batch_sizes:\n try:\n dummy = torch.randn(bs, *input_shape)\n lat = _forward_latencies(model, dummy, device=device, warmup=warmup, steps=steps)\n stats = _stats(lat, bs)\n rows.append({\"batch_size\": bs, \"latency_per_sample_ms\": stats[\"mean_ms\"] / bs, **stats})\n except (RuntimeError, torch.cuda.OutOfMemoryError) as e:\n warnings.warn(f\"Batch size {bs} failed (likely OOM): {e}\")\n rows.append({\"batch_size\": bs, \"mean_ms\": float(\"nan\"), \"throughput_s\": float(\"nan\")})\n return rows\n\n\n#| export\ndef sweep_latency(\n model: nn.Module, # model to benchmark\n shapes: Sequence[Sequence[int]], # input shapes to test, e.g. [(1,3,224,224), (1,3,384,384)]\n *,\n device: str | torch.device = \"cuda\", # device to run on\n warmup: int = 20, # warmup iterations per shape\n steps: int = 100, # measurement iterations per shape\n) -> list[dict]:\n \"\"\"Sweep input shapes to analyze latency vs resolution.\"\"\"\n rows = []\n for shape in shapes:\n dummy = torch.empty(*shape)\n lat = _forward_latencies(model, dummy, device=device, warmup=warmup, steps=steps)\n rows.append({\"shape\": \"×\".join(map(str, shape)), **_stats(lat, shape[0])})\n return rows" }, { "cell_type": "code", From 346f73de52ab1380969959c4ce64127af95a62b0 Mon Sep 17 00:00:00 2001 From: nathanhubens Date: Mon, 6 Jul 2026 15:09:48 +0200 Subject: [PATCH 2/2] fix(metrics): correct params & MACs for quantized models Quantized weights live in packed params (_packed_params), not .parameters(), so get_num_parameters returned 0; thop has no quantized-op handlers so MACs read ~0. But quantization keeps the SAME param & MAC counts (int8 vs fp32) -- only bytes and op-cost drop. Now correct: - get_num_parameters: when _is_quantized, add the numel of quantized modules' callable .weight()/.bias() accessors -> params(quant) == params(fp32). - compute_compute: route quantized models through a forward-hook MAC counter (handles quantized + fp Conv2d/Linear); the thop path is unchanged for fp models. Measured: params 0 -> 5220 (= fp32); MACs 0.008 -> 1.29M (0.99x fp32); disk still drops (int8). Protects the before/after report from fake 'params/MACs -> 0 (-100%)' rows on quantized outputs. nbdev-test green. --- fasterbench/_modidx.py | 4 +++ fasterbench/compute.py | 54 ++++++++++++++++++++++++++++++++ fasterbench/size.py | 40 +++++++++++++++++++++--- nbs/metrics/compute.ipynb | 31 ++++++------------- nbs/metrics/size.ipynb | 65 ++++++--------------------------------- 5 files changed, 114 insertions(+), 80 deletions(-) diff --git a/fasterbench/_modidx.py b/fasterbench/_modidx.py index f7ce56f..171e36b 100644 --- a/fasterbench/_modidx.py +++ b/fasterbench/_modidx.py @@ -35,6 +35,8 @@ 'fasterbench/compute.py'), 'fasterbench.compute.ComputeMetrics.macs_available': ( 'metrics/compute.html#computemetrics.macs_available', 'fasterbench/compute.py'), + 'fasterbench.compute._count_macs_hooks': ( 'metrics/compute.html#_count_macs_hooks', + 'fasterbench/compute.py'), 'fasterbench.compute.compute_compute': ( 'metrics/compute.html#compute_compute', 'fasterbench/compute.py')}, 'fasterbench.core': { 'fasterbench.core._bytes_to_mib': ('core/core.html#_bytes_to_mib', 'fasterbench/core.py'), @@ -166,6 +168,8 @@ 'fasterbench/roofline.py')}, 'fasterbench.size': { 'fasterbench.size.SizeMetrics': ('metrics/size.html#sizemetrics', 'fasterbench/size.py'), 'fasterbench.size.SizeMetrics.as_dict': ('metrics/size.html#sizemetrics.as_dict', 'fasterbench/size.py'), + 'fasterbench.size._quantized_param_count': ( 'metrics/size.html#_quantized_param_count', + 'fasterbench/size.py'), 'fasterbench.size.compute_size': ('metrics/size.html#compute_size', 'fasterbench/size.py'), 'fasterbench.size.get_model_size': ('metrics/size.html#get_model_size', 'fasterbench/size.py'), 'fasterbench.size.get_num_parameters': ('metrics/size.html#get_num_parameters', 'fasterbench/size.py')}, diff --git a/fasterbench/compute.py b/fasterbench/compute.py index 777b4bd..70b5e40 100644 --- a/fasterbench/compute.py +++ b/fasterbench/compute.py @@ -3,12 +3,15 @@ # %% ../nbs/metrics/compute.ipynb #0091d170 from __future__ import annotations +import math import warnings from dataclasses import dataclass import torch import torch.nn as nn +from .core import _is_quantized + try: from thop import profile as _thop_profile except ImportError: @@ -39,6 +42,48 @@ def as_dict(self) -> dict[str, float]: } +#| export +def _count_macs_hooks( + model: nn.Module, # model to profile (run on CPU) + sample: torch.Tensor, # input tensor (with batch dimension) +) -> int: + """Count Conv2d/Linear MACs via forward hooks (fp *and* quantized modules). + + thop has no handlers for quantized ops (it reports ~0), so we count manually. + Quantization does not change the number of multiply-accumulates: an int8 conv + does the same MACs as its fp32 twin. Uses module attributes (`out_channels`, + `kernel_size`, `in_features`, ...) which both fp and quantized modules expose, + and reads output spatial dims from the produced tensor. + """ + macs = 0 + handles = [] + + def conv_hook(m, inp, out): + nonlocal macs + out_spatial = math.prod(out.shape[2:]) # out_H * out_W + kernel = math.prod(m.kernel_size) # kH * kW + macs += out.shape[0] * m.out_channels * (m.in_channels // m.groups) * kernel * out_spatial + + def linear_hook(m, inp, out): + nonlocal macs + leading = out.numel() // m.out_features # batch/token dims + macs += leading * m.in_features * m.out_features + + for m in model.modules(): + if hasattr(m, "kernel_size") and hasattr(m, "out_channels") and hasattr(m, "groups"): + handles.append(m.register_forward_hook(conv_hook)) + elif hasattr(m, "in_features") and hasattr(m, "out_features"): + handles.append(m.register_forward_hook(linear_hook)) + + try: + with torch.no_grad(): + model(sample) + finally: + for h in handles: + h.remove() + return macs + + #| export def compute_compute( model: nn.Module, # model to analyze @@ -55,6 +100,15 @@ def compute_compute( macs_m: float | None = None + # Quantized models: thop has no quantized-op handlers (returns ~0). Count the + # Conv2d/Linear MACs manually — quantization preserves the MAC count. + if _is_quantized(model): + try: + macs_m = round(_count_macs_hooks(model, sample) / 1e6, 3) + except Exception as e: + warnings.warn(f"quantized MAC counting failed: {e}") + return ComputeMetrics(macs_m=macs_m) + if _thop_profile is not None: try: mac_raw, _ = _thop_profile(model, inputs=(sample,), verbose=False) diff --git a/fasterbench/size.py b/fasterbench/size.py index b767f25..e596b62 100644 --- a/fasterbench/size.py +++ b/fasterbench/size.py @@ -3,7 +3,7 @@ # %% ../nbs/metrics/size.ipynb #4c37777e from __future__ import annotations -from .core import _bytes_to_mib +from .core import _bytes_to_mib, _is_quantized import torch import io from dataclasses import dataclass, asdict @@ -22,15 +22,47 @@ def get_model_size(model: torch.nn.Module) -> int: # model to measure return buf.getbuffer().nbytes +#| export +def _quantized_param_count(model: torch.nn.Module) -> int: # model with quantized modules + """Sum the numel of packed quantized weights/biases (not in `.parameters()`). + + Quantized `Conv2d`/`Linear` keep their weights in `_packed_params`, exposed via + a *callable* `.weight()`/`.bias()` accessor (unlike a plain `nn.Parameter`, which + is a tensor). Quantization preserves the number of parameters, so counting these + makes a quantized model's param count match its fp32 twin. + """ + total = 0 + for m in model.modules(): + for name in ("weight", "bias"): + accessor = getattr(m, name, None) + if not callable(accessor): # plain nn.Parameter tensors are not callable + continue + try: + t = accessor() + if t is not None: + total += t.numel() + except Exception: + pass + return total + + #| export def get_num_parameters( model: torch.nn.Module, # model to count parameters trainable_only: bool = True, # if True, only count trainable parameters ) -> int: - """Count the number of (optionally trainable) parameters.""" + """Count the number of (optionally trainable) parameters. + + Quantized weights live in packed params outside `.parameters()`, so they are + added separately (regardless of `trainable_only`, as they are never trainable). + """ if trainable_only: - return sum(p.numel() for p in model.parameters() if p.requires_grad) - return sum(p.numel() for p in model.parameters()) + n = sum(p.numel() for p in model.parameters() if p.requires_grad) + else: + n = sum(p.numel() for p in model.parameters()) + if _is_quantized(model): + n += _quantized_param_count(model) + return n #| export diff --git a/nbs/metrics/compute.ipynb b/nbs/metrics/compute.ipynb index 261b566..3df3272 100644 --- a/nbs/metrics/compute.ipynb +++ b/nbs/metrics/compute.ipynb @@ -34,26 +34,7 @@ "id": "0091d170", "metadata": {}, "outputs": [], - "source": [ - "#| export\n", - "from __future__ import annotations\n", - "\n", - "import warnings\n", - "from dataclasses import dataclass\n", - "\n", - "import torch\n", - "import torch.nn as nn\n", - "\n", - "try:\n", - " from thop import profile as _thop_profile\n", - "except ImportError:\n", - " _thop_profile = None\n", - "\n", - "try:\n", - " from torchprofile import profile_macs as _profile_macs\n", - "except ImportError:\n", - " _profile_macs = None" - ] + "source": "#| export\nfrom __future__ import annotations\n\nimport math\nimport warnings\nfrom dataclasses import dataclass\n\nimport torch\nimport torch.nn as nn\n\nfrom fasterbench.core import _is_quantized\n\ntry:\n from thop import profile as _thop_profile\nexcept ImportError:\n _thop_profile = None\n\ntry:\n from torchprofile import profile_macs as _profile_macs\nexcept ImportError:\n _profile_macs = None" }, { "cell_type": "code", @@ -61,7 +42,7 @@ "id": "09401efb", "metadata": {}, "outputs": [], - "source": "#| export\n@dataclass(slots=True)\nclass ComputeMetrics:\n \"\"\"MACs (Multiply-Accumulate operations) in millions.\"\"\"\n macs_m: float | None # None if unavailable\n\n @property\n def macs_available(self) -> bool:\n \"\"\"Check if MACs measurement succeeded.\"\"\"\n return self.macs_m is not None\n\n def as_dict(self) -> dict[str, float]:\n return {\n \"macs_m\": self.macs_m if self.macs_m is not None else float(\"nan\"),\n }\n\n\n#| export\ndef compute_compute(\n model: nn.Module, # model to analyze\n sample: torch.Tensor, # input tensor (with batch dimension)\n) -> ComputeMetrics:\n \"\"\"Compute MACs for a single forward pass.\"\"\"\n try:\n model_device = next(model.parameters()).device\n except StopIteration:\n model_device = torch.device(\"cpu\")\n\n if sample.device != model_device:\n sample = sample.to(model_device)\n\n macs_m: float | None = None\n\n if _thop_profile is not None:\n try:\n mac_raw, _ = _thop_profile(model, inputs=(sample,), verbose=False)\n macs_m = round(mac_raw / 1e6, 3)\n except Exception as e:\n warnings.warn(f\"thop failed: {e}\")\n elif _profile_macs is not None:\n try:\n macs_m = round(_profile_macs(model, sample) / 1e6, 3)\n except Exception as e:\n warnings.warn(f\"torchprofile failed: {e}\")\n else:\n warnings.warn(\"No MAC-counting backend available – skipping MACs\")\n\n return ComputeMetrics(macs_m=macs_m)" + "source": "#| export\n@dataclass(slots=True)\nclass ComputeMetrics:\n \"\"\"MACs (Multiply-Accumulate operations) in millions.\"\"\"\n macs_m: float | None # None if unavailable\n\n @property\n def macs_available(self) -> bool:\n \"\"\"Check if MACs measurement succeeded.\"\"\"\n return self.macs_m is not None\n\n def as_dict(self) -> dict[str, float]:\n return {\n \"macs_m\": self.macs_m if self.macs_m is not None else float(\"nan\"),\n }\n\n\n#| export\ndef _count_macs_hooks(\n model: nn.Module, # model to profile (run on CPU)\n sample: torch.Tensor, # input tensor (with batch dimension)\n) -> int:\n \"\"\"Count Conv2d/Linear MACs via forward hooks (fp *and* quantized modules).\n\n thop has no handlers for quantized ops (it reports ~0), so we count manually.\n Quantization does not change the number of multiply-accumulates: an int8 conv\n does the same MACs as its fp32 twin. Uses module attributes (`out_channels`,\n `kernel_size`, `in_features`, ...) which both fp and quantized modules expose,\n and reads output spatial dims from the produced tensor.\n \"\"\"\n macs = 0\n handles = []\n\n def conv_hook(m, inp, out):\n nonlocal macs\n out_spatial = math.prod(out.shape[2:]) # out_H * out_W\n kernel = math.prod(m.kernel_size) # kH * kW\n macs += out.shape[0] * m.out_channels * (m.in_channels // m.groups) * kernel * out_spatial\n\n def linear_hook(m, inp, out):\n nonlocal macs\n leading = out.numel() // m.out_features # batch/token dims\n macs += leading * m.in_features * m.out_features\n\n for m in model.modules():\n if hasattr(m, \"kernel_size\") and hasattr(m, \"out_channels\") and hasattr(m, \"groups\"):\n handles.append(m.register_forward_hook(conv_hook))\n elif hasattr(m, \"in_features\") and hasattr(m, \"out_features\"):\n handles.append(m.register_forward_hook(linear_hook))\n\n try:\n with torch.no_grad():\n model(sample)\n finally:\n for h in handles:\n h.remove()\n return macs\n\n\n#| export\ndef compute_compute(\n model: nn.Module, # model to analyze\n sample: torch.Tensor, # input tensor (with batch dimension)\n) -> ComputeMetrics:\n \"\"\"Compute MACs for a single forward pass.\"\"\"\n try:\n model_device = next(model.parameters()).device\n except StopIteration:\n model_device = torch.device(\"cpu\")\n\n if sample.device != model_device:\n sample = sample.to(model_device)\n\n macs_m: float | None = None\n\n # Quantized models: thop has no quantized-op handlers (returns ~0). Count the\n # Conv2d/Linear MACs manually — quantization preserves the MAC count.\n if _is_quantized(model):\n try:\n macs_m = round(_count_macs_hooks(model, sample) / 1e6, 3)\n except Exception as e:\n warnings.warn(f\"quantized MAC counting failed: {e}\")\n return ComputeMetrics(macs_m=macs_m)\n\n if _thop_profile is not None:\n try:\n mac_raw, _ = _thop_profile(model, inputs=(sample,), verbose=False)\n macs_m = round(mac_raw / 1e6, 3)\n except Exception as e:\n warnings.warn(f\"thop failed: {e}\")\n elif _profile_macs is not None:\n try:\n macs_m = round(_profile_macs(model, sample) / 1e6, 3)\n except Exception as e:\n warnings.warn(f\"torchprofile failed: {e}\")\n else:\n warnings.warn(\"No MAC-counting backend available – skipping MACs\")\n\n return ComputeMetrics(macs_m=macs_m)" }, { "cell_type": "code", @@ -101,6 +82,14 @@ "#| hide\nfrom fastcore.test import *\n\nimport torch, torch.nn as nn\n_m = nn.Linear(10, 5)\n_x = torch.randn(1, 10)\n_c = compute_compute(_m, _x)\nassert isinstance(_c, ComputeMetrics)" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "b7d64a51", + "metadata": {}, + "outputs": [], + "source": "#| hide\n# Quantization preserves the *number* of MACs (int8 conv/linear do the same\n# multiply-accumulates as their fp32 twin) — only the ops get cheaper. But thop\n# has no quantized-op handlers, so it reports ~0. Verify the manual hook counter\n# recovers a MAC count ~equal to the fp32 twin's thop count.\nfrom torch.ao.quantization import get_default_qconfig, QConfigMapping\nfrom torch.ao.quantization.quantize_fx import prepare_fx, convert_fx\nfrom fasterbench.core import _is_quantized\n\nclass _MacNet(nn.Module):\n \"Conv-dominated net (linear/pool MACs negligible), statically quantizable via FX.\"\n def __init__(self):\n super().__init__()\n self.c1, self.c2 = nn.Conv2d(3, 16, 3, padding=1), nn.Conv2d(16, 32, 3, padding=1)\n self.pool, self.flat, self.fc = nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(32, 10)\n def forward(self, x):\n x = torch.relu(self.c1(x)); x = torch.relu(self.c2(x))\n return self.fc(self.flat(self.pool(x)))\n\n_fp = _MacNet().eval()\n_x = torch.randn(1, 3, 16, 16)\n_fp_macs = compute_compute(_fp, _x).macs_m # fp32 baseline (thop path)\nassert _fp_macs is not None and _fp_macs > 0\n\ntorch.backends.quantized.engine = \"x86\"\n_qmap = QConfigMapping().set_global(get_default_qconfig(\"x86\"))\n_prep = prepare_fx(_MacNet().eval(), _qmap, example_inputs=(_x,))\n_prep(_x) # calibrate\n_qm = convert_fx(_prep)\nassert _is_quantized(_qm)\n\n# BEFORE the fix, thop reports ~0 MACs for a fully quantized model.\n_thop_raw, _ = _thop_profile(_qm, inputs=(_x,), verbose=False)\nassert _thop_raw / 1e6 < 0.1 * _fp_macs # thop is effectively blind here\n\n_q = compute_compute(_qm, _x) # AFTER the fix (manual counter)\nassert _q.macs_m is not None and _q.macs_m > 0 # zero is gone\nassert abs(_q.macs_m - _fp_macs) <= 0.05 * _fp_macs # within ~5% of fp32 MACs" + }, { "cell_type": "markdown", "id": "see_also", diff --git a/nbs/metrics/size.ipynb b/nbs/metrics/size.ipynb index b7a1f31..b726a87 100644 --- a/nbs/metrics/size.ipynb +++ b/nbs/metrics/size.ipynb @@ -34,15 +34,7 @@ "id": "4c37777e", "metadata": {}, "outputs": [], - "source": [ - "#| export\n", - "from __future__ import annotations\n", - "\n", - "from fasterbench.core import _bytes_to_mib\n", - "import torch\n", - "import io\n", - "from dataclasses import dataclass, asdict" - ] + "source": "#| export\nfrom __future__ import annotations\n\nfrom fasterbench.core import _bytes_to_mib, _is_quantized\nimport torch\nimport io\nfrom dataclasses import dataclass, asdict" }, { "cell_type": "code", @@ -50,52 +42,7 @@ "id": "daa12120", "metadata": {}, "outputs": [], - "source": [ - "#| export\n", - "def get_model_size(model: torch.nn.Module) -> int: # model to measure\n", - " \"\"\"Return the on-disk size (bytes) of the serialized model.\"\"\"\n", - " buf = io.BytesIO()\n", - " try:\n", - " model.save(buf)\n", - " except Exception:\n", - " torch.save(model.state_dict(), buf)\n", - " return buf.getbuffer().nbytes\n", - "\n", - "\n", - "#| export\n", - "def get_num_parameters(\n", - " model: torch.nn.Module, # model to count parameters\n", - " trainable_only: bool = True, # if True, only count trainable parameters\n", - ") -> int:\n", - " \"\"\"Count the number of (optionally trainable) parameters.\"\"\"\n", - " if trainable_only:\n", - " return sum(p.numel() for p in model.parameters() if p.requires_grad)\n", - " return sum(p.numel() for p in model.parameters())\n", - "\n", - "\n", - "#| export\n", - "@dataclass(slots=True)\n", - "class SizeMetrics:\n", - " \"\"\"Model size metrics: disk size and parameter count.\"\"\"\n", - " disk_bytes: int\n", - " size_mib: float\n", - " num_params: int\n", - "\n", - " def as_dict(self) -> dict[str, float]:\n", - " return asdict(self)\n", - "\n", - "\n", - "#| export\n", - "def compute_size(\n", - " model: torch.nn.Module, # model to measure\n", - " *,\n", - " params_count: int | None = None, # pre-computed parameter count (avoids recount)\n", - ") -> SizeMetrics:\n", - " \"\"\"Compute size metrics for a model.\"\"\"\n", - " disk = get_model_size(model)\n", - " params = params_count if params_count is not None else get_num_parameters(model)\n", - " return SizeMetrics(disk_bytes=disk, size_mib=_bytes_to_mib(disk), num_params=params)" - ] + "source": "#| export\ndef get_model_size(model: torch.nn.Module) -> int: # model to measure\n \"\"\"Return the on-disk size (bytes) of the serialized model.\"\"\"\n buf = io.BytesIO()\n try:\n model.save(buf)\n except Exception:\n torch.save(model.state_dict(), buf)\n return buf.getbuffer().nbytes\n\n\n#| export\ndef _quantized_param_count(model: torch.nn.Module) -> int: # model with quantized modules\n \"\"\"Sum the numel of packed quantized weights/biases (not in `.parameters()`).\n\n Quantized `Conv2d`/`Linear` keep their weights in `_packed_params`, exposed via\n a *callable* `.weight()`/`.bias()` accessor (unlike a plain `nn.Parameter`, which\n is a tensor). Quantization preserves the number of parameters, so counting these\n makes a quantized model's param count match its fp32 twin.\n \"\"\"\n total = 0\n for m in model.modules():\n for name in (\"weight\", \"bias\"):\n accessor = getattr(m, name, None)\n if not callable(accessor): # plain nn.Parameter tensors are not callable\n continue\n try:\n t = accessor()\n if t is not None:\n total += t.numel()\n except Exception:\n pass\n return total\n\n\n#| export\ndef get_num_parameters(\n model: torch.nn.Module, # model to count parameters\n trainable_only: bool = True, # if True, only count trainable parameters\n) -> int:\n \"\"\"Count the number of (optionally trainable) parameters.\n\n Quantized weights live in packed params outside `.parameters()`, so they are\n added separately (regardless of `trainable_only`, as they are never trainable).\n \"\"\"\n if trainable_only:\n n = sum(p.numel() for p in model.parameters() if p.requires_grad)\n else:\n n = sum(p.numel() for p in model.parameters())\n if _is_quantized(model):\n n += _quantized_param_count(model)\n return n\n\n\n#| export\n@dataclass(slots=True)\nclass SizeMetrics:\n \"\"\"Model size metrics: disk size and parameter count.\"\"\"\n disk_bytes: int\n size_mib: float\n num_params: int\n\n def as_dict(self) -> dict[str, float]:\n return asdict(self)\n\n\n#| export\ndef compute_size(\n model: torch.nn.Module, # model to measure\n *,\n params_count: int | None = None, # pre-computed parameter count (avoids recount)\n) -> SizeMetrics:\n \"\"\"Compute size metrics for a model.\"\"\"\n disk = get_model_size(model)\n params = params_count if params_count is not None else get_num_parameters(model)\n return SizeMetrics(disk_bytes=disk, size_mib=_bytes_to_mib(disk), num_params=params)" }, { "cell_type": "code", @@ -147,6 +94,14 @@ "#| hide\nfrom fastcore.test import *\n\nimport torch.nn as nn\n_m = nn.Linear(10, 5)\n_s = compute_size(_m)\nassert isinstance(_s, SizeMetrics)\nassert _s.num_params > 0\ntest_eq(get_num_parameters(_m), 55)" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "52ea5e13", + "metadata": {}, + "outputs": [], + "source": "#| hide\n# Quantization preserves the *number* of params (int8 conv/linear have the same\n# weight count as their fp32 twin) — only the bytes shrink. But quantized weights\n# live in packed params outside `.parameters()`, so the naive sum reports 0.\n# Verify: quantized param count is non-zero and ~equal to the fp32 count.\nimport torch.nn as nn\nfrom torch.ao.quantization import get_default_qconfig, QConfigMapping\nfrom torch.ao.quantization.quantize_fx import prepare_fx, convert_fx\nfrom fasterbench.core import _is_quantized\n\nclass _SizeNet(nn.Module):\n \"Small conv+linear net, statically quantizable via FX.\"\n def __init__(self):\n super().__init__()\n self.c1, self.c2 = nn.Conv2d(3, 16, 3, padding=1), nn.Conv2d(16, 32, 3, padding=1)\n self.pool, self.flat, self.fc = nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(32, 10)\n def forward(self, x):\n x = torch.relu(self.c1(x)); x = torch.relu(self.c2(x))\n return self.fc(self.flat(self.pool(x)))\n\n_fp = _SizeNet().eval()\n_x = torch.randn(1, 3, 16, 16)\n_fp_params = get_num_parameters(_fp) # fp32 baseline\n\ntorch.backends.quantized.engine = \"x86\"\n_qmap = QConfigMapping().set_global(get_default_qconfig(\"x86\"))\n_prep = prepare_fx(_SizeNet().eval(), _qmap, example_inputs=(_x,))\n_prep(_x) # calibrate\n_qm = convert_fx(_prep)\nassert _is_quantized(_qm)\n\n# BEFORE the fix, the naive `.parameters()` sum is 0 for a fully quantized model.\ntest_eq(sum(p.numel() for p in _qm.parameters() if p.requires_grad), 0)\n\n_q_params = get_num_parameters(_qm) # AFTER the fix\nassert _q_params > 0 # zero is gone\nassert abs(_q_params - _fp_params) <= 0.05 * _fp_params # within ~5% of fp32\n# trainable_only must NOT hide quantized weights (they are never trainable)\ntest_eq(get_num_parameters(_qm, trainable_only=False), _q_params)" + }, { "cell_type": "markdown", "execution_count": null,