popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Track Kimi model metadata and cache path

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-01 12:38:31 +02:00
parent 78834e5045
commit bc760c1694
7 changed files with 231 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
packages/node/meshnet_node/model_backend.py
View File

@@ -4,6 +4,7 @@ from __future__ import annotations
import base64
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal
Quantization = Literal["bfloat16", "int8", "nf4"]
@@ -65,6 +66,7 @@ class TorchModelShard:
shard_start: int,
shard_end: int,
quantization: Quantization = "bfloat16",
cache_dir: Path | None = None,
) -> None:
if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
raise ValueError("shard_start must be <= shard_end and non-negative")
@@ -89,9 +91,9 @@ class TorchModelShard:
model_id,
quantization_config=quant_config,
device_map="auto" if quant_config is not None else None,
torch_dtype=torch.bfloat16,
dtype=torch.bfloat16,
low_cpu_mem_usage=True,
use_safetensors=True,
cache_dir=str(cache_dir) if cache_dir is not None else None,
)
if quant_config is None:
self.model.to(self.device)
@@ -104,7 +106,10 @@ class TorchModelShard:
raise
self.model.eval()
self.tokenizer = AutoTokenizer.from_pretrained(model_id)
self.tokenizer = AutoTokenizer.from_pretrained(
model_id,
cache_dir=str(cache_dir) if cache_dir is not None else None,
)
self.layers = _model_layers(self.model)
self.total_layers = len(self.layers)
# shard_end is INCLUSIVE (last layer index, 0-based), matching the CLI convention.
@@ -336,8 +341,9 @@ def load_torch_shard(
shard_start: int,
shard_end: int,
quantization: Quantization = "bfloat16",
cache_dir: Path | None = None,
) -> TorchModelShard:
return TorchModelShard(model_id, shard_start, shard_end, quantization)
return TorchModelShard(model_id, shard_start, shard_end, quantization, cache_dir)
def _model_layers(model: Any) -> Any:

71
packages/node/meshnet_node/model_catalog.py
View File

@@ -3,6 +3,7 @@
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
@dataclass
@@ -15,6 +16,7 @@ class ModelPreset:
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."""
@@ -123,6 +125,37 @@ CURATED_MODELS: list[ModelPreset] = [
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"],
},
),
]
@@ -140,6 +173,44 @@ def detect_num_layers(hf_repo: str) -> int | None:
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:

17
packages/node/meshnet_node/startup.py
View File

@@ -15,6 +15,7 @@ from typing import Any
from .downloader import compute_shard_checksum, download_shard
from .hardware import detect_hardware, benchmark_throughput_checked
from .model_catalog import model_metadata_for
from .relay_bridge import RelayHttpBridge, peer_id_from_wallet
from .server import StubNodeServer
from .torch_server import TorchNodeServer
@@ -422,7 +423,10 @@ def run_startup(
user_pinned_shard = shard_start is not None or shard_end is not None
# Auto-detect shard range from model config if not explicitly provided
if shard_start is None or shard_end is None:
detected = _detect_num_layers(model_id)
try:
detected = _detect_num_layers(model_id, cache_dir=cache_dir)
except TypeError:
detected = _detect_num_layers(model_id)
if detected is None:
raise ValueError(
f"Could not read num_hidden_layers from {model_id} config. "
@@ -459,6 +463,7 @@ def run_startup(
quantization=quantization,
tracker_url=tracker_url,
route_timeout=route_timeout,
cache_dir=cache_dir,
debug=debug,
)
_node_start_time = time.monotonic()
@@ -495,6 +500,7 @@ def run_startup(
"score": 1.0,
"tracker_mode": (shard_start == 0),
"managed_assignment": not user_pinned_shard,
"model_metadata": model_metadata_for(model_id, total_layers, cache_dir=cache_dir),
**registration_capabilities,
**relay_fields,
}
@@ -559,6 +565,7 @@ def run_startup(
quantization=quantization,
tracker_url=tracker_url,
route_timeout=route_timeout,
cache_dir=cache_dir,
debug=debug,
)
_node_start_time = time.monotonic()
@@ -587,6 +594,7 @@ def run_startup(
"score": 1.0,
"tracker_mode": (assigned_shard_start == 0),
"managed_assignment": True,
"model_metadata": model_metadata_for(assigned_hf_repo, assigned_num_layers, cache_dir=cache_dir),
**registration_capabilities,
**relay_fields,
}
@@ -722,11 +730,14 @@ def run_startup(
return node
def _detect_num_layers(model_id: str) -> int | None:
def _detect_num_layers(model_id: str, cache_dir: Path | None = None) -> int | None:
"""Fetch num_hidden_layers from HuggingFace model config (downloads ~1 KB config.json only)."""
try:
from transformers import AutoConfig # type: ignore[import]
cfg = AutoConfig.from_pretrained(model_id)
cfg = AutoConfig.from_pretrained(
model_id,
cache_dir=str(cache_dir) if cache_dir is not None else None,
)
return int(cfg.num_hidden_layers)
except Exception as exc:
print(f" Warning: could not read model config from HF: {exc}", flush=True)

12
packages/node/meshnet_node/torch_server.py
View File

@@ -12,6 +12,7 @@ import urllib.error
import urllib.parse
import urllib.request
import uuid
from pathlib import Path
from typing import Any
from .model_backend import (
@@ -682,6 +683,7 @@ class TorchNodeServer:
tracker_mode: bool | None = None,
tracker_url: str | None = None,
route_timeout: float = 30.0,
cache_dir: Path | None = None,
debug: bool = False,
) -> None:
self._host = host
@@ -691,11 +693,13 @@ class TorchNodeServer:
shard_start,
shard_end,
quantization,
cache_dir,
)
# Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
self._tracker_mode = tracker_mode if tracker_mode is not None else (shard_start == 0)
self._tracker_url = tracker_url
self._route_timeout = route_timeout
self._cache_dir = cache_dir
self._debug = debug
self._server: _TorchHTTPServer | None = None
self._thread: threading.Thread | None = None
@@ -745,7 +749,10 @@ class TorchNodeServer:
f" [node] loading reassigned shard: {model_id} layers {shard_start}-{shard_end}",
flush=True,
)
new_backend = _load_backend(model_id, shard_start, shard_end, quantization)
try:
new_backend = _load_backend(model_id, shard_start, shard_end, quantization, self._cache_dir)
except TypeError:
new_backend = _load_backend(model_id, shard_start, shard_end, quantization)
self._backend = new_backend
self._tracker_mode = shard_start == 0
if self._server is not None:
@@ -797,12 +804,13 @@ def _load_backend(
shard_start: int,
shard_end: int,
quantization: str,
cache_dir: Path | None = None,
) -> TorchModelShard:
from .model_backend import load_torch_shard
quant = validate_quantization(quantization)
try:
return load_torch_shard(model_id, shard_start, shard_end, quant)
return load_torch_shard(model_id, shard_start, shard_end, quant, cache_dir)
except MissingModelDependencyError:
raise
except InsufficientVRAMError as exc: