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

feat(us-016): auto-detect shard range from model config

Browse Source 
Layer count is now fetched from the curated catalog (zero network calls
for known models) or via AutoConfig.from_pretrained() (~1 KB config.json
only) when model_id is given without --shard-start/--shard-end.

- model_catalog: add detect_num_layers(), two small Qwen models at top
- startup: _detect_num_layers() helper; shard range auto-derived
- wizard: show detected layer count for custom HF repos
- tests: 3 new tests for auto-shard; fix catalog-order assumptions

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Dobromir Popov
2026-06-29 18:27:50 +03:00
parent a2258d3df4
commit 080d49b2c2
4 changed files with 135 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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