SPB Git forge

spb/ai-atlas

Public
41commits 1branches 0releases
4.6 MBsize
maindefault branch
12 days agolast push
HTML 77.2% TypeScript 10.5% Python 9.6% JavaScript 2.5%
8.8 KB · 164 lines python
Raw Blame History
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               — or the OBSERVED `file_size_gb` of an artifact when one is available (flagged `weights_source: observed`)5    kv_cache = per layer: 2 (K and V) × kv_heads × head_dim × bytes × context × batch — when the architecture metadata is known6               (`num_hidden_layers`, `num_key_value_heads` | `num_attention_heads`, `head_dim` | `hidden_size`); otherwise the documented7               heuristic 0.5 GB per 8 192 tokens (× batch), architecture-agnostic8    fits     = estimated_memory_gb ≤ hardware memory − 2 GB (OS / framework headroom); multi-GPU = sum of device memory, interconnect ignored9"""10from __future__ import annotations1112from typing import Any1314BYTES_PER_PARAM = {"4bit": 0.5, "8bit": 1.0, "fp16": 2.0, "bf16": 2.0, "fp8": 1.0, "int4": 0.5, "int8": 1.0, "fp32": 4.0}15KV_BYTES = {"fp16": 2.0, "bf16": 2.0, "fp8": 1.0, "int8": 1.0, "8bit": 1.0, "4bit": 2.0, "int4": 2.0, "fp32": 4.0}  # KV cache dtype (4-bit weights usually keep fp16 KV)16OVERHEAD = 1.1517KV_GB_PER_8K = 0.518RESERVED_GB = 2.019ASSUMPTIONS = [20    "Estimated, not measured: weights = parameters × bytes/param × 1.15 runtime overhead (or the observed artifact file size when one is recorded).",21    "bytes/param: 4bit = 0.5, 8bit = 1.0, fp16 = 2.0 (uniform quantization, no per-layer exceptions).",22    "KV cache: 2 × layers × kv_heads × head_dim × 2 bytes × context × batch when the architecture is known; otherwise 0.5 GB per 8 192 tokens "23    "(× batch), independent of architecture (GQA/MLA models need less).",24    "A model 'fits' when the estimate is at most the device memory minus 2 GB reserved for the OS and framework.",25    "Mixture-of-experts models are estimated on total parameters (all experts must be resident); active parameters are ignored.",26    "Device memory uses the largest configuration when several are listed (e.g. Apple silicon tiers).",27    "Multi-GPU: device memories are summed; interconnect bandwidth, tensor-parallel replication and pipeline bubbles are not modelled.",28]29QUANT_ALIASES = {"q4": "4bit", "int4": "4bit", "nf4": "4bit", "gguf-q4": "4bit", "mlx-4bit": "4bit", "q8": "8bit", "int8": "8bit", "fp8": "8bit",30                 "f16": "fp16", "half": "fp16", "bf16": "fp16", "fp16": "fp16", "4bit": "4bit", "8bit": "8bit", "fp32": "fp32", "f32": "fp32"}313233def normalize_quant(q: str | None) -> str:34    s = (q or "4bit").strip().lower()35    return QUANT_ALIASES.get(s, s if s in BYTES_PER_PARAM else "4bit")363738def estimate_memory_gb(parameter_count: float, quant: str = "4bit", context: int = 8192) -> float:39    bpp = BYTES_PER_PARAM.get(quant, BYTES_PER_PARAM["4bit"])40    weights = parameter_count * bpp * OVERHEAD / 1e941    kv = KV_GB_PER_8K * max(0.0, float(context)) / 8192.042    return round(weights + kv, 2)434445def hardware_memory_gb(attrs: dict[str, Any] | None) -> float | None:46    """`memory_gb` may be a number or a list of configurations — use the largest."""47    v = (attrs or {}).get("memory_gb")48    if isinstance(v, list):49        nums = [float(x) for x in v if isinstance(x, (int, float)) and not isinstance(x, bool)]50        return max(nums) if nums else None51    if isinstance(v, (int, float)) and not isinstance(v, bool):52        return float(v)53    if isinstance(v, str):54        try:55            return float(v)56        except ValueError:57            return None58    return None596061def hardware_memory_options(attrs: dict[str, Any] | None) -> list[float]:62    v = (attrs or {}).get("memory_gb")63    if isinstance(v, list):64        return sorted(float(x) for x in v if isinstance(x, (int, float)) and not isinstance(x, bool))65    m = hardware_memory_gb(attrs)66    return [m] if m is not None else []676869def _num(v: Any) -> float | None:70    if isinstance(v, bool):71        return None72    if isinstance(v, (int, float)):73        return float(v)74    if isinstance(v, str):75        try:76            return float(v.replace(",", ""))77        except ValueError:78            return None79    return None808182def parameter_count(attrs: dict[str, Any] | None) -> float | None:83    v = _num((attrs or {}).get("parameter_count"))84    return v if v is not None and v > 0 else None858687def file_size_gb(attrs: dict[str, Any] | None) -> float | None:88    """Observed weights size (Hugging Face repo file total) when recorded."""89    v = _num((attrs or {}).get("file_size_gb"))90    return v if v is not None and v > 0 else None919293def architecture(attrs: dict[str, Any] | None) -> dict[str, float] | None:94    """{layers, kv_heads, head_dim} from Hugging Face-style config keys when all three can be derived, else None."""95    a = attrs or {}96    layers = _num(a.get("num_hidden_layers") or a.get("n_layer") or a.get("num_layers"))97    heads = _num(a.get("num_attention_heads") or a.get("n_head"))98    kv_heads = _num(a.get("num_key_value_heads")) or heads99    head_dim = _num(a.get("head_dim"))100    hidden = _num(a.get("hidden_size") or a.get("n_embd"))101    if head_dim is None and hidden and heads:102        head_dim = hidden / heads103    if layers and kv_heads and head_dim:104        return {"layers": layers, "kv_heads": kv_heads, "head_dim": head_dim}105    return None106107108def kv_cache_gb(context: int, *, batch: int = 1, quant: str = "4bit", arch: dict[str, float] | None = None) -> tuple[float, str]:109    """(GB, method). `method` is 'architecture' when computed from layers × kv_heads × head_dim, else 'heuristic'."""110    ctx = max(0.0, float(context))111    b = max(1, int(batch))112    if arch:113        bytes_ = KV_BYTES.get(quant, 2.0)114        gb = 2.0 * arch["layers"] * arch["kv_heads"] * arch["head_dim"] * bytes_ * ctx * b / 1e9115        return round(gb, 3), "architecture"116    return round(KV_GB_PER_8K * ctx / 8192.0 * b, 3), "heuristic"117118119def fit(parameter_count_: float, memory_gb: float, quant: str = "4bit", context: int = 8192) -> dict[str, Any]:120    est = estimate_memory_gb(parameter_count_, quant, context)121    headroom = round(memory_gb - RESERVED_GB - est, 2)122    return {"quantization": quant, "estimated_memory_gb": est, "fits": headroom >= 0, "headroom_gb": headroom, "estimated": True,123            "note": f"estimated: {parameter_count_ / 1e9:.1f}B params × {BYTES_PER_PARAM.get(quant, 0.5)} B × 1.15 + KV cache for {context} tokens"}124125126def fit_detailed(attrs: dict[str, Any] | None, memory_gb: float, *, quant: str = "4bit", context: int = 8192, batch: int = 1, gpu_count: int = 1,127                 observed_size_gb: float | None = None) -> dict[str, Any] | None:128    """Full breakdown (weights / KV cache / overhead / headroom). `observed_size_gb` (artifact file size) replaces the parameter-based weights129    estimate when given. Returns None when neither parameters nor an observed size are known — nothing is estimated from thin air."""130    quant = normalize_quant(quant)131    params = parameter_count(attrs)132    observed = observed_size_gb if observed_size_gb else None133    if params is None and observed is None:134        return None135    if observed is not None:136        weights_raw = observed137        weights_source = "observed"138    else:139        assert params is not None140        weights_raw = params * BYTES_PER_PARAM[quant] / 1e9141        weights_source = "estimated"142    overhead = round(weights_raw * (OVERHEAD - 1), 2)143    kv, kv_method = kv_cache_gb(context, batch=batch, quant=quant, arch=architecture(attrs))144    total = round(weights_raw + overhead + kv, 2)145    total_memory = float(memory_gb) * max(1, int(gpu_count))146    headroom = round(total_memory - RESERVED_GB - total, 2)147    out: dict[str, Any] = {148        "quantization": quant, "estimated": True, "fits": headroom >= 0, "estimated_memory_gb": total, "headroom_gb": headroom,149        "breakdown": {"weights_gb": round(weights_raw, 2), "weights_source": weights_source, "overhead_gb": overhead, "kv_cache_gb": kv, "kv_cache_method": kv_method,150                      "reserved_gb": RESERVED_GB, "context": int(context), "batch": int(batch)},151        "device": {"memory_gb": float(memory_gb), "gpu_count": int(gpu_count), "total_memory_gb": total_memory},152        "note": (f"{'observed' if weights_source == 'observed' else 'estimated'} weights {weights_raw:.1f} GB + 15% overhead + KV cache {kv:.2f} GB ({kv_method}) "153                 f"for {context} tokens × batch {batch}; {RESERVED_GB:.0f} GB reserved for the OS/framework"),154    }155    if params is not None:156        out["parameter_count"] = params157    if gpu_count > 1:158        out["multi_gpu_note"] = "Device memories summed; interconnect bandwidth, tensor-parallel replication and pipeline bubbles are not modelled."159    return out160161162__all__ = ["ASSUMPTIONS", "BYTES_PER_PARAM", "KV_BYTES", "RESERVED_GB", "architecture", "estimate_memory_gb", "file_size_gb", "fit", "fit_detailed",163           "hardware_memory_gb", "hardware_memory_options", "kv_cache_gb", "normalize_quant", "parameter_count"]164