1"""ESTIMATED memory footprint of a model on a piece of hardware. Transparent formula, always labelled as an estimate:23 weights = parameter_count × bytes_per_param × 1.15 (runtime overhead: activations, buffers, fragmentation)4 kv_cache = 0.5 GB per 8 192 tokens of context (coarse, architecture-agnostic)5 fits = estimated_memory_gb ≤ hardware memory − 2 GB (OS / framework headroom)6"""7from __future__ import annotations89from typing import Any1011BYTES_PER_PARAM = {"4bit": 0.5, "8bit": 1.0, "fp16": 2.0}12OVERHEAD = 1.1513KV_GB_PER_8K = 0.514RESERVED_GB = 2.015ASSUMPTIONS = [16 "Estimated, not measured: weights = parameters × bytes/param × 1.15 runtime overhead.",17 "bytes/param: 4bit = 0.5, 8bit = 1.0, fp16 = 2.0 (uniform quantization, no per-layer exceptions).",18 "KV cache approximated at 0.5 GB per 8 192 tokens of context, independent of architecture (GQA/MLA models need less).",19 "A model 'fits' when the estimate is at most the device memory minus 2 GB reserved for the OS and framework.",20 "Mixture-of-experts models are estimated on total parameters (all experts must be resident); active parameters are ignored.",21 "Device memory uses the largest configuration when several are listed (e.g. Apple silicon tiers).",22]232425def estimate_memory_gb(parameter_count: float, quant: str = "4bit", context: int = 8192) -> float:26 bpp = BYTES_PER_PARAM.get(quant, BYTES_PER_PARAM["4bit"])27 weights = parameter_count * bpp * OVERHEAD / 1e928 kv = KV_GB_PER_8K * max(0.0, float(context)) / 8192.029 return round(weights + kv, 2)303132def hardware_memory_gb(attrs: dict[str, Any] | None) -> float | None:33 """`memory_gb` may be a number or a list of configurations — use the largest."""34 v = (attrs or {}).get("memory_gb")35 if isinstance(v, list):36 nums = [float(x) for x in v if isinstance(x, (int, float)) and not isinstance(x, bool)]37 return max(nums) if nums else None38 if isinstance(v, (int, float)) and not isinstance(v, bool):39 return float(v)40 if isinstance(v, str):41 try:42 return float(v)43 except ValueError:44 return None45 return None464748def parameter_count(attrs: dict[str, Any] | None) -> float | None:49 v = (attrs or {}).get("parameter_count")50 if isinstance(v, (int, float)) and not isinstance(v, bool) and v > 0:51 return float(v)52 if isinstance(v, str):53 try:54 f = float(v)55 return f if f > 0 else None56 except ValueError:57 return None58 return None596061def fit(parameter_count_: float, memory_gb: float, quant: str = "4bit", context: int = 8192) -> dict[str, Any]:62 est = estimate_memory_gb(parameter_count_, quant, context)63 headroom = round(memory_gb - RESERVED_GB - est, 2)64 return {"quantization": quant, "estimated_memory_gb": est, "fits": headroom >= 0, "headroom_gb": headroom,65 "note": f"estimated: {parameter_count_ / 1e9:.1f}B params × {BYTES_PER_PARAM.get(quant, 0.5)} B × 1.15 + KV cache for {context} tokens"}666768__all__ = ["ASSUMPTIONS", "BYTES_PER_PARAM", "estimate_memory_gb", "fit", "hardware_memory_gb", "parameter_count"]69