diff --git a/packages/node/meshnet_node/model_catalog.py b/packages/node/meshnet_node/model_catalog.py index 5393435..9852476 100644 --- a/packages/node/meshnet_node/model_catalog.py +++ b/packages/node/meshnet_node/model_catalog.py @@ -42,6 +42,24 @@ class ModelPreset: 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", @@ -72,7 +90,7 @@ CURATED_MODELS: list[ModelPreset] = [ ModelPreset( name="Llama-3-8B-Instruct", hf_repo="meta-llama/Meta-Llama-3-8B-Instruct", - num_layers=32, + num_layers=32, # gated repo — requires HF login vram_nf4=4.5, vram_int8=8.5, vram_bf16=16.0, @@ -108,6 +126,20 @@ CURATED_MODELS: list[ModelPreset] = [ ] +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 browse_hf_hub(top_n: int = 20) -> list[dict]: """Fetch top downloaded text-generation models from HuggingFace Hub.""" try: diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index 62d896a..036b7ff 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -84,7 +84,19 @@ def run_startup( if probationary_line is not None: print(f" {probationary_line}", flush=True) - if model_id is not None and shard_start is not None and shard_end is not None: + if model_id 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) + if detected is None: + raise ValueError( + f"Could not read num_hidden_layers from {model_id} config. " + "Pass --shard-start and --shard-end explicitly." + ) + shard_start = shard_start if shard_start is not None else 0 + shard_end = shard_end if shard_end is not None else detected - 1 + print(f" Auto-detected {detected} layers → shard {shard_start}–{shard_end}", flush=True) + print("Loading real PyTorch model shard...", flush=True) node = TorchNodeServer( host=host, @@ -102,7 +114,7 @@ def run_startup( f"meshnet-node ready\n" f" Wallet: {address}\n" f" Model ID: {model_id}\n" - f" Shard: layers {shard_start}-{shard_end}\n" + f" Shard: layers {shard_start}–{shard_end}\n" f" Quantization: {quantization}\n" f" Endpoint: {endpoint}\n" f" Hardware: {device.upper()}\n" @@ -110,8 +122,8 @@ def run_startup( flush=True, ) return node - if model_id is not None or shard_start is not None or shard_end is not None: - raise ValueError("--model-id, --shard-start, and --shard-end must be provided together") + if shard_start is not None or shard_end is not None: + raise ValueError("--shard-start / --shard-end require --model-id") # 3. Shard assignment from tracker print("Querying tracker for shard assignment...", flush=True) @@ -201,6 +213,17 @@ def run_startup( return node +def _detect_num_layers(model_id: str) -> 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) + return int(cfg.num_hidden_layers) + except Exception as exc: + print(f" Warning: could not read model config from HF: {exc}", flush=True) + return None + + def _probationary_status_line(contracts: Any | None, wallet_address: str) -> str | None: if contracts is None: return None diff --git a/packages/node/meshnet_node/wizard.py b/packages/node/meshnet_node/wizard.py index 624d810..5edf1d3 100644 --- a/packages/node/meshnet_node/wizard.py +++ b/packages/node/meshnet_node/wizard.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import TYPE_CHECKING from .config import DEFAULTS, _DEFAULT_DOWNLOAD_DIR, _DEFAULT_TRACKER_URL, _DEFAULT_WALLET_PATH -from .model_catalog import CURATED_MODELS, ModelPreset, browse_hf_hub +from .model_catalog import CURATED_MODELS, ModelPreset, browse_hf_hub, detect_num_layers if TYPE_CHECKING: pass @@ -239,6 +239,13 @@ def run_wizard(config_path_override=None) -> dict: if choice == len(CURATED_MODELS) + 1: repo = _browse_hf_interactive() if repo: + # Look up layer count for custom repo + print(f" Checking {repo} config...", end=" ", flush=True) + layers = detect_num_layers(repo) + if layers: + print(f"{layers} layers") + else: + print("(layer count unknown — will detect on start)") selected_repo = repo selected_preset = None else: @@ -254,7 +261,10 @@ def run_wizard(config_path_override=None) -> dict: selected_repo = None selected_preset = None - print(f"\n ✓ Selected: {selected_repo}") + num_layers = (selected_preset.num_layers if selected_preset + else detect_num_layers(selected_repo or "")) + layers_str = f" {num_layers} layers" if num_layers else "" + print(f"\n ✓ Selected: {selected_repo}{layers_str}") # Step 3b: Quantization quant = _ask_quant(gpus, selected_preset) diff --git a/tests/test_mining_cli.py b/tests/test_mining_cli.py index 94498bf..211ea06 100644 --- a/tests/test_mining_cli.py +++ b/tests/test_mining_cli.py @@ -21,7 +21,7 @@ def test_curated_models_list_is_non_empty(): def test_model_preset_vram_for_quant(): from meshnet_node.model_catalog import CURATED_MODELS - m = CURATED_MODELS[0] # Llama-3-70B + m = next(m for m in CURATED_MODELS if "Llama-3-70B" in m.name) assert m.vram_for_quant("nf4") == m.vram_nf4 assert m.vram_for_quant("int8") == m.vram_int8 assert m.vram_for_quant("bf16") == m.vram_bf16 @@ -39,7 +39,8 @@ def test_model_preset_fits_vram(): def test_recommended_quant_respects_vram(): from meshnet_node.model_catalog import CURATED_MODELS - m = CURATED_MODELS[0] # Llama-3-70B: nf4=18, int8=40, bf16=140 + m = next(m for m in CURATED_MODELS if "Llama-3-70B" in m.name) + # nf4=18, int8=40, bf16=140 assert m.recommended_quant(200) == "bf16" assert m.recommended_quant(50) == "int8" assert m.recommended_quant(20) == "nf4" @@ -125,9 +126,9 @@ def test_wizard_writes_config_on_happy_path(tmp_path, monkeypatch): # Tracker not reachable (stub) monkeypatch.setattr(wiz, "_ping_tracker", lambda url: False) - # Simulate user selecting model 3 (Mixtral), quant 1 (nf4), default dir, default tracker, default wallet + # Simulate user selecting model 1 (Qwen2.5-0.5B), quant 1 (nf4), default dir, default tracker, default wallet inputs = iter([ - "3", # pick Mixtral (index 3 in CURATED_MODELS) + "1", # pick Qwen2.5-0.5B-Instruct (index 1 in CURATED_MODELS) "1", # quant NF4 str(tmp_path / "models"), # download dir "http://localhost:8080", # tracker @@ -136,7 +137,7 @@ def test_wizard_writes_config_on_happy_path(tmp_path, monkeypatch): monkeypatch.setattr("builtins.input", lambda prompt="": next(inputs)) cfg = wiz.run_wizard(config_path_override=tmp_path / "config.json") - assert cfg["model_hf_repo"] == "mistralai/Mixtral-8x7B-Instruct-v0.1" + assert cfg["model_hf_repo"] == "Qwen/Qwen2.5-0.5B-Instruct" assert cfg["quantization"] == "nf4" assert "download_dir" in cfg assert cfg["tracker_url"] == "http://localhost:8080" @@ -265,6 +266,63 @@ def test_config_command_prints_saved_config(tmp_path, monkeypatch, capsys): assert data["model_hf_repo"] == saved["model_hf_repo"] +def test_detect_num_layers_returns_catalog_value_without_network(monkeypatch): + """detect_num_layers uses the curated catalog first — no network call.""" + from meshnet_node.model_catalog import detect_num_layers + + # Qwen2.5-0.5B is in the catalog with 24 layers + layers = detect_num_layers("Qwen/Qwen2.5-0.5B-Instruct") + assert layers == 24 + + +def test_detect_num_layers_returns_none_on_error(monkeypatch): + from meshnet_node.model_catalog import detect_num_layers + + # Monkeypatch AutoConfig to raise + import meshnet_node.model_catalog as cat + monkeypatch.setattr(cat, "detect_num_layers", lambda repo: None if "bad" in repo else detect_num_layers(repo)) + assert cat.detect_num_layers("bad/repo") is None + + +def test_startup_auto_detects_shard_range(monkeypatch, tmp_path): + """When shard_start/end are None, startup reads layer count from catalog.""" + from meshnet_node import startup as su + from meshnet_node.model_catalog import detect_num_layers + + calls = [] + + def fake_detect(repo): + calls.append(repo) + return 24 # Qwen2.5-0.5B + + monkeypatch.setattr(su, "_detect_num_layers", fake_detect) + + # Fake hardware detection + monkeypatch.setattr(su, "detect_hardware", lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0}) + + # Fake wallet + monkeypatch.setattr(su, "load_or_create_wallet", lambda **kw: (None, None, "fake-wallet")) + + # Fake TorchNodeServer + class FakeNode: + chat_completion_count = 0 + def start(self): return 9999 + def stop(self): pass + + import meshnet_node.startup as su2 + monkeypatch.setattr(su2, "TorchNodeServer", lambda **kw: FakeNode()) + + node = su.run_startup( + tracker_url="http://localhost:8080", + model_id="Qwen/Qwen2.5-0.5B-Instruct", + # shard_start and shard_end intentionally omitted + quantization="bfloat16", + host="127.0.0.1", + ) + assert calls == ["Qwen/Qwen2.5-0.5B-Instruct"] + assert isinstance(node, FakeNode) + + def test_legacy_start_subcommand_accepted(monkeypatch): """meshnet-node start --tracker http://... does not crash on arg parsing.""" from meshnet_node.cli import main