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

new tasks, model pricing, auto quantisation, etc...

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-06 17:11:53 +03:00
parent 7f67e29d76
commit ccb69c41e3
14 changed files with 466 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
packages/node/meshnet_node/cli.py
View File

@@ -56,7 +56,7 @@ def _run_node(cfg: dict) -> None:
model_id=cfg.get("model_hf_repo") or None,
shard_start=cfg.get("shard_start"),
shard_end=cfg.get("shard_end"),
quantization=cfg.get("quantization", "int8").replace("bf16", "bfloat16"),
quantization=cfg.get("quantization", "auto").replace("bf16", "bfloat16"),
wallet_path=Path(cfg["wallet_path"]) if cfg.get("wallet_path") else None,
cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
host=cfg.get("host", "0.0.0.0"),
@@ -199,17 +199,19 @@ def _cmd_config(args) -> int:
def _cmd_start(args) -> int:
"""Legacy `start` subcommand — preserves backward compatibility with existing tests."""
from .config import load_config, DEFAULTS
from .config import DEFAULTS
# Build a transient config from flags (don't write to disk)
cfg = dict(DEFAULTS)
cfg["tracker_url"] = args.tracker
if args.tracker:
cfg["tracker_url"] = args.tracker
cfg["port"] = args.port if args.port is not None else _first_available_port(args.host)
if args.model_id is None and "/" in args.model:
cfg["model_hf_repo"] = args.model
cfg["model_name"] = args.model.split("/")[-1]
model = args.model or cfg.get("model_hf_repo") or cfg.get("model_name") or "stub-model"
if args.model_id is None and "/" in model:
cfg["model_hf_repo"] = model
cfg["model_name"] = model.split("/")[-1]
else:
cfg["model_name"] = args.model
cfg["model_name"] = model
cfg["quantization"] = args.quantization
cfg["host"] = args.host
if args.model_id:
@@ -307,13 +309,13 @@ def main() -> None:
# start subcommand (legacy / backward-compat)
start_cmd = subparsers.add_parser("start", help="Start node (legacy flags)")
start_cmd.add_argument("--tracker", default="http://localhost:8080")
start_cmd.add_argument("--tracker")
start_cmd.add_argument("--port", type=int)
start_cmd.add_argument("--model", default="stub-model")
start_cmd.add_argument("--model")
start_cmd.add_argument("--model-id", help="HuggingFace repo ID")
start_cmd.add_argument("--shard-start", type=int)
start_cmd.add_argument("--shard-end", type=int)
start_cmd.add_argument("--quantization", choices=["bfloat16", "int8", "nf4", "bf16"], default="int8")
start_cmd.add_argument("--quantization", choices=["auto", "bfloat16", "int8", "nf4", "bf16"], default="auto")
start_cmd.add_argument("--host", default="0.0.0.0")
start_cmd.add_argument("--advertise-host")
start_cmd.add_argument("--tracker-mode", action="store_true")

11
packages/node/meshnet_node/config.py
View File

@@ -14,13 +14,16 @@ _DEFAULT_CONFIG_FILE = _DEFAULT_CONFIG_DIR / "config.json"
_DEFAULT_DOWNLOAD_DIR = Path(
os.environ.get("MESHNET_DOWNLOAD_DIR", str(Path.home() / ".meshnet" / "models"))
)
_DEFAULT_TRACKER_URL = "http://localhost:8080"
_DEFAULT_TRACKER_URL = os.environ.get("MESHNET_TRACKER_URL", "http://localhost:8080")
_DEFAULT_WALLET_PATH = str(Path.home() / ".config" / "meshnet" / "wallet.json")
_DEFAULT_QUANTIZATION = "nf4"
_DEFAULT_QUANTIZATION = "auto"
_DEFAULT_MODEL = os.environ.get("MESHNET_MODEL_ID") or os.environ.get("MESHNET_MODEL", "")
_DEFAULT_MODEL_HF_REPO = _DEFAULT_MODEL if "/" in _DEFAULT_MODEL else ""
_DEFAULT_MODEL_NAME = _DEFAULT_MODEL.split("/")[-1] if "/" in _DEFAULT_MODEL else _DEFAULT_MODEL
DEFAULTS = {
"model_hf_repo": "",
"model_name": "",
"model_hf_repo": _DEFAULT_MODEL_HF_REPO,
"model_name": _DEFAULT_MODEL_NAME,
"quantization": _DEFAULT_QUANTIZATION,
"download_dir": str(_DEFAULT_DOWNLOAD_DIR),
"tracker_url": _DEFAULT_TRACKER_URL,

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

@@ -7,7 +7,7 @@ from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal
Quantization = Literal["bfloat16", "int8", "nf4"]
Quantization = Literal["auto", "bfloat16", "int8", "nf4"]
class ModelBackendError(RuntimeError):
@@ -31,14 +31,14 @@ class TensorPayload:
def validate_quantization(value: str) -> Quantization:
if value not in {"bfloat16", "int8", "nf4"}:
raise ValueError("quantization must be one of: bfloat16, int8, nf4")
if value not in {"auto", "bfloat16", "int8", "nf4"}:
raise ValueError("quantization must be one of: auto, bfloat16, int8, nf4")
return value # type: ignore[return-value]
def build_quantization_config(quantization: Quantization) -> Any | None:
"""Return a transformers BitsAndBytesConfig for quantized weights."""
if quantization == "bfloat16":
if quantization in {"auto", "bfloat16"}:
return None
try:
import torch
@@ -65,7 +65,7 @@ class TorchModelShard:
model_id: str,
shard_start: int,
shard_end: int,
quantization: Quantization = "bfloat16",
quantization: Quantization = "auto",
cache_dir: Path | None = None,
) -> None:
if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
@@ -77,7 +77,7 @@ class TorchModelShard:
try:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
except ModuleNotFoundError as exc:
raise MissingModelDependencyError(
"real model backend requires torch, transformers, safetensors, accelerate, and bitsandbytes"
@@ -85,17 +85,27 @@ class TorchModelShard:
self.torch = torch
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
quant_config = build_quantization_config(quantization)
quant_config, dtype, uses_quantized_weights = _model_load_plan(
AutoConfig,
model_id,
quantization,
torch,
cache_dir,
)
try:
load_kwargs = {
"device_map": "auto" if uses_quantized_weights else None,
"dtype": dtype,
"low_cpu_mem_usage": True,
"cache_dir": str(cache_dir) if cache_dir is not None else None,
}
if quant_config is not None:
load_kwargs["quantization_config"] = quant_config
self.model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config=quant_config,
device_map="auto" if quant_config is not None else None,
dtype=torch.bfloat16,
low_cpu_mem_usage=True,
cache_dir=str(cache_dir) if cache_dir is not None else None,
**load_kwargs,
)
if quant_config is None:
if not uses_quantized_weights:
self.model.to(self.device)
except Exception as exc:
if _looks_like_oom(exc):
@@ -340,12 +350,71 @@ def load_torch_shard(
model_id: str,
shard_start: int,
shard_end: int,
quantization: Quantization = "bfloat16",
quantization: Quantization = "auto",
cache_dir: Path | None = None,
) -> TorchModelShard:
return TorchModelShard(model_id, shard_start, shard_end, quantization, cache_dir)
def _model_load_plan(
auto_config: Any,
model_id: str,
quantization: Quantization,
torch: Any,
cache_dir: Path | None = None,
) -> tuple[Any | None, Any, bool]:
"""Return (explicit quant config, dtype, uses quantized weights)."""
if quantization != "auto":
quant_config = build_quantization_config(quantization)
return quant_config, torch.bfloat16, quant_config is not None
cfg = auto_config.from_pretrained(
model_id,
cache_dir=str(cache_dir) if cache_dir is not None else None,
)
if _native_quantization_config(cfg) is not None:
return None, _native_torch_dtype(cfg, torch), True
return None, _native_torch_dtype(cfg, torch), False
def _config_candidates(cfg: Any) -> list[Any]:
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
text_config = getattr(cfg, "text_config", None)
if text_config is not None:
candidates.append(text_config)
return candidates
def _native_quantization_config(cfg: Any) -> Any | None:
for candidate in _config_candidates(cfg):
quant_config = getattr(candidate, "quantization_config", None)
if quant_config:
return quant_config
return None
def _native_torch_dtype(cfg: Any, torch: Any) -> Any:
for candidate in _config_candidates(cfg):
for attr in ("dtype", "torch_dtype"):
dtype = getattr(candidate, attr, None)
if dtype is None:
continue
if isinstance(dtype, str):
dtype_name = dtype.removeprefix("torch.")
dtype_value = getattr(torch, dtype_name, None)
if dtype_value is not None:
return dtype_value
else:
return dtype
return torch.bfloat16
def _model_layers(model: Any) -> Any:
if hasattr(model, "model") and hasattr(model.model, "layers"):
return model.model.layers

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

@@ -52,6 +52,7 @@ def _max_assignable_layers(memory_mb: int, total_layers: int | None) -> int:
def _shard_budget_line(memory_mb: int, memory_source: str, total_layers: int | None, quantization: str) -> str:
memory_gb = memory_mb / 1024
gb_str = f"{memory_gb:.1f} GB"
budget_quantization = "bfloat16" if quantization == "auto" else quantization
if total_layers is None or total_layers <= 0:
return f"Memory budget: {gb_str} {memory_source}; shard budget: unknown model layer count"
max_layers = _max_assignable_layers(memory_mb, total_layers)
@@ -61,7 +62,7 @@ def _shard_budget_line(memory_mb: int, memory_source: str, total_layers: int | N
remaining_str = f"; {remaining_gb:.1f} GB remaining after full load" if remaining_gb > 1 else ""
return (
f"Memory budget: {gb_str} {memory_source}; "
f"Shard budget: up to {max_layers}/{total_layers} layers at {quantization}"
f"Shard budget: up to {max_layers}/{total_layers} layers at {budget_quantization}"
f"{remaining_str}"
)
@@ -306,7 +307,7 @@ def run_startup(
model_id: str | None = None,
shard_start: int | None = None,
shard_end: int | None = None,
quantization: str = "bfloat16",
quantization: str = "auto",
wallet_path: Path | None = None,
cache_dir: Path | None = None,
host: str = "127.0.0.1",