merge: feat/64-vram-estimate into freeze/0.22.0
# Conflicts: # CHANGELOG.md
This commit is contained in:
commit
8d70783896
5 changed files with 383 additions and 0 deletions
10
CHANGELOG.md
10
CHANGELOG.md
|
|
@ -24,6 +24,16 @@ Accepts a multipart video file upload, writes it to a temp file, captions it via
|
|||
|
||||
- `mkdocs.yml` — new blue-grey/cyan palette, wired to a central `docs/stylesheets/theme.css` for consistent theme-aware styling across the site (light and dark mode).
|
||||
|
||||
**`circuitforge_core.hardware.model_vram_estimate`** — model-to-hardware VRAM fit check (closes #64)
|
||||
|
||||
Answers "can this hardware run model X at quantization level Y?" — the missing capability noted against `cf_core.hardware`, which detects available VRAM but had no way to cross-reference model requirements. Queries the HuggingFace Hub API for parameter count (`safetensors.total`) and architecture (`config.json`: `num_hidden_layers`, `hidden_size`, `num_attention_heads`, `num_key_value_heads`), then applies the standard formula: `vram_gb = params * bytes_per_param(quant) + kv_cache_gb(ctx_len, arch) + overhead_gb`. Reference algorithm from [LLMcalc](https://github.com/Raskoll2/LLMcalc) (unlicensed upstream — algorithm reference only, no code copied, no dependency added).
|
||||
|
||||
- `model_vram_estimate(hf_model_id, quant_level, *, ctx_len=4096, available_vram_mb=None, overhead_gb=0.6, timeout=10.0) -> VramEstimate`
|
||||
- Supports common quant levels: `fp32`, `fp16`/`bf16`, `int8`/`q8`/`q8_0`, `q6_k`, `q5_k_m`/`q5_0`, `int4`/`q4`/`q4_k_m`/`q4_0`, `q3_k_m`, `q2_k`.
|
||||
- KV cache sizing accounts for GQA (`num_key_value_heads`); falls back to 0 GB when the model's `config.json` lacks standard architecture fields, rather than failing the whole estimate — weights dominate VRAM use regardless.
|
||||
- Raises `ModelVramLookupError` on HF Hub API failures or missing safetensors metadata; raises `ValueError` for unrecognized quant levels.
|
||||
- Application points noted in the ticket: Avocet preflight (verify a checkpoint fits before benchmarking), cf-orch worker assignment (match model to GPU by VRAM fit), Peregrine/Kiwi onboarding wizard ("your GPU has X GB — here are models that will run well").
|
||||
|
||||
---
|
||||
|
||||
## [0.20.0] — 2026-05-05
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from .detect import detect_hardware, detect_hardware_json
|
|||
from .generator import generate_profile
|
||||
from .models import HardwareSpec, LLMBackendConfig, LLMConfig
|
||||
from .tiers import VRAM_TIERS, VramTier, select_tier
|
||||
from .vram_estimate import ModelVramLookupError, VramEstimate, model_vram_estimate
|
||||
|
||||
__all__ = [
|
||||
"detect_hardware",
|
||||
|
|
@ -27,4 +28,7 @@ __all__ = [
|
|||
"VRAM_TIERS",
|
||||
"VramTier",
|
||||
"select_tier",
|
||||
"ModelVramLookupError",
|
||||
"VramEstimate",
|
||||
"model_vram_estimate",
|
||||
]
|
||||
|
|
|
|||
185
circuitforge_core/hardware/vram_estimate.py
Normal file
185
circuitforge_core/hardware/vram_estimate.py
Normal file
|
|
@ -0,0 +1,185 @@
|
|||
# circuitforge_core/hardware/vram_estimate.py
|
||||
"""
|
||||
Model VRAM fit estimation — cf-core #64.
|
||||
|
||||
`cf_core.hardware` can detect available VRAM but has no way to answer "can this
|
||||
hardware run model X at quantization level Y?". This module closes that gap by
|
||||
querying the HuggingFace Hub API for parameter count and architecture, then
|
||||
applying the standard VRAM estimation formula:
|
||||
|
||||
vram_gb = params * bytes_per_param(quant) + kv_cache_gb(ctx_len, arch) + overhead_gb
|
||||
|
||||
The formula is the reference algorithm used by LLMcalc
|
||||
(https://github.com/Raskoll2/LLMcalc, no license — algorithm reference only,
|
||||
not a dependency). No LLMcalc code is copied here.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import requests
|
||||
|
||||
_HF_API_MODEL_URL = "https://huggingface.co/api/models/{model_id}"
|
||||
_HF_CONFIG_URL = "https://huggingface.co/{model_id}/resolve/main/config.json"
|
||||
|
||||
_DEFAULT_OVERHEAD_GB = 0.6 # CUDA context + activation buffers, rough constant
|
||||
_KV_CACHE_DTYPE_BYTES = 2 # KV cache is stored fp16 in the common case
|
||||
|
||||
# Bits per parameter for common quantization levels (weights-only, excludes KV cache).
|
||||
_QUANT_BITS: dict[str, float] = {
|
||||
"fp32": 32.0,
|
||||
"fp16": 16.0,
|
||||
"bf16": 16.0,
|
||||
"int8": 8.0,
|
||||
"q8": 8.0,
|
||||
"q8_0": 8.0,
|
||||
"q6_k": 6.0,
|
||||
"q5_k_m": 5.0,
|
||||
"q5_0": 5.0,
|
||||
"int4": 4.0,
|
||||
"q4": 4.0,
|
||||
"q4_k_m": 4.0,
|
||||
"q4_0": 4.0,
|
||||
"q3_k_m": 3.0,
|
||||
"q2_k": 2.0,
|
||||
}
|
||||
|
||||
|
||||
class ModelVramLookupError(RuntimeError):
|
||||
"""Raised when the HuggingFace Hub API can't supply data needed to estimate VRAM."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VramEstimate:
|
||||
"""Result of a `model_vram_estimate()` call."""
|
||||
|
||||
hf_model_id: str
|
||||
quant_level: str
|
||||
params_billions: float
|
||||
weights_gb: float
|
||||
kv_cache_gb: float
|
||||
overhead_gb: float
|
||||
total_vram_gb: float
|
||||
fits: bool | None # None when available_vram_mb wasn't supplied
|
||||
|
||||
|
||||
def _quant_bits(quant_level: str) -> float:
|
||||
key = quant_level.strip().lower()
|
||||
try:
|
||||
return _QUANT_BITS[key]
|
||||
except KeyError:
|
||||
raise ValueError(
|
||||
f"Unknown quant_level {quant_level!r}. Known levels: {sorted(_QUANT_BITS)}"
|
||||
) from None
|
||||
|
||||
|
||||
def _fetch_param_count(hf_model_id: str, *, timeout: float) -> float:
|
||||
"""Return total parameter count via the HF Hub API's safetensors metadata."""
|
||||
url = _HF_API_MODEL_URL.format(model_id=hf_model_id)
|
||||
try:
|
||||
resp = requests.get(url, params={"expand": ["safetensors"]}, timeout=timeout)
|
||||
except requests.RequestException as exc:
|
||||
raise ModelVramLookupError(
|
||||
f"HF Hub API request failed for {hf_model_id!r}: {exc}"
|
||||
) from exc
|
||||
|
||||
if resp.status_code != 200:
|
||||
raise ModelVramLookupError(
|
||||
f"HF Hub API returned {resp.status_code} for {hf_model_id!r}"
|
||||
)
|
||||
|
||||
data = resp.json()
|
||||
total = (data.get("safetensors") or {}).get("total")
|
||||
if not total:
|
||||
raise ModelVramLookupError(
|
||||
f"{hf_model_id!r} has no safetensors parameter metadata on the HF Hub "
|
||||
"(model may not publish safetensors weights)"
|
||||
)
|
||||
return float(total)
|
||||
|
||||
|
||||
def _fetch_arch_config(hf_model_id: str, *, timeout: float) -> dict:
|
||||
"""Return the model's config.json (architecture fields used for KV cache sizing)."""
|
||||
url = _HF_CONFIG_URL.format(model_id=hf_model_id)
|
||||
try:
|
||||
resp = requests.get(url, timeout=timeout)
|
||||
except requests.RequestException as exc:
|
||||
raise ModelVramLookupError(
|
||||
f"config.json fetch failed for {hf_model_id!r}: {exc}"
|
||||
) from exc
|
||||
|
||||
if resp.status_code != 200:
|
||||
raise ModelVramLookupError(
|
||||
f"config.json unavailable for {hf_model_id!r} (HTTP {resp.status_code})"
|
||||
)
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _kv_cache_gb(config: dict, ctx_len: int) -> float:
|
||||
"""Estimate KV cache size in GB from architecture fields, 0.0 if unavailable."""
|
||||
num_layers = config.get("num_hidden_layers")
|
||||
hidden_size = config.get("hidden_size")
|
||||
num_heads = config.get("num_attention_heads")
|
||||
num_kv_heads = config.get("num_key_value_heads", num_heads)
|
||||
|
||||
if not (num_layers and hidden_size and num_heads):
|
||||
# Non-standard config (missing architecture fields) — skip the KV
|
||||
# estimate rather than fail the whole call. Weights dominate VRAM use.
|
||||
return 0.0
|
||||
|
||||
head_dim = hidden_size / num_heads
|
||||
bytes_total = 2 * num_layers * num_kv_heads * head_dim * ctx_len * _KV_CACHE_DTYPE_BYTES
|
||||
return bytes_total / 1e9
|
||||
|
||||
|
||||
def model_vram_estimate(
|
||||
hf_model_id: str,
|
||||
quant_level: str,
|
||||
*,
|
||||
ctx_len: int = 4096,
|
||||
available_vram_mb: int | None = None,
|
||||
overhead_gb: float = _DEFAULT_OVERHEAD_GB,
|
||||
timeout: float = 10.0,
|
||||
) -> VramEstimate:
|
||||
"""
|
||||
Estimate VRAM required to run `hf_model_id` at `quant_level`.
|
||||
|
||||
Args:
|
||||
hf_model_id: HuggingFace model repo ID, e.g. "Qwen/Qwen2.5-7B-Instruct".
|
||||
quant_level: One of the known quant levels (see `_QUANT_BITS`), e.g. "q4_k_m".
|
||||
ctx_len: Context length used for KV cache sizing.
|
||||
available_vram_mb: If given, populates `VramEstimate.fits`.
|
||||
overhead_gb: Fixed overhead for CUDA context / activation buffers.
|
||||
timeout: Per-request timeout in seconds for HF Hub API calls.
|
||||
|
||||
Raises:
|
||||
ModelVramLookupError: HF Hub API request failed or returned unusable data.
|
||||
ValueError: `quant_level` isn't a recognized quantization level.
|
||||
"""
|
||||
bits = _quant_bits(quant_level)
|
||||
params = _fetch_param_count(hf_model_id, timeout=timeout)
|
||||
weights_gb = (params * bits / 8) / 1e9
|
||||
|
||||
try:
|
||||
config = _fetch_arch_config(hf_model_id, timeout=timeout)
|
||||
kv_gb = _kv_cache_gb(config, ctx_len)
|
||||
except ModelVramLookupError:
|
||||
# Architecture lookup is best-effort — weights_gb alone is still useful.
|
||||
kv_gb = 0.0
|
||||
|
||||
total_gb = weights_gb + kv_gb + overhead_gb
|
||||
|
||||
fits = None
|
||||
if available_vram_mb is not None:
|
||||
fits = total_gb <= (available_vram_mb / 1024)
|
||||
|
||||
return VramEstimate(
|
||||
hf_model_id=hf_model_id,
|
||||
quant_level=quant_level,
|
||||
params_billions=params / 1e9,
|
||||
weights_gb=weights_gb,
|
||||
kv_cache_gb=kv_gb,
|
||||
overhead_gb=overhead_gb,
|
||||
total_vram_gb=total_gb,
|
||||
fits=fits,
|
||||
)
|
||||
|
|
@ -49,3 +49,18 @@ Profile selection rules:
|
|||
## HardwareProfile
|
||||
|
||||
The `HardwareProfile` dataclass is written to `compose.override.yml` by `preflight.py` at product startup, making GPU capabilities available to Docker Compose without hardcoding.
|
||||
|
||||
## Model VRAM fit estimation
|
||||
|
||||
`model_vram_estimate()` answers "can this hardware run model X at quantization level Y?" by querying the HuggingFace Hub API for parameter count and architecture, then applying the standard VRAM formula (weights + KV cache + overhead). Reference algorithm from [LLMcalc](https://github.com/Raskoll2/LLMcalc) — no code copied, since LLMcalc has no published license.
|
||||
|
||||
```python
|
||||
from circuitforge_core.hardware import model_vram_estimate
|
||||
|
||||
est = model_vram_estimate("Qwen/Qwen2.5-7B-Instruct", "q4_k_m", available_vram_mb=8_000)
|
||||
print(est.total_vram_gb, est.fits) # e.g. 5.2 True
|
||||
```
|
||||
|
||||
Use cases: Avocet preflight (verify a checkpoint fits before benchmarking), cf-orch worker assignment (match model to GPU by VRAM fit), and onboarding wizards ("your GPU has X GB — here are models that will run well").
|
||||
|
||||
Raises `ModelVramLookupError` if the HF Hub API request fails or the model has no safetensors metadata; raises `ValueError` for an unrecognized `quant_level`.
|
||||
|
|
|
|||
169
tests/test_hardware/test_vram_estimate.py
Normal file
169
tests/test_hardware/test_vram_estimate.py
Normal file
|
|
@ -0,0 +1,169 @@
|
|||
"""Tests for circuitforge_core.hardware.vram_estimate (cf-core #64)."""
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from circuitforge_core.hardware.vram_estimate import (
|
||||
ModelVramLookupError,
|
||||
model_vram_estimate,
|
||||
)
|
||||
|
||||
_QWEN_CONFIG = {
|
||||
"num_hidden_layers": 28,
|
||||
"hidden_size": 3584,
|
||||
"num_attention_heads": 28,
|
||||
"num_key_value_heads": 4,
|
||||
}
|
||||
|
||||
|
||||
def _api_response(status_code=200, safetensors_total=7_000_000_000):
|
||||
resp = MagicMock()
|
||||
resp.status_code = status_code
|
||||
resp.json.return_value = {"safetensors": {"total": safetensors_total}} if safetensors_total else {}
|
||||
return resp
|
||||
|
||||
|
||||
def _config_response(status_code=200, config=None):
|
||||
resp = MagicMock()
|
||||
resp.status_code = status_code
|
||||
resp.json.return_value = config if config is not None else _QWEN_CONFIG
|
||||
return resp
|
||||
|
||||
|
||||
class TestModelVramEstimate:
|
||||
def test_estimates_weights_gb_for_fp16(self):
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
side_effect=[_api_response(), _config_response()],
|
||||
):
|
||||
est = model_vram_estimate("Qwen/Qwen2.5-7B-Instruct", "fp16")
|
||||
|
||||
assert est.params_billions == pytest.approx(7.0)
|
||||
# 7e9 params * 16 bits / 8 bits-per-byte / 1e9 = 14 GB
|
||||
assert est.weights_gb == pytest.approx(14.0)
|
||||
|
||||
def test_lower_bit_quant_uses_less_vram(self):
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
side_effect=[_api_response(), _config_response()],
|
||||
):
|
||||
fp16 = model_vram_estimate("Qwen/Qwen2.5-7B-Instruct", "fp16")
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
side_effect=[_api_response(), _config_response()],
|
||||
):
|
||||
q4 = model_vram_estimate("Qwen/Qwen2.5-7B-Instruct", "q4_k_m")
|
||||
|
||||
assert q4.weights_gb < fp16.weights_gb
|
||||
|
||||
def test_kv_cache_included_when_config_available(self):
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
side_effect=[_api_response(), _config_response()],
|
||||
):
|
||||
est = model_vram_estimate("Qwen/Qwen2.5-7B-Instruct", "fp16", ctx_len=8192)
|
||||
|
||||
assert est.kv_cache_gb > 0.0
|
||||
assert est.total_vram_gb == pytest.approx(
|
||||
est.weights_gb + est.kv_cache_gb + est.overhead_gb
|
||||
)
|
||||
|
||||
def test_kv_cache_zero_when_config_missing_arch_fields(self):
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
side_effect=[_api_response(), _config_response(config={"some_other_field": 1})],
|
||||
):
|
||||
est = model_vram_estimate("weird/model", "fp16")
|
||||
|
||||
assert est.kv_cache_gb == 0.0
|
||||
|
||||
def test_kv_cache_zero_when_config_fetch_fails(self):
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
side_effect=[_api_response(), _config_response(status_code=404)],
|
||||
):
|
||||
est = model_vram_estimate("Qwen/Qwen2.5-7B-Instruct", "fp16")
|
||||
|
||||
assert est.kv_cache_gb == 0.0
|
||||
|
||||
def test_fits_true_when_vram_sufficient(self):
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
side_effect=[_api_response(), _config_response()],
|
||||
):
|
||||
est = model_vram_estimate(
|
||||
"Qwen/Qwen2.5-7B-Instruct", "q4_k_m", available_vram_mb=24_000
|
||||
)
|
||||
|
||||
assert est.fits is True
|
||||
|
||||
def test_fits_false_when_vram_insufficient(self):
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
side_effect=[_api_response(), _config_response()],
|
||||
):
|
||||
est = model_vram_estimate(
|
||||
"Qwen/Qwen2.5-7B-Instruct", "fp32", available_vram_mb=4_000
|
||||
)
|
||||
|
||||
assert est.fits is False
|
||||
|
||||
def test_fits_none_when_available_vram_not_supplied(self):
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
side_effect=[_api_response(), _config_response()],
|
||||
):
|
||||
est = model_vram_estimate("Qwen/Qwen2.5-7B-Instruct", "fp16")
|
||||
|
||||
assert est.fits is None
|
||||
|
||||
def test_unknown_quant_level_raises_value_error(self):
|
||||
with pytest.raises(ValueError):
|
||||
model_vram_estimate("Qwen/Qwen2.5-7B-Instruct", "not-a-real-quant")
|
||||
|
||||
def test_quant_level_case_insensitive(self):
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
side_effect=[_api_response(), _config_response()],
|
||||
):
|
||||
est = model_vram_estimate("Qwen/Qwen2.5-7B-Instruct", "FP16")
|
||||
|
||||
assert est.weights_gb == pytest.approx(14.0)
|
||||
|
||||
def test_raises_on_non_200_from_hf_api(self):
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
return_value=_api_response(status_code=404),
|
||||
):
|
||||
with pytest.raises(ModelVramLookupError):
|
||||
model_vram_estimate("nonexistent/model", "fp16")
|
||||
|
||||
def test_raises_on_missing_safetensors_metadata(self):
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
return_value=_api_response(safetensors_total=None),
|
||||
):
|
||||
with pytest.raises(ModelVramLookupError):
|
||||
model_vram_estimate("gguf-only/model", "fp16")
|
||||
|
||||
def test_raises_on_request_exception(self):
|
||||
import requests
|
||||
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
side_effect=requests.ConnectionError("no network"),
|
||||
):
|
||||
with pytest.raises(ModelVramLookupError):
|
||||
model_vram_estimate("Qwen/Qwen2.5-7B-Instruct", "fp16")
|
||||
|
||||
def test_result_includes_hf_model_id_and_quant_level(self):
|
||||
with patch(
|
||||
"circuitforge_core.hardware.vram_estimate.requests.get",
|
||||
side_effect=[_api_response(), _config_response()],
|
||||
):
|
||||
est = model_vram_estimate("Qwen/Qwen2.5-7B-Instruct", "q4_k_m")
|
||||
|
||||
assert est.hf_model_id == "Qwen/Qwen2.5-7B-Instruct"
|
||||
assert est.quant_level == "q4_k_m"
|
||||
Loading…
Reference in a new issue