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",