rpx_benchmark.profiler
Hardware-agnostic model efficiency profiling for RPX.
Three-tier efficiency reporting
Tier 1 — Model properties (fully hardware-agnostic): Parameters (M), FLOPs (G), MACs (G), estimated DRAM traffic (GB), arithmetic intensity (FLOP/Byte). Identical on any hardware.
Tier 2 — Roofline latency bounds (hardware-parametric):
Given Tier-1 numbers and a GPUSpec, compute theoretical
lower-bound latency on any GPU via the roofline model:
latency = max(FLOPs/peak_flops, traffic/peak_bw).
Published for reference GPUs; anyone can recompute for their own.
Tier 3 — Measured latency (hardware-specific):
Wall-clock p50/p95/p99, peak memory per backend, throughput.
Always accompanied by a SystemCard.
FLOPs convention (for reproducibility):
- Single forward pass, batch_size=1
- Input resolution: task default or caller-specified
- Precision: FP32 for counting (hardware-agnostic)
- API-only models: report None (shown as "N/A (API)" in tables)
1"""Hardware-agnostic model efficiency profiling for RPX. 2 3Three-tier efficiency reporting 4=============================== 5 6**Tier 1 — Model properties (fully hardware-agnostic):** 7 Parameters (M), FLOPs (G), MACs (G), estimated DRAM traffic (GB), 8 arithmetic intensity (FLOP/Byte). Identical on any hardware. 9 10**Tier 2 — Roofline latency bounds (hardware-parametric):** 11 Given Tier-1 numbers and a :class:`GPUSpec`, compute theoretical 12 lower-bound latency on any GPU via the roofline model: 13 ``latency = max(FLOPs/peak_flops, traffic/peak_bw)``. 14 Published for reference GPUs; anyone can recompute for their own. 15 16**Tier 3 — Measured latency (hardware-specific):** 17 Wall-clock p50/p95/p99, peak memory per backend, throughput. 18 Always accompanied by a :class:`SystemCard`. 19 20FLOPs convention (for reproducibility): 21 - Single forward pass, batch_size=1 22 - Input resolution: task default or caller-specified 23 - Precision: FP32 for counting (hardware-agnostic) 24 - API-only models: report None (shown as "N/A (API)" in tables) 25""" 26 27from __future__ import annotations 28 29import platform 30from dataclasses import dataclass, field 31from typing import Any, Dict, List, Sequence, Tuple 32 33# =========================================================================== # 34# Tier 1: Hardware-agnostic model properties 35# =========================================================================== # 36 37 38@dataclass 39class EfficiencyMetadata: 40 """Hardware-agnostic efficiency metadata for a model. 41 42 Tier-1 fields (identical on any hardware): 43 ``params_m``, ``flops_g``, ``macs_g``, ``memory_traffic_gb``, 44 ``arithmetic_intensity``. 45 46 Tier-3 fields (hardware-specific, filled post-run): 47 ``latency_*``, ``peak_*_mb``. 48 """ 49 50 # --- Tier 1: hardware-agnostic ---------------------------------------- 51 params_m: float | None = None # Total parameters in millions 52 flops_g: float | None = None # FLOPs (giga) at batch=1, FP32 53 macs_g: float | None = None # MACs (giga) = FLOPs / 2 54 actmem_gb_fp16: float | None = None # Activation memory at FP16 (GB) 55 memory_traffic_gb: float | None = None # Estimated DRAM read+write (GB) 56 arithmetic_intensity: float | None = None # FLOPs / Bytes (FLOP/Byte) 57 58 latency_ms_per_sample: float | None = None # Filled post-run by the task runner 59 model_type: str = "local" # "local" | "api" 60 notes: str = "" # e.g. precision mode, resolution override 61 62 # --- Tier 2: roofline bounds (computed on demand) --------------------- 63 roofline: Dict[str, "RooflineBound"] | None = None 64 65 # --- Tier 3: measured (hardware-specific) ----------------------------- 66 peak_cpu_mb: float | None = None 67 peak_cuda_mb: float | None = None 68 peak_mps_mb: float | None = None 69 70 latency_p50_ms: float | None = None 71 latency_p95_ms: float | None = None 72 latency_p99_ms: float | None = None 73 74 # --- System card (attached when measured) ------------------------------ 75 system_card: "SystemCard | None" = None 76 77 def derive_tier1(self) -> None: 78 """Fill in derived Tier-1 fields from primary measurements. 79 80 Call after ``flops_g`` and ``params_m`` are populated. 81 """ 82 if self.flops_g is not None and self.macs_g is None: 83 self.macs_g = round(self.flops_g / 2.0, 3) 84 85 if self.memory_traffic_gb is None and self.params_m is not None: 86 # Conservative estimate: read all params (FP32 = 4 bytes) once 87 # + read/write activations ~ 2x param bytes. This is a 88 # floor; actual traffic depends on layer structure and caching. 89 param_bytes = self.params_m * 1e6 * 4 # FP32 90 self.memory_traffic_gb = round(param_bytes * 3 / 1e9, 4) 91 92 if ( 93 self.arithmetic_intensity is None 94 and self.flops_g is not None 95 and self.memory_traffic_gb is not None 96 and self.memory_traffic_gb > 0 97 ): 98 total_flops = self.flops_g * 1e9 99 total_bytes = self.memory_traffic_gb * 1e9 100 self.arithmetic_intensity = round(total_flops / total_bytes, 2) 101 102 def compute_roofline( 103 self, 104 specs: Sequence["GPUSpec"] | None = None, 105 ) -> Dict[str, "RooflineBound"]: 106 """Compute Tier-2 roofline latency bounds for reference GPUs. 107 108 Parameters 109 ---------- 110 specs : sequence of GPUSpec, optional 111 GPU specifications to compute bounds for. Defaults to 112 :data:`REFERENCE_GPUS` (A100, RTX 4090, Jetson Orin). 113 114 Returns 115 ------- 116 dict mapping GPU name to :class:`RooflineBound`. 117 """ 118 if specs is None: 119 specs = list(REFERENCE_GPUS.values()) 120 121 self.derive_tier1() 122 if self.flops_g is None or self.memory_traffic_gb is None: 123 return {} 124 125 bounds: Dict[str, RooflineBound] = {} 126 for gpu in specs: 127 bounds[gpu.name] = RooflineBound.from_model_and_gpu( 128 flops_g=self.flops_g, 129 traffic_gb=self.memory_traffic_gb, 130 gpu=gpu, 131 ) 132 self.roofline = bounds 133 return bounds 134 135 def to_table_row(self) -> dict: 136 """Produce result-table-ready dict (None -> 'N/A (API)' for API models).""" 137 na = "N/A (API)" if self.model_type == "api" else None 138 row = { 139 "type": self.model_type, 140 "params_m": self.params_m if self.params_m is not None else na, 141 "flops_g": self.flops_g if self.flops_g is not None else na, 142 "macs_g": self.macs_g if self.macs_g is not None else na, 143 "memory_traffic_gb": self.memory_traffic_gb, 144 "arithmetic_intensity": self.arithmetic_intensity, 145 "actmem_gb_fp16": self.actmem_gb_fp16, 146 "latency_ms_per_sample": self.latency_ms_per_sample, 147 "latency_p50_ms": self.latency_p50_ms, 148 "latency_p95_ms": self.latency_p95_ms, 149 "latency_p99_ms": self.latency_p99_ms, 150 "peak_cpu_mb": self.peak_cpu_mb, 151 "peak_cuda_mb": self.peak_cuda_mb, 152 "peak_mps_mb": self.peak_mps_mb, 153 } 154 if self.roofline: 155 for gpu_name, bound in self.roofline.items(): 156 key = gpu_name.lower().replace(" ", "_").replace("-", "_") 157 row[f"roofline_{key}_ms"] = bound.latency_ms 158 row[f"bottleneck_{key}"] = bound.bottleneck 159 if self.system_card is not None: 160 row["system_card"] = self.system_card.to_dict() 161 return row 162 163 164# =========================================================================== # 165# Tier 2: GPU specs + Roofline bounds 166# =========================================================================== # 167 168 169@dataclass(frozen=True) 170class GPUSpec: 171 """Hardware specification for a reference GPU. 172 173 Used to compute roofline latency bounds. All numbers are for the 174 precision specified by ``precision`` (default FP32). 175 176 Attributes 177 ---------- 178 name : str 179 Human-readable name, e.g. ``"A100-80GB"``. 180 peak_tflops : float 181 Peak throughput in tera-FLOP/s at the stated precision. 182 memory_bw_gbps : float 183 Peak DRAM bandwidth in GB/s. 184 memory_gb : float 185 Total device memory in GB. 186 precision : str 187 Precision these specs refer to (``"fp32"``, ``"fp16"``, etc.). 188 """ 189 190 name: str 191 peak_tflops: float 192 memory_bw_gbps: float 193 memory_gb: float 194 precision: str = "fp32" 195 196 197#: Reference GPUs covering datacenter, workstation, and edge-robot scenarios. 198#: Sources: official NVIDIA data sheets. 199REFERENCE_GPUS: Dict[str, GPUSpec] = { 200 "A100-80GB": GPUSpec( 201 name="A100-80GB", 202 peak_tflops=19.5, 203 memory_bw_gbps=2039.0, 204 memory_gb=80.0, 205 ), 206 "RTX 4090": GPUSpec( 207 name="RTX 4090", 208 peak_tflops=82.6, 209 memory_bw_gbps=1008.0, 210 memory_gb=24.0, 211 ), 212 "Jetson Orin 64GB": GPUSpec( 213 name="Jetson Orin 64GB", 214 peak_tflops=5.3, 215 memory_bw_gbps=204.8, 216 memory_gb=64.0, 217 precision="fp32", 218 ), 219} 220 221 222@dataclass(frozen=True) 223class RooflineBound: 224 """Roofline-model latency bound for one model on one GPU. 225 226 The roofline model (Williams et al., 2009) gives a lower-bound 227 latency based on whichever of compute or memory bandwidth is the 228 bottleneck: 229 230 .. math:: 231 232 T_{\\text{roofline}} = \\max\\!\\Big( 233 \\frac{\\text{FLOPs}}{\\text{peak FLOP/s}},\\; 234 \\frac{\\text{Traffic}}{\\text{peak BW}} 235 \\Big) 236 237 Attributes 238 ---------- 239 gpu_name : str 240 compute_ms : float 241 Time if fully compute-bound (FLOPs / peak_tflops). 242 memory_ms : float 243 Time if fully memory-bound (traffic / peak_bw). 244 latency_ms : float 245 ``max(compute_ms, memory_ms)`` — the roofline lower bound. 246 bottleneck : str 247 ``"compute"`` or ``"memory"``. 248 """ 249 250 gpu_name: str 251 compute_ms: float 252 memory_ms: float 253 latency_ms: float 254 bottleneck: str 255 256 @classmethod 257 def from_model_and_gpu( 258 cls, 259 flops_g: float, 260 traffic_gb: float, 261 gpu: GPUSpec, 262 ) -> "RooflineBound": 263 """Compute the roofline bound for a model on a GPU. 264 265 Parameters 266 ---------- 267 flops_g : float 268 Model forward-pass FLOPs in giga-ops. 269 traffic_gb : float 270 Estimated DRAM traffic in GB. 271 gpu : GPUSpec 272 Target GPU specification. 273 """ 274 # peak_tflops is tera-FLOP/s = 1e12 FLOP/s 275 # flops_g is giga-FLOP = 1e9 FLOP 276 # compute_s = (flops_g * 1e9) / (peak_tflops * 1e12) 277 # = flops_g / (peak_tflops * 1e3) 278 compute_s = flops_g / (gpu.peak_tflops * 1e3) 279 memory_s = traffic_gb / gpu.memory_bw_gbps 280 281 compute_ms = round(compute_s * 1e3, 4) 282 memory_ms = round(memory_s * 1e3, 4) 283 latency_ms = round(max(compute_ms, memory_ms), 4) 284 bottleneck = "compute" if compute_ms >= memory_ms else "memory" 285 286 return cls( 287 gpu_name=gpu.name, 288 compute_ms=compute_ms, 289 memory_ms=memory_ms, 290 latency_ms=latency_ms, 291 bottleneck=bottleneck, 292 ) 293 294 295# =========================================================================== # 296# Tier 3: System Card (attached to measured results) 297# =========================================================================== # 298 299 300@dataclass 301class SystemCard: 302 """Full hardware/software description for reproducibility. 303 304 Attached to every set of Tier-3 measurements so readers know 305 exactly what hardware the numbers were collected on. Follows the 306 MLPerf convention of reporting system details alongside results. 307 """ 308 309 gpu_name: str = "" 310 gpu_memory_gb: float = 0.0 311 gpu_count: int = 1 312 cpu_name: str = "" 313 ram_gb: float = 0.0 314 cuda_version: str = "" 315 pytorch_version: str = "" 316 python_version: str = "" 317 precision: str = "fp32" 318 batch_size: int = 1 319 input_resolution: str = "640x480" 320 os: str = "" 321 driver_version: str = "" 322 323 @classmethod 324 def auto_detect( 325 cls, 326 precision: str = "fp32", 327 batch_size: int = 1, 328 input_resolution: str = "640x480", 329 ) -> "SystemCard": 330 """Best-effort auto-detection of the current system. 331 332 Fills every field it can; leaves empty strings for fields 333 that require libraries not installed (e.g. torch for CUDA info). 334 """ 335 card = cls( 336 precision=precision, 337 batch_size=batch_size, 338 input_resolution=input_resolution, 339 python_version=platform.python_version(), 340 os=f"{platform.system()} {platform.release()}", 341 ) 342 343 # CPU 344 card.cpu_name = platform.processor() or platform.machine() 345 346 # RAM 347 try: 348 import psutil # noqa: PLC0415 349 350 card.ram_gb = round(psutil.virtual_memory().total / (1024**3), 1) 351 except ImportError: 352 pass 353 354 # GPU (torch) 355 try: 356 import torch # noqa: PLC0415 357 358 card.pytorch_version = torch.__version__ 359 if torch.cuda.is_available(): 360 props = torch.cuda.get_device_properties(0) 361 card.gpu_name = props.name 362 card.gpu_memory_gb = round(props.total_memory / (1024**3), 1) 363 card.gpu_count = torch.cuda.device_count() 364 card.cuda_version = torch.version.cuda or "" 365 # Driver version from nvidia-smi is not in torch; leave blank 366 # unless the user fills it manually. 367 except ImportError: 368 pass 369 370 return card 371 372 def to_dict(self) -> Dict[str, Any]: 373 return { 374 "gpu_name": self.gpu_name, 375 "gpu_memory_gb": self.gpu_memory_gb, 376 "gpu_count": self.gpu_count, 377 "cpu_name": self.cpu_name, 378 "ram_gb": self.ram_gb, 379 "cuda_version": self.cuda_version, 380 "pytorch_version": self.pytorch_version, 381 "python_version": self.python_version, 382 "precision": self.precision, 383 "batch_size": self.batch_size, 384 "input_resolution": self.input_resolution, 385 "os": self.os, 386 "driver_version": self.driver_version, 387 } 388 389 def summary(self) -> str: 390 """One-line summary for table footers / report headers.""" 391 parts = [] 392 if self.gpu_name: 393 parts.append(f"{self.gpu_name} ({self.gpu_memory_gb}GB)") 394 if self.cpu_name: 395 parts.append(self.cpu_name) 396 parts.append(f"PyTorch {self.pytorch_version}" if self.pytorch_version else "") 397 parts.append(f"CUDA {self.cuda_version}" if self.cuda_version else "") 398 parts.append(self.precision) 399 return " | ".join(p for p in parts if p) 400 401 402# =========================================================================== # 403# Memory traffic estimation 404# =========================================================================== # 405 406 407def estimate_memory_traffic_gb( 408 model: Any, 409 input_shape: Tuple[int, ...] = (3, 480, 640), 410 device: str = "cpu", 411 precision_bytes: int = 4, 412) -> float | None: 413 """Estimate total DRAM traffic (GB) for a single forward pass. 414 415 Strategy (in order of preference): 416 417 1. **torch.profiler** with ``profile_memory=True`` — most accurate, 418 captures actual CUDA memory events. Requires torch >= 2.0 and 419 a CUDA device. 420 2. **Analytical estimate** — conservative lower bound: 421 ``param_bytes + 2 * activation_bytes``. Works on any backend. 422 423 Parameters 424 ---------- 425 model : nn.Module or similar 426 The model to profile. 427 input_shape : tuple 428 ``(C, H, W)`` — batch dim is prepended. 429 device : str 430 Device for the dummy input. 431 precision_bytes : int 432 Bytes per element (4 for FP32, 2 for FP16/BF16). 433 434 Returns 435 ------- 436 float or None 437 Estimated DRAM traffic in GB, or None if estimation fails. 438 """ 439 # Strategy 1: torch.profiler 440 try: 441 import torch # noqa: PLC0415 442 443 if device != "cpu" and torch.cuda.is_available(): 444 dummy = torch.zeros(1, *input_shape, device=device) 445 with torch.profiler.profile( 446 activities=[ 447 torch.profiler.ProfilerActivity.CPU, 448 torch.profiler.ProfilerActivity.CUDA, 449 ], 450 profile_memory=True, 451 record_shapes=True, 452 ) as prof: 453 with torch.no_grad(): 454 model(dummy) 455 456 # Sum all CUDA memory allocations as a proxy for traffic. 457 # This overcounts (includes temporaries freed mid-forward) 458 # but is more realistic than the analytical estimate. 459 total_bytes = 0 460 for evt in prof.key_averages(): 461 mem = getattr(evt, "cuda_memory_usage", 0) or 0 462 if mem > 0: 463 total_bytes += mem 464 if total_bytes > 0: 465 return round(total_bytes / 1e9, 4) 466 except (ImportError, RuntimeError, AttributeError): 467 pass 468 469 # Strategy 2: analytical estimate 470 try: 471 params_bytes = sum(p.numel() * p.element_size() for p in model.parameters()) 472 # Heuristic: activations ~ 2x param memory for typical vision 473 # models (each layer reads previous activation, writes new one). 474 activation_bytes = params_bytes * 2 475 total = params_bytes + activation_bytes # read params + r/w activations 476 return round(total / 1e9, 4) 477 except (AttributeError, TypeError): 478 pass 479 480 return None 481 482 483# --------------------------------------------------------------------------- # 484# Memory profiler — CPU / CUDA / MPS backends 485# --------------------------------------------------------------------------- # 486 487 488def _cpu_peak_rss_mb() -> float | None: 489 """Current process peak RSS in MB; None if neither psutil nor resource 490 can produce a number.""" 491 # psutil first (cross-platform, no import of ``resource``). 492 try: 493 import psutil # noqa: PLC0415 494 495 proc = psutil.Process() 496 info = proc.memory_info() 497 # ``peak_wset`` (Windows) or ``peak_rss`` (macOS 12+) are rare; 498 # fall back to ``rss`` if the peak variant isn't available. 499 peak = getattr(info, "peak_wset", None) or getattr(info, "peak_rss", None) 500 if peak is not None: 501 return round(peak / (1024 * 1024), 3) 502 return round(info.rss / (1024 * 1024), 3) 503 except ImportError: 504 pass 505 try: 506 import resource # noqa: PLC0415 — POSIX only 507 508 usage = resource.getrusage(resource.RUSAGE_SELF) 509 # ``ru_maxrss`` units: KB on Linux, bytes on macOS. We normalise 510 # by checking the magnitude — a process using hundreds of GB is 511 # exceedingly unlikely, so > 10M means the value is in bytes. 512 maxrss = usage.ru_maxrss 513 if maxrss > 10_000_000: 514 return round(maxrss / (1024 * 1024), 3) 515 return round(maxrss / 1024, 3) 516 except ImportError: 517 return None 518 519 520def _cuda_peak_alloc_mb() -> float | None: 521 try: 522 import torch # noqa: PLC0415 523 except ImportError: 524 return None 525 if not hasattr(torch, "cuda") or not torch.cuda.is_available(): 526 return None 527 try: 528 return round(torch.cuda.max_memory_allocated() / (1024 * 1024), 3) 529 except Exception: 530 return None 531 532 533def _mps_current_alloc_mb() -> float | None: 534 """Apple Silicon MPS memory — uses ``torch.mps.current_allocated_memory``. 535 536 ``torch.mps`` has no ``max_memory_allocated`` until torch 2.3; we 537 fall back to the current alloc so we still surface *something*. 538 """ 539 try: 540 import torch # noqa: PLC0415 541 except ImportError: 542 return None 543 mps_mod = getattr(torch, "mps", None) 544 if mps_mod is None: 545 return None 546 try: 547 if not torch.backends.mps.is_available(): # type: ignore[attr-defined] 548 return None 549 except AttributeError: 550 return None 551 fn = getattr(mps_mod, "current_allocated_memory", None) 552 if fn is None: 553 return None 554 try: 555 return round(fn() / (1024 * 1024), 3) 556 except Exception: 557 return None 558 559 560@dataclass 561class MemoryProfiler: 562 """Hardware-agnostic memory profiler covering CPU / CUDA / MPS. 563 564 Usage:: 565 566 mp = MemoryProfiler().reset() 567 # ... run inference ... 568 peaks = mp.peaks() # {"cpu_mb": ..., "cuda_mb": ..., "mps_mb": ...} 569 570 Each backend returns ``None`` when it cannot be sampled (library 571 missing, backend unavailable, measurement unsupported). The caller 572 is responsible for deciding whether ``None`` is informative 573 (report as ``"n/a"``) or an error. 574 """ 575 576 sample_cpu: bool = True 577 sample_cuda: bool = True 578 sample_mps: bool = True 579 580 def reset(self) -> "MemoryProfiler": 581 """Clear accumulated peaks on every enabled backend.""" 582 if self.sample_cuda: 583 try: 584 import torch # noqa: PLC0415 585 586 if hasattr(torch, "cuda") and torch.cuda.is_available(): 587 torch.cuda.reset_peak_memory_stats() 588 except ImportError: 589 pass 590 # CPU and MPS backends don't expose a reset primitive — the 591 # sampled value is "current" (MPS) or "peak-since-process-start" 592 # (CPU ru_maxrss). Document the limitation in the docstring. 593 return self 594 595 def peaks(self) -> dict: 596 return { 597 "cpu_mb": _cpu_peak_rss_mb() if self.sample_cpu else None, 598 "cuda_mb": _cuda_peak_alloc_mb() if self.sample_cuda else None, 599 "mps_mb": _mps_current_alloc_mb() if self.sample_mps else None, 600 } 601 602 603# --------------------------------------------------------------------------- # 604# Latency profiler — percentiles + mean 605# --------------------------------------------------------------------------- # 606 607 608@dataclass 609class LatencyProfiler: 610 """Accumulate per-sample latencies and report percentiles. 611 612 Percentiles use linear interpolation (``numpy.percentile`` default) 613 and skip a user-configurable number of warmup samples so the 614 reported numbers reflect steady-state throughput rather than the 615 first-batch JIT compile. 616 """ 617 618 warmup: int = 1 619 _samples_ms: List[float] = field(default_factory=list) 620 621 def add_sample_seconds(self, seconds: float) -> None: 622 self._samples_ms.append(float(seconds) * 1000.0) 623 624 def add_batch_seconds(self, batch_seconds: float, batch_size: int) -> None: 625 """Amortise a batch timing over ``batch_size`` sample entries.""" 626 if batch_size <= 0: 627 return 628 per = batch_seconds / batch_size 629 self._samples_ms.extend([per * 1000.0] * batch_size) 630 631 def percentiles(self) -> dict: 632 """Return ``{p50_ms, p95_ms, p99_ms, mean_ms}`` in milliseconds. 633 634 Returns ``{k: None for k in keys}`` when fewer than one sample 635 exists after warmup trimming. 636 """ 637 trimmed = self._samples_ms[self.warmup :] or self._samples_ms 638 if not trimmed: 639 return {"p50_ms": None, "p95_ms": None, "p99_ms": None, "mean_ms": None} 640 try: 641 import numpy as np # noqa: PLC0415 642 643 p50, p95, p99 = np.percentile(trimmed, [50, 95, 99]) 644 return { 645 "p50_ms": round(float(p50), 3), 646 "p95_ms": round(float(p95), 3), 647 "p99_ms": round(float(p99), 3), 648 "mean_ms": round(float(np.mean(trimmed)), 3), 649 } 650 except ImportError: 651 trimmed_sorted = sorted(trimmed) 652 n = len(trimmed_sorted) 653 654 def _q(frac: float) -> float: 655 return trimmed_sorted[min(n - 1, int(round(frac * (n - 1))))] 656 657 return { 658 "p50_ms": round(_q(0.50), 3), 659 "p95_ms": round(_q(0.95), 3), 660 "p99_ms": round(_q(0.99), 3), 661 "mean_ms": round(sum(trimmed) / n, 3), 662 } 663 664 def samples_ms(self) -> Sequence[float]: 665 """Return the raw latency samples (post-warmup not trimmed).""" 666 return list(self._samples_ms) 667 668 669def count_parameters(model: Any) -> float: 670 """Count trainable parameters in millions. 671 672 Supports PyTorch nn.Module and any object with a ``parameters()`` method. 673 Returns None if the model type is not supported. 674 """ 675 try: 676 params = sum(p.numel() for p in model.parameters() if p.requires_grad) 677 return round(params / 1e6, 3) 678 except AttributeError: 679 pass 680 681 # JAX / Flax: model may expose a ``params`` pytree 682 try: 683 import jax 684 685 leaves = jax.tree_util.tree_leaves(model.params) 686 params = sum(leaf.size for leaf in leaves) 687 return round(params / 1e6, 3) 688 except (AttributeError, ImportError): 689 pass 690 691 return None 692 693 694def count_flops_torch( 695 model: Any, 696 input_shape: Tuple[int, ...], 697 device: str = "cpu", 698) -> float | None: 699 """Estimate FLOPs (giga) using ``torch.utils.flop_counter`` (PyTorch >= 2.1). 700 701 Falls back to ``fvcore`` if available. 702 703 **OOM-safe**: when the FLOP counter runs out of CUDA memory (the 704 dispatch metadata alone can consume ~3 GB on top of the model), the 705 function automatically retries on CPU. The FLOP count is identical 706 regardless of device — it is a graph property, not a runtime one. 707 708 Args: 709 model: PyTorch nn.Module. 710 input_shape: (C, H, W) — batch dimension is added automatically. 711 device: device string for the dummy input tensor. 712 713 Returns: 714 FLOPs in giga-ops, or None if neither backend is available. 715 """ 716 717 def _try_flop_count(mod: Any, dev: str) -> float | None: 718 try: 719 import torch # noqa: PLC0415 720 except ImportError: 721 return None 722 723 dummy = torch.zeros(1, *input_shape, device=dev) 724 725 # PyTorch 2.1+ native counter 726 try: 727 from torch.utils.flop_counter import FlopCounterMode # noqa: PLC0415 728 729 with torch.no_grad(): 730 with FlopCounterMode(display=False) as fcm: 731 mod(dummy) 732 total = _sum_nested_counts(fcm.get_flop_counts()) 733 if total > 0: 734 return round(total / 1e9, 3) 735 except ImportError: 736 pass 737 738 # fvcore fallback 739 try: 740 from fvcore.nn import FlopCountAnalysis # noqa: PLC0415 741 742 flops = FlopCountAnalysis(mod, dummy) 743 return round(flops.total() / 1e9, 3) 744 except ImportError: 745 pass 746 747 return None 748 749 # Helper for nested FlopCounterMode dicts (same logic as runner.py) 750 def _sum_nested_counts(node: Any) -> int: 751 if isinstance(node, dict): 752 return sum(_sum_nested_counts(v) for v in node.values()) 753 try: 754 return int(node) 755 except (TypeError, ValueError): 756 return 0 757 758 # Try on the requested device first 759 try: 760 result = _try_flop_count(model, device) 761 if result is not None: 762 return result 763 except RuntimeError as exc: 764 # Catch CUDA OOM and fall through to CPU retry 765 if "out of memory" not in str(exc).lower(): 766 return None 767 768 # CPU fallback: FLOPs are a graph property — the count is device- 769 # independent. Move the model to CPU, count, then move it back. 770 if device != "cpu": 771 try: 772 original_device = next(model.parameters()).device 773 model_cpu = model.cpu() 774 result = _try_flop_count(model_cpu, "cpu") 775 model.to(original_device) 776 return result 777 except Exception: 778 # If anything goes wrong restoring the model, still return 779 # None rather than crashing the caller. 780 try: 781 model.to(device) 782 except Exception: 783 pass 784 return None 785 786 return None 787 788 789def profile_model( 790 model: Any, 791 input_shape: Tuple[int, ...] = (3, 480, 640), 792 device: str = "cpu", 793 model_type: str = "local", 794 notes: str = "", 795 compute_roofline: bool = True, 796 attach_system_card: bool = True, 797 precision: str = "fp32", 798) -> EfficiencyMetadata: 799 """Auto-profile a model and return EfficiencyMetadata. 800 801 Fills all three tiers of efficiency reporting: 802 803 - **Tier 1**: params, FLOPs, MACs, memory traffic, arithmetic 804 intensity (hardware-agnostic). 805 - **Tier 2**: roofline latency bounds for reference GPUs 806 (hardware-parametric). 807 - **Tier 3**: system card for the current machine 808 (hardware-specific; measured latency is added later by the runner). 809 810 Parameters 811 ---------- 812 model : nn.Module or similar 813 The model to profile. 814 input_shape : tuple 815 ``(C, H, W)`` for FLOPs counting. Default 640x480 RGB. 816 device : str 817 Device for the dummy input tensor. 818 model_type : str 819 ``"local"`` or ``"api"``. 820 notes : str 821 Free-text notes (e.g. ``"ViT-L/14, FP16 inference"``). 822 compute_roofline : bool 823 If True, compute Tier-2 roofline bounds for reference GPUs. 824 attach_system_card : bool 825 If True, auto-detect and attach a SystemCard. 826 precision : str 827 Precision string for the system card (``"fp32"``, ``"fp16"``). 828 829 Returns 830 ------- 831 EfficiencyMetadata 832 With Tier 1/2/3 fields populated where possible. 833 """ 834 if model_type == "api": 835 return EfficiencyMetadata(model_type="api", notes=notes) 836 837 params_m = count_parameters(model) 838 flops_g = count_flops_torch(model, input_shape, device=device) 839 traffic_gb = estimate_memory_traffic_gb(model, input_shape, device=device) 840 841 h, w = input_shape[-2], input_shape[-1] 842 em = EfficiencyMetadata( 843 params_m=params_m, 844 flops_g=flops_g, 845 memory_traffic_gb=traffic_gb, 846 model_type=model_type, 847 notes=notes, 848 ) 849 850 # Derive Tier 1 (MACs, arithmetic intensity) 851 em.derive_tier1() 852 853 # Tier 2: roofline bounds 854 if compute_roofline: 855 em.compute_roofline() 856 857 # Tier 3: system card 858 if attach_system_card: 859 em.system_card = SystemCard.auto_detect( 860 precision=precision, 861 input_resolution=f"{w}x{h}", 862 ) 863 864 return em
@dataclass
class
EfficiencyMetadata:
39@dataclass 40class EfficiencyMetadata: 41 """Hardware-agnostic efficiency metadata for a model. 42 43 Tier-1 fields (identical on any hardware): 44 ``params_m``, ``flops_g``, ``macs_g``, ``memory_traffic_gb``, 45 ``arithmetic_intensity``. 46 47 Tier-3 fields (hardware-specific, filled post-run): 48 ``latency_*``, ``peak_*_mb``. 49 """ 50 51 # --- Tier 1: hardware-agnostic ---------------------------------------- 52 params_m: float | None = None # Total parameters in millions 53 flops_g: float | None = None # FLOPs (giga) at batch=1, FP32 54 macs_g: float | None = None # MACs (giga) = FLOPs / 2 55 actmem_gb_fp16: float | None = None # Activation memory at FP16 (GB) 56 memory_traffic_gb: float | None = None # Estimated DRAM read+write (GB) 57 arithmetic_intensity: float | None = None # FLOPs / Bytes (FLOP/Byte) 58 59 latency_ms_per_sample: float | None = None # Filled post-run by the task runner 60 model_type: str = "local" # "local" | "api" 61 notes: str = "" # e.g. precision mode, resolution override 62 63 # --- Tier 2: roofline bounds (computed on demand) --------------------- 64 roofline: Dict[str, "RooflineBound"] | None = None 65 66 # --- Tier 3: measured (hardware-specific) ----------------------------- 67 peak_cpu_mb: float | None = None 68 peak_cuda_mb: float | None = None 69 peak_mps_mb: float | None = None 70 71 latency_p50_ms: float | None = None 72 latency_p95_ms: float | None = None 73 latency_p99_ms: float | None = None 74 75 # --- System card (attached when measured) ------------------------------ 76 system_card: "SystemCard | None" = None 77 78 def derive_tier1(self) -> None: 79 """Fill in derived Tier-1 fields from primary measurements. 80 81 Call after ``flops_g`` and ``params_m`` are populated. 82 """ 83 if self.flops_g is not None and self.macs_g is None: 84 self.macs_g = round(self.flops_g / 2.0, 3) 85 86 if self.memory_traffic_gb is None and self.params_m is not None: 87 # Conservative estimate: read all params (FP32 = 4 bytes) once 88 # + read/write activations ~ 2x param bytes. This is a 89 # floor; actual traffic depends on layer structure and caching. 90 param_bytes = self.params_m * 1e6 * 4 # FP32 91 self.memory_traffic_gb = round(param_bytes * 3 / 1e9, 4) 92 93 if ( 94 self.arithmetic_intensity is None 95 and self.flops_g is not None 96 and self.memory_traffic_gb is not None 97 and self.memory_traffic_gb > 0 98 ): 99 total_flops = self.flops_g * 1e9 100 total_bytes = self.memory_traffic_gb * 1e9 101 self.arithmetic_intensity = round(total_flops / total_bytes, 2) 102 103 def compute_roofline( 104 self, 105 specs: Sequence["GPUSpec"] | None = None, 106 ) -> Dict[str, "RooflineBound"]: 107 """Compute Tier-2 roofline latency bounds for reference GPUs. 108 109 Parameters 110 ---------- 111 specs : sequence of GPUSpec, optional 112 GPU specifications to compute bounds for. Defaults to 113 :data:`REFERENCE_GPUS` (A100, RTX 4090, Jetson Orin). 114 115 Returns 116 ------- 117 dict mapping GPU name to :class:`RooflineBound`. 118 """ 119 if specs is None: 120 specs = list(REFERENCE_GPUS.values()) 121 122 self.derive_tier1() 123 if self.flops_g is None or self.memory_traffic_gb is None: 124 return {} 125 126 bounds: Dict[str, RooflineBound] = {} 127 for gpu in specs: 128 bounds[gpu.name] = RooflineBound.from_model_and_gpu( 129 flops_g=self.flops_g, 130 traffic_gb=self.memory_traffic_gb, 131 gpu=gpu, 132 ) 133 self.roofline = bounds 134 return bounds 135 136 def to_table_row(self) -> dict: 137 """Produce result-table-ready dict (None -> 'N/A (API)' for API models).""" 138 na = "N/A (API)" if self.model_type == "api" else None 139 row = { 140 "type": self.model_type, 141 "params_m": self.params_m if self.params_m is not None else na, 142 "flops_g": self.flops_g if self.flops_g is not None else na, 143 "macs_g": self.macs_g if self.macs_g is not None else na, 144 "memory_traffic_gb": self.memory_traffic_gb, 145 "arithmetic_intensity": self.arithmetic_intensity, 146 "actmem_gb_fp16": self.actmem_gb_fp16, 147 "latency_ms_per_sample": self.latency_ms_per_sample, 148 "latency_p50_ms": self.latency_p50_ms, 149 "latency_p95_ms": self.latency_p95_ms, 150 "latency_p99_ms": self.latency_p99_ms, 151 "peak_cpu_mb": self.peak_cpu_mb, 152 "peak_cuda_mb": self.peak_cuda_mb, 153 "peak_mps_mb": self.peak_mps_mb, 154 } 155 if self.roofline: 156 for gpu_name, bound in self.roofline.items(): 157 key = gpu_name.lower().replace(" ", "_").replace("-", "_") 158 row[f"roofline_{key}_ms"] = bound.latency_ms 159 row[f"bottleneck_{key}"] = bound.bottleneck 160 if self.system_card is not None: 161 row["system_card"] = self.system_card.to_dict() 162 return row
Hardware-agnostic efficiency metadata for a model.
Tier-1 fields (identical on any hardware):
params_m, flops_g, macs_g, memory_traffic_gb,
arithmetic_intensity.
Tier-3 fields (hardware-specific, filled post-run):
latency_*, peak_*_mb.
EfficiencyMetadata( params_m: float | None = None, flops_g: float | None = None, macs_g: float | None = None, actmem_gb_fp16: float | None = None, memory_traffic_gb: float | None = None, arithmetic_intensity: float | None = None, latency_ms_per_sample: float | None = None, model_type: str = 'local', notes: str = '', roofline: Optional[Dict[str, RooflineBound]] = None, peak_cpu_mb: float | None = None, peak_cuda_mb: float | None = None, peak_mps_mb: float | None = None, latency_p50_ms: float | None = None, latency_p95_ms: float | None = None, latency_p99_ms: float | None = None, system_card: SystemCard | None = None)
def
derive_tier1(self) -> None:
78 def derive_tier1(self) -> None: 79 """Fill in derived Tier-1 fields from primary measurements. 80 81 Call after ``flops_g`` and ``params_m`` are populated. 82 """ 83 if self.flops_g is not None and self.macs_g is None: 84 self.macs_g = round(self.flops_g / 2.0, 3) 85 86 if self.memory_traffic_gb is None and self.params_m is not None: 87 # Conservative estimate: read all params (FP32 = 4 bytes) once 88 # + read/write activations ~ 2x param bytes. This is a 89 # floor; actual traffic depends on layer structure and caching. 90 param_bytes = self.params_m * 1e6 * 4 # FP32 91 self.memory_traffic_gb = round(param_bytes * 3 / 1e9, 4) 92 93 if ( 94 self.arithmetic_intensity is None 95 and self.flops_g is not None 96 and self.memory_traffic_gb is not None 97 and self.memory_traffic_gb > 0 98 ): 99 total_flops = self.flops_g * 1e9 100 total_bytes = self.memory_traffic_gb * 1e9 101 self.arithmetic_intensity = round(total_flops / total_bytes, 2)
103 def compute_roofline( 104 self, 105 specs: Sequence["GPUSpec"] | None = None, 106 ) -> Dict[str, "RooflineBound"]: 107 """Compute Tier-2 roofline latency bounds for reference GPUs. 108 109 Parameters 110 ---------- 111 specs : sequence of GPUSpec, optional 112 GPU specifications to compute bounds for. Defaults to 113 :data:`REFERENCE_GPUS` (A100, RTX 4090, Jetson Orin). 114 115 Returns 116 ------- 117 dict mapping GPU name to :class:`RooflineBound`. 118 """ 119 if specs is None: 120 specs = list(REFERENCE_GPUS.values()) 121 122 self.derive_tier1() 123 if self.flops_g is None or self.memory_traffic_gb is None: 124 return {} 125 126 bounds: Dict[str, RooflineBound] = {} 127 for gpu in specs: 128 bounds[gpu.name] = RooflineBound.from_model_and_gpu( 129 flops_g=self.flops_g, 130 traffic_gb=self.memory_traffic_gb, 131 gpu=gpu, 132 ) 133 self.roofline = bounds 134 return bounds
Compute Tier-2 roofline latency bounds for reference GPUs.
Parameters
specs : sequence of GPUSpec, optional
GPU specifications to compute bounds for. Defaults to
REFERENCE_GPUS (A100, RTX 4090, Jetson Orin).
Returns
dict mapping GPU name to RooflineBound.
def
to_table_row(self) -> dict:
136 def to_table_row(self) -> dict: 137 """Produce result-table-ready dict (None -> 'N/A (API)' for API models).""" 138 na = "N/A (API)" if self.model_type == "api" else None 139 row = { 140 "type": self.model_type, 141 "params_m": self.params_m if self.params_m is not None else na, 142 "flops_g": self.flops_g if self.flops_g is not None else na, 143 "macs_g": self.macs_g if self.macs_g is not None else na, 144 "memory_traffic_gb": self.memory_traffic_gb, 145 "arithmetic_intensity": self.arithmetic_intensity, 146 "actmem_gb_fp16": self.actmem_gb_fp16, 147 "latency_ms_per_sample": self.latency_ms_per_sample, 148 "latency_p50_ms": self.latency_p50_ms, 149 "latency_p95_ms": self.latency_p95_ms, 150 "latency_p99_ms": self.latency_p99_ms, 151 "peak_cpu_mb": self.peak_cpu_mb, 152 "peak_cuda_mb": self.peak_cuda_mb, 153 "peak_mps_mb": self.peak_mps_mb, 154 } 155 if self.roofline: 156 for gpu_name, bound in self.roofline.items(): 157 key = gpu_name.lower().replace(" ", "_").replace("-", "_") 158 row[f"roofline_{key}_ms"] = bound.latency_ms 159 row[f"bottleneck_{key}"] = bound.bottleneck 160 if self.system_card is not None: 161 row["system_card"] = self.system_card.to_dict() 162 return row
Produce result-table-ready dict (None -> 'N/A (API)' for API models).
@dataclass(frozen=True)
class
GPUSpec:
170@dataclass(frozen=True) 171class GPUSpec: 172 """Hardware specification for a reference GPU. 173 174 Used to compute roofline latency bounds. All numbers are for the 175 precision specified by ``precision`` (default FP32). 176 177 Attributes 178 ---------- 179 name : str 180 Human-readable name, e.g. ``"A100-80GB"``. 181 peak_tflops : float 182 Peak throughput in tera-FLOP/s at the stated precision. 183 memory_bw_gbps : float 184 Peak DRAM bandwidth in GB/s. 185 memory_gb : float 186 Total device memory in GB. 187 precision : str 188 Precision these specs refer to (``"fp32"``, ``"fp16"``, etc.). 189 """ 190 191 name: str 192 peak_tflops: float 193 memory_bw_gbps: float 194 memory_gb: float 195 precision: str = "fp32"
Hardware specification for a reference GPU.
Used to compute roofline latency bounds. All numbers are for the
precision specified by precision (default FP32).
Attributes
name : str
Human-readable name, e.g. "A100-80GB".
peak_tflops : float
Peak throughput in tera-FLOP/s at the stated precision.
memory_bw_gbps : float
Peak DRAM bandwidth in GB/s.
memory_gb : float
Total device memory in GB.
precision : str
Precision these specs refer to ("fp32", "fp16", etc.).
REFERENCE_GPUS: Dict[str, GPUSpec] =
{'A100-80GB': GPUSpec(name='A100-80GB', peak_tflops=19.5, memory_bw_gbps=2039.0, memory_gb=80.0, precision='fp32'), 'RTX 4090': GPUSpec(name='RTX 4090', peak_tflops=82.6, memory_bw_gbps=1008.0, memory_gb=24.0, precision='fp32'), 'Jetson Orin 64GB': GPUSpec(name='Jetson Orin 64GB', peak_tflops=5.3, memory_bw_gbps=204.8, memory_gb=64.0, precision='fp32')}
@dataclass(frozen=True)
class
RooflineBound:
223@dataclass(frozen=True) 224class RooflineBound: 225 """Roofline-model latency bound for one model on one GPU. 226 227 The roofline model (Williams et al., 2009) gives a lower-bound 228 latency based on whichever of compute or memory bandwidth is the 229 bottleneck: 230 231 .. math:: 232 233 T_{\\text{roofline}} = \\max\\!\\Big( 234 \\frac{\\text{FLOPs}}{\\text{peak FLOP/s}},\\; 235 \\frac{\\text{Traffic}}{\\text{peak BW}} 236 \\Big) 237 238 Attributes 239 ---------- 240 gpu_name : str 241 compute_ms : float 242 Time if fully compute-bound (FLOPs / peak_tflops). 243 memory_ms : float 244 Time if fully memory-bound (traffic / peak_bw). 245 latency_ms : float 246 ``max(compute_ms, memory_ms)`` — the roofline lower bound. 247 bottleneck : str 248 ``"compute"`` or ``"memory"``. 249 """ 250 251 gpu_name: str 252 compute_ms: float 253 memory_ms: float 254 latency_ms: float 255 bottleneck: str 256 257 @classmethod 258 def from_model_and_gpu( 259 cls, 260 flops_g: float, 261 traffic_gb: float, 262 gpu: GPUSpec, 263 ) -> "RooflineBound": 264 """Compute the roofline bound for a model on a GPU. 265 266 Parameters 267 ---------- 268 flops_g : float 269 Model forward-pass FLOPs in giga-ops. 270 traffic_gb : float 271 Estimated DRAM traffic in GB. 272 gpu : GPUSpec 273 Target GPU specification. 274 """ 275 # peak_tflops is tera-FLOP/s = 1e12 FLOP/s 276 # flops_g is giga-FLOP = 1e9 FLOP 277 # compute_s = (flops_g * 1e9) / (peak_tflops * 1e12) 278 # = flops_g / (peak_tflops * 1e3) 279 compute_s = flops_g / (gpu.peak_tflops * 1e3) 280 memory_s = traffic_gb / gpu.memory_bw_gbps 281 282 compute_ms = round(compute_s * 1e3, 4) 283 memory_ms = round(memory_s * 1e3, 4) 284 latency_ms = round(max(compute_ms, memory_ms), 4) 285 bottleneck = "compute" if compute_ms >= memory_ms else "memory" 286 287 return cls( 288 gpu_name=gpu.name, 289 compute_ms=compute_ms, 290 memory_ms=memory_ms, 291 latency_ms=latency_ms, 292 bottleneck=bottleneck, 293 )
Roofline-model latency bound for one model on one GPU.
The roofline model (Williams et al., 2009) gives a lower-bound latency based on whichever of compute or memory bandwidth is the bottleneck:
$$T_{\text{roofline}} = \max!\Big( \frac{\text{FLOPs}}{\text{peak FLOP/s}},\; \frac{\text{Traffic}}{\text{peak BW}} \Big)$$
Attributes
gpu_name : str
compute_ms : float
Time if fully compute-bound (FLOPs / peak_tflops).
memory_ms : float
Time if fully memory-bound (traffic / peak_bw).
latency_ms : float
max(compute_ms, memory_ms) — the roofline lower bound.
bottleneck : str
"compute" or "memory".
RooflineBound( gpu_name: str, compute_ms: float, memory_ms: float, latency_ms: float, bottleneck: str)
@classmethod
def
from_model_and_gpu( cls, flops_g: float, traffic_gb: float, gpu: GPUSpec) -> RooflineBound:
257 @classmethod 258 def from_model_and_gpu( 259 cls, 260 flops_g: float, 261 traffic_gb: float, 262 gpu: GPUSpec, 263 ) -> "RooflineBound": 264 """Compute the roofline bound for a model on a GPU. 265 266 Parameters 267 ---------- 268 flops_g : float 269 Model forward-pass FLOPs in giga-ops. 270 traffic_gb : float 271 Estimated DRAM traffic in GB. 272 gpu : GPUSpec 273 Target GPU specification. 274 """ 275 # peak_tflops is tera-FLOP/s = 1e12 FLOP/s 276 # flops_g is giga-FLOP = 1e9 FLOP 277 # compute_s = (flops_g * 1e9) / (peak_tflops * 1e12) 278 # = flops_g / (peak_tflops * 1e3) 279 compute_s = flops_g / (gpu.peak_tflops * 1e3) 280 memory_s = traffic_gb / gpu.memory_bw_gbps 281 282 compute_ms = round(compute_s * 1e3, 4) 283 memory_ms = round(memory_s * 1e3, 4) 284 latency_ms = round(max(compute_ms, memory_ms), 4) 285 bottleneck = "compute" if compute_ms >= memory_ms else "memory" 286 287 return cls( 288 gpu_name=gpu.name, 289 compute_ms=compute_ms, 290 memory_ms=memory_ms, 291 latency_ms=latency_ms, 292 bottleneck=bottleneck, 293 )
Compute the roofline bound for a model on a GPU.
Parameters
flops_g : float Model forward-pass FLOPs in giga-ops. traffic_gb : float Estimated DRAM traffic in GB. gpu : GPUSpec Target GPU specification.
@dataclass
class
SystemCard:
301@dataclass 302class SystemCard: 303 """Full hardware/software description for reproducibility. 304 305 Attached to every set of Tier-3 measurements so readers know 306 exactly what hardware the numbers were collected on. Follows the 307 MLPerf convention of reporting system details alongside results. 308 """ 309 310 gpu_name: str = "" 311 gpu_memory_gb: float = 0.0 312 gpu_count: int = 1 313 cpu_name: str = "" 314 ram_gb: float = 0.0 315 cuda_version: str = "" 316 pytorch_version: str = "" 317 python_version: str = "" 318 precision: str = "fp32" 319 batch_size: int = 1 320 input_resolution: str = "640x480" 321 os: str = "" 322 driver_version: str = "" 323 324 @classmethod 325 def auto_detect( 326 cls, 327 precision: str = "fp32", 328 batch_size: int = 1, 329 input_resolution: str = "640x480", 330 ) -> "SystemCard": 331 """Best-effort auto-detection of the current system. 332 333 Fills every field it can; leaves empty strings for fields 334 that require libraries not installed (e.g. torch for CUDA info). 335 """ 336 card = cls( 337 precision=precision, 338 batch_size=batch_size, 339 input_resolution=input_resolution, 340 python_version=platform.python_version(), 341 os=f"{platform.system()} {platform.release()}", 342 ) 343 344 # CPU 345 card.cpu_name = platform.processor() or platform.machine() 346 347 # RAM 348 try: 349 import psutil # noqa: PLC0415 350 351 card.ram_gb = round(psutil.virtual_memory().total / (1024**3), 1) 352 except ImportError: 353 pass 354 355 # GPU (torch) 356 try: 357 import torch # noqa: PLC0415 358 359 card.pytorch_version = torch.__version__ 360 if torch.cuda.is_available(): 361 props = torch.cuda.get_device_properties(0) 362 card.gpu_name = props.name 363 card.gpu_memory_gb = round(props.total_memory / (1024**3), 1) 364 card.gpu_count = torch.cuda.device_count() 365 card.cuda_version = torch.version.cuda or "" 366 # Driver version from nvidia-smi is not in torch; leave blank 367 # unless the user fills it manually. 368 except ImportError: 369 pass 370 371 return card 372 373 def to_dict(self) -> Dict[str, Any]: 374 return { 375 "gpu_name": self.gpu_name, 376 "gpu_memory_gb": self.gpu_memory_gb, 377 "gpu_count": self.gpu_count, 378 "cpu_name": self.cpu_name, 379 "ram_gb": self.ram_gb, 380 "cuda_version": self.cuda_version, 381 "pytorch_version": self.pytorch_version, 382 "python_version": self.python_version, 383 "precision": self.precision, 384 "batch_size": self.batch_size, 385 "input_resolution": self.input_resolution, 386 "os": self.os, 387 "driver_version": self.driver_version, 388 } 389 390 def summary(self) -> str: 391 """One-line summary for table footers / report headers.""" 392 parts = [] 393 if self.gpu_name: 394 parts.append(f"{self.gpu_name} ({self.gpu_memory_gb}GB)") 395 if self.cpu_name: 396 parts.append(self.cpu_name) 397 parts.append(f"PyTorch {self.pytorch_version}" if self.pytorch_version else "") 398 parts.append(f"CUDA {self.cuda_version}" if self.cuda_version else "") 399 parts.append(self.precision) 400 return " | ".join(p for p in parts if p)
Full hardware/software description for reproducibility.
Attached to every set of Tier-3 measurements so readers know exactly what hardware the numbers were collected on. Follows the MLPerf convention of reporting system details alongside results.
SystemCard( gpu_name: str = '', gpu_memory_gb: float = 0.0, gpu_count: int = 1, cpu_name: str = '', ram_gb: float = 0.0, cuda_version: str = '', pytorch_version: str = '', python_version: str = '', precision: str = 'fp32', batch_size: int = 1, input_resolution: str = '640x480', os: str = '', driver_version: str = '')
@classmethod
def
auto_detect( cls, precision: str = 'fp32', batch_size: int = 1, input_resolution: str = '640x480') -> SystemCard:
324 @classmethod 325 def auto_detect( 326 cls, 327 precision: str = "fp32", 328 batch_size: int = 1, 329 input_resolution: str = "640x480", 330 ) -> "SystemCard": 331 """Best-effort auto-detection of the current system. 332 333 Fills every field it can; leaves empty strings for fields 334 that require libraries not installed (e.g. torch for CUDA info). 335 """ 336 card = cls( 337 precision=precision, 338 batch_size=batch_size, 339 input_resolution=input_resolution, 340 python_version=platform.python_version(), 341 os=f"{platform.system()} {platform.release()}", 342 ) 343 344 # CPU 345 card.cpu_name = platform.processor() or platform.machine() 346 347 # RAM 348 try: 349 import psutil # noqa: PLC0415 350 351 card.ram_gb = round(psutil.virtual_memory().total / (1024**3), 1) 352 except ImportError: 353 pass 354 355 # GPU (torch) 356 try: 357 import torch # noqa: PLC0415 358 359 card.pytorch_version = torch.__version__ 360 if torch.cuda.is_available(): 361 props = torch.cuda.get_device_properties(0) 362 card.gpu_name = props.name 363 card.gpu_memory_gb = round(props.total_memory / (1024**3), 1) 364 card.gpu_count = torch.cuda.device_count() 365 card.cuda_version = torch.version.cuda or "" 366 # Driver version from nvidia-smi is not in torch; leave blank 367 # unless the user fills it manually. 368 except ImportError: 369 pass 370 371 return card
Best-effort auto-detection of the current system.
Fills every field it can; leaves empty strings for fields that require libraries not installed (e.g. torch for CUDA info).
def
to_dict(self) -> Dict[str, Any]:
373 def to_dict(self) -> Dict[str, Any]: 374 return { 375 "gpu_name": self.gpu_name, 376 "gpu_memory_gb": self.gpu_memory_gb, 377 "gpu_count": self.gpu_count, 378 "cpu_name": self.cpu_name, 379 "ram_gb": self.ram_gb, 380 "cuda_version": self.cuda_version, 381 "pytorch_version": self.pytorch_version, 382 "python_version": self.python_version, 383 "precision": self.precision, 384 "batch_size": self.batch_size, 385 "input_resolution": self.input_resolution, 386 "os": self.os, 387 "driver_version": self.driver_version, 388 }
def
summary(self) -> str:
390 def summary(self) -> str: 391 """One-line summary for table footers / report headers.""" 392 parts = [] 393 if self.gpu_name: 394 parts.append(f"{self.gpu_name} ({self.gpu_memory_gb}GB)") 395 if self.cpu_name: 396 parts.append(self.cpu_name) 397 parts.append(f"PyTorch {self.pytorch_version}" if self.pytorch_version else "") 398 parts.append(f"CUDA {self.cuda_version}" if self.cuda_version else "") 399 parts.append(self.precision) 400 return " | ".join(p for p in parts if p)
One-line summary for table footers / report headers.
def
estimate_memory_traffic_gb( model: Any, input_shape: Tuple[int, ...] = (3, 480, 640), device: str = 'cpu', precision_bytes: int = 4) -> float | None:
408def estimate_memory_traffic_gb( 409 model: Any, 410 input_shape: Tuple[int, ...] = (3, 480, 640), 411 device: str = "cpu", 412 precision_bytes: int = 4, 413) -> float | None: 414 """Estimate total DRAM traffic (GB) for a single forward pass. 415 416 Strategy (in order of preference): 417 418 1. **torch.profiler** with ``profile_memory=True`` — most accurate, 419 captures actual CUDA memory events. Requires torch >= 2.0 and 420 a CUDA device. 421 2. **Analytical estimate** — conservative lower bound: 422 ``param_bytes + 2 * activation_bytes``. Works on any backend. 423 424 Parameters 425 ---------- 426 model : nn.Module or similar 427 The model to profile. 428 input_shape : tuple 429 ``(C, H, W)`` — batch dim is prepended. 430 device : str 431 Device for the dummy input. 432 precision_bytes : int 433 Bytes per element (4 for FP32, 2 for FP16/BF16). 434 435 Returns 436 ------- 437 float or None 438 Estimated DRAM traffic in GB, or None if estimation fails. 439 """ 440 # Strategy 1: torch.profiler 441 try: 442 import torch # noqa: PLC0415 443 444 if device != "cpu" and torch.cuda.is_available(): 445 dummy = torch.zeros(1, *input_shape, device=device) 446 with torch.profiler.profile( 447 activities=[ 448 torch.profiler.ProfilerActivity.CPU, 449 torch.profiler.ProfilerActivity.CUDA, 450 ], 451 profile_memory=True, 452 record_shapes=True, 453 ) as prof: 454 with torch.no_grad(): 455 model(dummy) 456 457 # Sum all CUDA memory allocations as a proxy for traffic. 458 # This overcounts (includes temporaries freed mid-forward) 459 # but is more realistic than the analytical estimate. 460 total_bytes = 0 461 for evt in prof.key_averages(): 462 mem = getattr(evt, "cuda_memory_usage", 0) or 0 463 if mem > 0: 464 total_bytes += mem 465 if total_bytes > 0: 466 return round(total_bytes / 1e9, 4) 467 except (ImportError, RuntimeError, AttributeError): 468 pass 469 470 # Strategy 2: analytical estimate 471 try: 472 params_bytes = sum(p.numel() * p.element_size() for p in model.parameters()) 473 # Heuristic: activations ~ 2x param memory for typical vision 474 # models (each layer reads previous activation, writes new one). 475 activation_bytes = params_bytes * 2 476 total = params_bytes + activation_bytes # read params + r/w activations 477 return round(total / 1e9, 4) 478 except (AttributeError, TypeError): 479 pass 480 481 return None
Estimate total DRAM traffic (GB) for a single forward pass.
Strategy (in order of preference):
- torch.profiler with
profile_memory=True— most accurate, captures actual CUDA memory events. Requires torch >= 2.0 and a CUDA device. - Analytical estimate — conservative lower bound:
param_bytes + 2 * activation_bytes. Works on any backend.
Parameters
model : nn.Module or similar
The model to profile.
input_shape : tuple
(C, H, W) — batch dim is prepended.
device : str
Device for the dummy input.
precision_bytes : int
Bytes per element (4 for FP32, 2 for FP16/BF16).
Returns
float or None Estimated DRAM traffic in GB, or None if estimation fails.
@dataclass
class
MemoryProfiler:
561@dataclass 562class MemoryProfiler: 563 """Hardware-agnostic memory profiler covering CPU / CUDA / MPS. 564 565 Usage:: 566 567 mp = MemoryProfiler().reset() 568 # ... run inference ... 569 peaks = mp.peaks() # {"cpu_mb": ..., "cuda_mb": ..., "mps_mb": ...} 570 571 Each backend returns ``None`` when it cannot be sampled (library 572 missing, backend unavailable, measurement unsupported). The caller 573 is responsible for deciding whether ``None`` is informative 574 (report as ``"n/a"``) or an error. 575 """ 576 577 sample_cpu: bool = True 578 sample_cuda: bool = True 579 sample_mps: bool = True 580 581 def reset(self) -> "MemoryProfiler": 582 """Clear accumulated peaks on every enabled backend.""" 583 if self.sample_cuda: 584 try: 585 import torch # noqa: PLC0415 586 587 if hasattr(torch, "cuda") and torch.cuda.is_available(): 588 torch.cuda.reset_peak_memory_stats() 589 except ImportError: 590 pass 591 # CPU and MPS backends don't expose a reset primitive — the 592 # sampled value is "current" (MPS) or "peak-since-process-start" 593 # (CPU ru_maxrss). Document the limitation in the docstring. 594 return self 595 596 def peaks(self) -> dict: 597 return { 598 "cpu_mb": _cpu_peak_rss_mb() if self.sample_cpu else None, 599 "cuda_mb": _cuda_peak_alloc_mb() if self.sample_cuda else None, 600 "mps_mb": _mps_current_alloc_mb() if self.sample_mps else None, 601 }
Hardware-agnostic memory profiler covering CPU / CUDA / MPS.
Usage::
mp = MemoryProfiler().reset()
# ... run inference ...
peaks = mp.peaks() # {"cpu_mb": ..., "cuda_mb": ..., "mps_mb": ...}
Each backend returns None when it cannot be sampled (library
missing, backend unavailable, measurement unsupported). The caller
is responsible for deciding whether None is informative
(report as "n/a") or an error.
581 def reset(self) -> "MemoryProfiler": 582 """Clear accumulated peaks on every enabled backend.""" 583 if self.sample_cuda: 584 try: 585 import torch # noqa: PLC0415 586 587 if hasattr(torch, "cuda") and torch.cuda.is_available(): 588 torch.cuda.reset_peak_memory_stats() 589 except ImportError: 590 pass 591 # CPU and MPS backends don't expose a reset primitive — the 592 # sampled value is "current" (MPS) or "peak-since-process-start" 593 # (CPU ru_maxrss). Document the limitation in the docstring. 594 return self
Clear accumulated peaks on every enabled backend.
@dataclass
class
LatencyProfiler:
609@dataclass 610class LatencyProfiler: 611 """Accumulate per-sample latencies and report percentiles. 612 613 Percentiles use linear interpolation (``numpy.percentile`` default) 614 and skip a user-configurable number of warmup samples so the 615 reported numbers reflect steady-state throughput rather than the 616 first-batch JIT compile. 617 """ 618 619 warmup: int = 1 620 _samples_ms: List[float] = field(default_factory=list) 621 622 def add_sample_seconds(self, seconds: float) -> None: 623 self._samples_ms.append(float(seconds) * 1000.0) 624 625 def add_batch_seconds(self, batch_seconds: float, batch_size: int) -> None: 626 """Amortise a batch timing over ``batch_size`` sample entries.""" 627 if batch_size <= 0: 628 return 629 per = batch_seconds / batch_size 630 self._samples_ms.extend([per * 1000.0] * batch_size) 631 632 def percentiles(self) -> dict: 633 """Return ``{p50_ms, p95_ms, p99_ms, mean_ms}`` in milliseconds. 634 635 Returns ``{k: None for k in keys}`` when fewer than one sample 636 exists after warmup trimming. 637 """ 638 trimmed = self._samples_ms[self.warmup :] or self._samples_ms 639 if not trimmed: 640 return {"p50_ms": None, "p95_ms": None, "p99_ms": None, "mean_ms": None} 641 try: 642 import numpy as np # noqa: PLC0415 643 644 p50, p95, p99 = np.percentile(trimmed, [50, 95, 99]) 645 return { 646 "p50_ms": round(float(p50), 3), 647 "p95_ms": round(float(p95), 3), 648 "p99_ms": round(float(p99), 3), 649 "mean_ms": round(float(np.mean(trimmed)), 3), 650 } 651 except ImportError: 652 trimmed_sorted = sorted(trimmed) 653 n = len(trimmed_sorted) 654 655 def _q(frac: float) -> float: 656 return trimmed_sorted[min(n - 1, int(round(frac * (n - 1))))] 657 658 return { 659 "p50_ms": round(_q(0.50), 3), 660 "p95_ms": round(_q(0.95), 3), 661 "p99_ms": round(_q(0.99), 3), 662 "mean_ms": round(sum(trimmed) / n, 3), 663 } 664 665 def samples_ms(self) -> Sequence[float]: 666 """Return the raw latency samples (post-warmup not trimmed).""" 667 return list(self._samples_ms)
Accumulate per-sample latencies and report percentiles.
Percentiles use linear interpolation (numpy.percentile default)
and skip a user-configurable number of warmup samples so the
reported numbers reflect steady-state throughput rather than the
first-batch JIT compile.
def
add_batch_seconds(self, batch_seconds: float, batch_size: int) -> None:
625 def add_batch_seconds(self, batch_seconds: float, batch_size: int) -> None: 626 """Amortise a batch timing over ``batch_size`` sample entries.""" 627 if batch_size <= 0: 628 return 629 per = batch_seconds / batch_size 630 self._samples_ms.extend([per * 1000.0] * batch_size)
Amortise a batch timing over batch_size sample entries.
def
percentiles(self) -> dict:
632 def percentiles(self) -> dict: 633 """Return ``{p50_ms, p95_ms, p99_ms, mean_ms}`` in milliseconds. 634 635 Returns ``{k: None for k in keys}`` when fewer than one sample 636 exists after warmup trimming. 637 """ 638 trimmed = self._samples_ms[self.warmup :] or self._samples_ms 639 if not trimmed: 640 return {"p50_ms": None, "p95_ms": None, "p99_ms": None, "mean_ms": None} 641 try: 642 import numpy as np # noqa: PLC0415 643 644 p50, p95, p99 = np.percentile(trimmed, [50, 95, 99]) 645 return { 646 "p50_ms": round(float(p50), 3), 647 "p95_ms": round(float(p95), 3), 648 "p99_ms": round(float(p99), 3), 649 "mean_ms": round(float(np.mean(trimmed)), 3), 650 } 651 except ImportError: 652 trimmed_sorted = sorted(trimmed) 653 n = len(trimmed_sorted) 654 655 def _q(frac: float) -> float: 656 return trimmed_sorted[min(n - 1, int(round(frac * (n - 1))))] 657 658 return { 659 "p50_ms": round(_q(0.50), 3), 660 "p95_ms": round(_q(0.95), 3), 661 "p99_ms": round(_q(0.99), 3), 662 "mean_ms": round(sum(trimmed) / n, 3), 663 }
Return {p50_ms, p95_ms, p99_ms, mean_ms} in milliseconds.
Returns {k: None for k in keys} when fewer than one sample
exists after warmup trimming.
def
count_parameters(model: Any) -> float:
670def count_parameters(model: Any) -> float: 671 """Count trainable parameters in millions. 672 673 Supports PyTorch nn.Module and any object with a ``parameters()`` method. 674 Returns None if the model type is not supported. 675 """ 676 try: 677 params = sum(p.numel() for p in model.parameters() if p.requires_grad) 678 return round(params / 1e6, 3) 679 except AttributeError: 680 pass 681 682 # JAX / Flax: model may expose a ``params`` pytree 683 try: 684 import jax 685 686 leaves = jax.tree_util.tree_leaves(model.params) 687 params = sum(leaf.size for leaf in leaves) 688 return round(params / 1e6, 3) 689 except (AttributeError, ImportError): 690 pass 691 692 return None
Count trainable parameters in millions.
Supports PyTorch nn.Module and any object with a parameters() method.
Returns None if the model type is not supported.
def
count_flops_torch( model: Any, input_shape: Tuple[int, ...], device: str = 'cpu') -> float | None:
695def count_flops_torch( 696 model: Any, 697 input_shape: Tuple[int, ...], 698 device: str = "cpu", 699) -> float | None: 700 """Estimate FLOPs (giga) using ``torch.utils.flop_counter`` (PyTorch >= 2.1). 701 702 Falls back to ``fvcore`` if available. 703 704 **OOM-safe**: when the FLOP counter runs out of CUDA memory (the 705 dispatch metadata alone can consume ~3 GB on top of the model), the 706 function automatically retries on CPU. The FLOP count is identical 707 regardless of device — it is a graph property, not a runtime one. 708 709 Args: 710 model: PyTorch nn.Module. 711 input_shape: (C, H, W) — batch dimension is added automatically. 712 device: device string for the dummy input tensor. 713 714 Returns: 715 FLOPs in giga-ops, or None if neither backend is available. 716 """ 717 718 def _try_flop_count(mod: Any, dev: str) -> float | None: 719 try: 720 import torch # noqa: PLC0415 721 except ImportError: 722 return None 723 724 dummy = torch.zeros(1, *input_shape, device=dev) 725 726 # PyTorch 2.1+ native counter 727 try: 728 from torch.utils.flop_counter import FlopCounterMode # noqa: PLC0415 729 730 with torch.no_grad(): 731 with FlopCounterMode(display=False) as fcm: 732 mod(dummy) 733 total = _sum_nested_counts(fcm.get_flop_counts()) 734 if total > 0: 735 return round(total / 1e9, 3) 736 except ImportError: 737 pass 738 739 # fvcore fallback 740 try: 741 from fvcore.nn import FlopCountAnalysis # noqa: PLC0415 742 743 flops = FlopCountAnalysis(mod, dummy) 744 return round(flops.total() / 1e9, 3) 745 except ImportError: 746 pass 747 748 return None 749 750 # Helper for nested FlopCounterMode dicts (same logic as runner.py) 751 def _sum_nested_counts(node: Any) -> int: 752 if isinstance(node, dict): 753 return sum(_sum_nested_counts(v) for v in node.values()) 754 try: 755 return int(node) 756 except (TypeError, ValueError): 757 return 0 758 759 # Try on the requested device first 760 try: 761 result = _try_flop_count(model, device) 762 if result is not None: 763 return result 764 except RuntimeError as exc: 765 # Catch CUDA OOM and fall through to CPU retry 766 if "out of memory" not in str(exc).lower(): 767 return None 768 769 # CPU fallback: FLOPs are a graph property — the count is device- 770 # independent. Move the model to CPU, count, then move it back. 771 if device != "cpu": 772 try: 773 original_device = next(model.parameters()).device 774 model_cpu = model.cpu() 775 result = _try_flop_count(model_cpu, "cpu") 776 model.to(original_device) 777 return result 778 except Exception: 779 # If anything goes wrong restoring the model, still return 780 # None rather than crashing the caller. 781 try: 782 model.to(device) 783 except Exception: 784 pass 785 return None 786 787 return None
Estimate FLOPs (giga) using torch.utils.flop_counter (PyTorch >= 2.1).
Falls back to fvcore if available.
OOM-safe: when the FLOP counter runs out of CUDA memory (the dispatch metadata alone can consume ~3 GB on top of the model), the function automatically retries on CPU. The FLOP count is identical regardless of device — it is a graph property, not a runtime one.
Args: model: PyTorch nn.Module. input_shape: (C, H, W) — batch dimension is added automatically. device: device string for the dummy input tensor.
Returns: FLOPs in giga-ops, or None if neither backend is available.
def
profile_model( model: Any, input_shape: Tuple[int, ...] = (3, 480, 640), device: str = 'cpu', model_type: str = 'local', notes: str = '', compute_roofline: bool = True, attach_system_card: bool = True, precision: str = 'fp32') -> EfficiencyMetadata:
790def profile_model( 791 model: Any, 792 input_shape: Tuple[int, ...] = (3, 480, 640), 793 device: str = "cpu", 794 model_type: str = "local", 795 notes: str = "", 796 compute_roofline: bool = True, 797 attach_system_card: bool = True, 798 precision: str = "fp32", 799) -> EfficiencyMetadata: 800 """Auto-profile a model and return EfficiencyMetadata. 801 802 Fills all three tiers of efficiency reporting: 803 804 - **Tier 1**: params, FLOPs, MACs, memory traffic, arithmetic 805 intensity (hardware-agnostic). 806 - **Tier 2**: roofline latency bounds for reference GPUs 807 (hardware-parametric). 808 - **Tier 3**: system card for the current machine 809 (hardware-specific; measured latency is added later by the runner). 810 811 Parameters 812 ---------- 813 model : nn.Module or similar 814 The model to profile. 815 input_shape : tuple 816 ``(C, H, W)`` for FLOPs counting. Default 640x480 RGB. 817 device : str 818 Device for the dummy input tensor. 819 model_type : str 820 ``"local"`` or ``"api"``. 821 notes : str 822 Free-text notes (e.g. ``"ViT-L/14, FP16 inference"``). 823 compute_roofline : bool 824 If True, compute Tier-2 roofline bounds for reference GPUs. 825 attach_system_card : bool 826 If True, auto-detect and attach a SystemCard. 827 precision : str 828 Precision string for the system card (``"fp32"``, ``"fp16"``). 829 830 Returns 831 ------- 832 EfficiencyMetadata 833 With Tier 1/2/3 fields populated where possible. 834 """ 835 if model_type == "api": 836 return EfficiencyMetadata(model_type="api", notes=notes) 837 838 params_m = count_parameters(model) 839 flops_g = count_flops_torch(model, input_shape, device=device) 840 traffic_gb = estimate_memory_traffic_gb(model, input_shape, device=device) 841 842 h, w = input_shape[-2], input_shape[-1] 843 em = EfficiencyMetadata( 844 params_m=params_m, 845 flops_g=flops_g, 846 memory_traffic_gb=traffic_gb, 847 model_type=model_type, 848 notes=notes, 849 ) 850 851 # Derive Tier 1 (MACs, arithmetic intensity) 852 em.derive_tier1() 853 854 # Tier 2: roofline bounds 855 if compute_roofline: 856 em.compute_roofline() 857 858 # Tier 3: system card 859 if attach_system_card: 860 em.system_card = SystemCard.auto_detect( 861 precision=precision, 862 input_resolution=f"{w}x{h}", 863 ) 864 865 return em
Auto-profile a model and return EfficiencyMetadata.
Fills all three tiers of efficiency reporting:
- Tier 1: params, FLOPs, MACs, memory traffic, arithmetic intensity (hardware-agnostic).
- Tier 2: roofline latency bounds for reference GPUs (hardware-parametric).
- Tier 3: system card for the current machine (hardware-specific; measured latency is added later by the runner).
Parameters
model : nn.Module or similar
The model to profile.
input_shape : tuple
(C, H, W) for FLOPs counting. Default 640x480 RGB.
device : str
Device for the dummy input tensor.
model_type : str
"local" or "api".
notes : str
Free-text notes (e.g. "ViT-L/14, FP16 inference").
compute_roofline : bool
If True, compute Tier-2 roofline bounds for reference GPUs.
attach_system_card : bool
If True, auto-detect and attach a SystemCard.
precision : str
Precision string for the system card ("fp32", "fp16").
Returns
EfficiencyMetadata With Tier 1/2/3 fields populated where possible.