282 lines
8.8 KiB
Python
282 lines
8.8 KiB
Python
"""Curated list of models supported by the network with VRAM requirements."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import dataclass
|
|
from importlib.resources import files
|
|
from pathlib import Path
|
|
|
|
|
|
@dataclass
|
|
class ModelPreset:
|
|
name: str
|
|
hf_repo: str
|
|
num_layers: int
|
|
# VRAM in GB at each quantization level (None = too large to quantize this way)
|
|
vram_nf4: float
|
|
vram_int8: float
|
|
vram_bf16: float
|
|
description: str
|
|
metadata: dict | None = None
|
|
|
|
def vram_for_quant(self, quant: str) -> float:
|
|
"""Return VRAM requirement in GB for the given quantization."""
|
|
q = quant.lower().replace("bfloat16", "bf16")
|
|
if q == "nf4":
|
|
return self.vram_nf4
|
|
if q in ("int8", "int8"):
|
|
return self.vram_int8
|
|
if q in ("bf16", "bfloat16"):
|
|
return self.vram_bf16
|
|
raise ValueError(f"unknown quantization: {quant!r}")
|
|
|
|
def fits_vram(self, available_gb: float, quant: str) -> bool:
|
|
return self.vram_for_quant(quant) <= available_gb
|
|
|
|
def recommended_quant(self, available_gb: float) -> str | None:
|
|
"""Return the highest-quality quantization that fits available VRAM, or None."""
|
|
if self.vram_bf16 <= available_gb:
|
|
return "bf16"
|
|
if self.vram_int8 <= available_gb:
|
|
return "int8"
|
|
if self.vram_nf4 <= available_gb:
|
|
return "nf4"
|
|
return None
|
|
|
|
|
|
def _load_model_metadata() -> dict[str, dict]:
|
|
try:
|
|
raw = files("meshnet_node").joinpath("model_metadata.json").read_text()
|
|
data = json.loads(raw)
|
|
except Exception:
|
|
return {}
|
|
models = data.get("models", {})
|
|
if not isinstance(models, dict):
|
|
return {}
|
|
return {
|
|
str(repo): metadata
|
|
for repo, metadata in models.items()
|
|
if isinstance(metadata, dict)
|
|
}
|
|
|
|
|
|
_MODEL_METADATA = _load_model_metadata()
|
|
|
|
|
|
def _local_model_path(hf_repo: str, cache_dir: Path | None) -> Path | None:
|
|
if cache_dir is None:
|
|
return None
|
|
if (cache_dir / "config.json").exists():
|
|
return cache_dir
|
|
candidate = cache_dir / hf_repo.split("/")[-1]
|
|
if (candidate / "config.json").exists():
|
|
return candidate
|
|
return None
|
|
|
|
|
|
CURATED_MODELS: list[ModelPreset] = [
|
|
ModelPreset(
|
|
name="Qwen2.5-0.5B-Instruct",
|
|
hf_repo="Qwen/Qwen2.5-0.5B-Instruct",
|
|
num_layers=24,
|
|
vram_nf4=0.4,
|
|
vram_int8=0.6,
|
|
vram_bf16=1.0,
|
|
description="Smallest no-gating model — great for testing, ~1 GB",
|
|
),
|
|
ModelPreset(
|
|
name="Qwen2.5-1.5B-Instruct",
|
|
hf_repo="Qwen/Qwen2.5-1.5B-Instruct",
|
|
num_layers=28,
|
|
vram_nf4=1.0,
|
|
vram_int8=1.8,
|
|
vram_bf16=3.2,
|
|
description="Fast no-gating model — good quality, ~3 GB",
|
|
),
|
|
ModelPreset(
|
|
name="Llama-3-70B-Instruct",
|
|
hf_repo="meta-llama/Meta-Llama-3-70B-Instruct",
|
|
num_layers=80,
|
|
vram_nf4=18.0,
|
|
vram_int8=40.0,
|
|
vram_bf16=140.0,
|
|
description="Meta's flagship 70B instruction model",
|
|
),
|
|
ModelPreset(
|
|
name="Qwen2.5-72B-Instruct",
|
|
hf_repo="Qwen/Qwen2.5-72B-Instruct",
|
|
num_layers=80,
|
|
vram_nf4=19.0,
|
|
vram_int8=41.0,
|
|
vram_bf16=145.0,
|
|
description="Alibaba's 72B multilingual instruction model",
|
|
),
|
|
ModelPreset(
|
|
name="Mixtral-8x7B-Instruct",
|
|
hf_repo="mistralai/Mixtral-8x7B-Instruct-v0.1",
|
|
num_layers=32,
|
|
vram_nf4=7.0,
|
|
vram_int8=14.0,
|
|
vram_bf16=27.0,
|
|
description="Mistral's sparse MoE — fast and efficient",
|
|
),
|
|
ModelPreset(
|
|
name="Llama-3-8B-Instruct",
|
|
hf_repo="meta-llama/Meta-Llama-3-8B-Instruct",
|
|
num_layers=32, # gated repo — requires HF login
|
|
vram_nf4=4.5,
|
|
vram_int8=8.5,
|
|
vram_bf16=16.0,
|
|
description="Meta's compact 8B model — good for low-VRAM nodes",
|
|
),
|
|
ModelPreset(
|
|
name="Phi-3-medium-128k",
|
|
hf_repo="microsoft/Phi-3-medium-128k-instruct",
|
|
num_layers=40,
|
|
vram_nf4=4.0,
|
|
vram_int8=8.0,
|
|
vram_bf16=15.0,
|
|
description="Microsoft's efficient 14B model with 128k context",
|
|
),
|
|
ModelPreset(
|
|
name="Gemma-2-27B-IT",
|
|
hf_repo="google/gemma-2-27b-it",
|
|
num_layers=46,
|
|
vram_nf4=10.0,
|
|
vram_int8=20.0,
|
|
vram_bf16=54.0,
|
|
description="Google's 27B instruction-tuned model",
|
|
),
|
|
ModelPreset(
|
|
name="DeepSeek-V2-Lite-Chat",
|
|
hf_repo="deepseek-ai/DeepSeek-V2-Lite-Chat",
|
|
num_layers=27,
|
|
vram_nf4=5.0,
|
|
vram_int8=9.0,
|
|
vram_bf16=16.0,
|
|
description="DeepSeek's efficient MoE — strong coding + reasoning",
|
|
),
|
|
ModelPreset(
|
|
name="Kimi-K2.7-Code",
|
|
hf_repo="unsloth/Kimi-K2.7-Code",
|
|
num_layers=61,
|
|
vram_nf4=500.0,
|
|
vram_int8=1000.0,
|
|
vram_bf16=2000.0,
|
|
description="Large coding-focused MoE model",
|
|
metadata=_MODEL_METADATA.get("unsloth/Kimi-K2.7-Code"),
|
|
),
|
|
]
|
|
|
|
|
|
def layers_from_config(cfg) -> int | None:
|
|
"""Extract the transformer layer count from a HuggingFace config object.
|
|
|
|
Composite configs (vision-language, some MoE) nest the decoder settings in
|
|
``text_config`` — e.g. Qwen3.5-MoE has no top-level ``num_hidden_layers``.
|
|
"""
|
|
candidates = [cfg]
|
|
get_text_config = getattr(cfg, "get_text_config", None)
|
|
if callable(get_text_config):
|
|
try:
|
|
candidates.append(get_text_config())
|
|
except Exception:
|
|
pass
|
|
nested = getattr(cfg, "text_config", None)
|
|
if nested is not None:
|
|
candidates.append(nested)
|
|
for candidate in candidates:
|
|
for attr in ("num_hidden_layers", "num_layers", "n_layer", "n_layers"):
|
|
value = getattr(candidate, attr, None)
|
|
if value is not None:
|
|
return int(value)
|
|
return None
|
|
|
|
|
|
def detect_num_layers(hf_repo: str) -> int | None:
|
|
"""Return num_hidden_layers from HuggingFace config.json (downloads ~1 KB only)."""
|
|
# Check curated list first (no network call)
|
|
for m in CURATED_MODELS:
|
|
if m.hf_repo == hf_repo:
|
|
return m.num_layers
|
|
try:
|
|
from transformers import AutoConfig # type: ignore[import]
|
|
cfg = AutoConfig.from_pretrained(hf_repo)
|
|
return layers_from_config(cfg)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def model_metadata_for(
|
|
hf_repo: str,
|
|
num_layers: int | None = None,
|
|
cache_dir: Path | None = None,
|
|
) -> dict:
|
|
"""Return operator-facing model metadata for a HuggingFace repo."""
|
|
for model in CURATED_MODELS:
|
|
if model.hf_repo == hf_repo:
|
|
metadata = dict(model.metadata or {})
|
|
metadata.setdefault("num_layers", model.num_layers)
|
|
return metadata
|
|
|
|
metadata: dict = {}
|
|
if num_layers is not None:
|
|
metadata["num_layers"] = num_layers
|
|
try:
|
|
from transformers import AutoConfig # type: ignore[import]
|
|
|
|
local_model = _local_model_path(hf_repo, cache_dir)
|
|
load_source = str(local_model) if local_model is not None else hf_repo
|
|
cfg = AutoConfig.from_pretrained(
|
|
load_source,
|
|
cache_dir=str(cache_dir) if cache_dir is not None and local_model is None else None,
|
|
)
|
|
# Composite configs (VLM/MoE) nest decoder fields in text_config.
|
|
text_cfg = getattr(cfg, "text_config", None) or cfg
|
|
for attr, key in (
|
|
("model_type", "architecture"),
|
|
("num_hidden_layers", "num_layers"),
|
|
("hidden_size", "hidden_size"),
|
|
("num_attention_heads", "attention_heads"),
|
|
("vocab_size", "vocabulary_size"),
|
|
("max_position_embeddings", "context_length"),
|
|
):
|
|
value = getattr(cfg, attr, None)
|
|
if value is None:
|
|
value = getattr(text_cfg, attr, None)
|
|
if value is not None:
|
|
metadata[key] = value
|
|
if "num_layers" not in metadata:
|
|
layers = layers_from_config(cfg)
|
|
if layers is not None:
|
|
metadata["num_layers"] = layers
|
|
except Exception:
|
|
pass
|
|
return metadata
|
|
|
|
|
|
def browse_hf_hub(top_n: int = 20) -> list[dict]:
|
|
"""Fetch top downloaded text-generation models from HuggingFace Hub."""
|
|
try:
|
|
from huggingface_hub import list_models # type: ignore[import]
|
|
|
|
models = list(
|
|
list_models(
|
|
pipeline_tag="text-generation",
|
|
library="transformers",
|
|
sort="downloads",
|
|
direction=-1,
|
|
limit=top_n,
|
|
)
|
|
)
|
|
return [
|
|
{
|
|
"repo": m.id,
|
|
"downloads": getattr(m, "downloads", 0) or 0,
|
|
}
|
|
for m in models
|
|
]
|
|
except Exception as exc:
|
|
raise RuntimeError(f"HuggingFace Hub lookup failed: {exc}") from exc
|