fix(benchmark): don't segfault benchmarking quantized models on CUDA

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.
 

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')},

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,13 +113,23 @@ 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()
 
     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:

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)
 

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)
 

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)
 

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,

nbs/core/core.ipynb

@@ -48,15 +48,15 @@
    "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",
    "execution_count": null,
    "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",