237 lines
7.4 KiB
Python
237 lines
7.4 KiB
Python
"""Curated list of models supported by the network with VRAM requirements."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
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
|
|
|
|
|
|
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="Moonshot/Unsloth coding-focused MoE model; 1T total, 32B activated",
|
|
metadata={
|
|
"architecture": "Mixture-of-Experts (MoE)",
|
|
"total_parameters": "1T",
|
|
"activated_parameters": "32B",
|
|
"num_layers": 61,
|
|
"dense_layers": 1,
|
|
"attention_hidden_dimension": 7168,
|
|
"moe_hidden_dimension_per_expert": 2048,
|
|
"attention_heads": 64,
|
|
"experts": 384,
|
|
"selected_experts_per_token": 8,
|
|
"shared_experts": 1,
|
|
"vocabulary_size": 160000,
|
|
"context_length": 256000,
|
|
"attention_mechanism": "MLA",
|
|
"activation_function": "SwiGLU",
|
|
"vision_encoder": "MoonViT",
|
|
"vision_encoder_parameters": "400M",
|
|
"license": "modified-mit",
|
|
"native_quantization": "int4",
|
|
"recommended_engines": ["vLLM", "SGLang", "KTransformers"],
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
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 int(cfg.num_hidden_layers)
|
|
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]
|
|
|
|
cfg = AutoConfig.from_pretrained(
|
|
hf_repo,
|
|
cache_dir=str(cache_dir) if cache_dir is not None else None,
|
|
)
|
|
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 not None:
|
|
metadata[key] = value
|
|
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
|