diff --git a/packages/node/meshnet_node/model_backend.py b/packages/node/meshnet_node/model_backend.py index 00d04f6..ef189fa 100644 --- a/packages/node/meshnet_node/model_backend.py +++ b/packages/node/meshnet_node/model_backend.py @@ -134,8 +134,9 @@ class TorchModelShard: self.model.to(self.device) except Exception as exc: if _looks_like_oom(exc): + memory_kind = "VRAM" if self.device.type == "cuda" else "RAM" raise InsufficientVRAMError( - f"insufficient VRAM to load {model_id} layers {shard_start}:{shard_end} " + f"insufficient {memory_kind} to load {model_id} layers {shard_start}:{shard_end} " f"with {quantization} quantization; choose a smaller shard or lower quantization" ) from exc raise @@ -411,7 +412,7 @@ def _should_partial_materialize_shard( return False if total_layers_hint is None: return False - return not (shard_start == 0 and shard_end >= total_layers_hint - 1) + return True def _load_partial_model_from_snapshot( @@ -476,7 +477,7 @@ def _load_partial_model_from_snapshot( ) with init_empty_weights_fn(): - model = auto_model_for_causal_lm.from_config(cfg, torch_dtype=dtype) + model = auto_model_for_causal_lm.from_config(_causal_lm_config(cfg), torch_dtype=dtype) tie_weights = getattr(model, "tie_weights", None) if callable(tie_weights): tie_weights() @@ -498,7 +499,7 @@ def _load_partial_model_from_snapshot( for tensor_name in names: set_tensor_fn( model, - tensor_name, + _checkpoint_tensor_name_for_model(model, tensor_name), device, value=handle.get_tensor(tensor_name), dtype=dtype, @@ -569,38 +570,74 @@ def _native_torch_dtype(cfg: Any, torch: Any) -> Any: return torch.bfloat16 +def _causal_lm_config(cfg: Any) -> Any: + """Use the text decoder config for composite VLM/MoE presets.""" + get_text_config = getattr(cfg, "get_text_config", None) + if callable(get_text_config): + try: + return get_text_config() + except Exception: + pass + text_config = getattr(cfg, "text_config", None) + if text_config is not None: + return text_config + return cfg + + +def _checkpoint_tensor_name_for_model(model: Any, tensor_name: str) -> str: + """Map multimodal checkpoint keys onto text-only CausalLM modules when needed.""" + inner = getattr(model, "model", None) + if inner is not None and hasattr(inner, "language_model"): + return tensor_name + if ".language_model." in tensor_name: + return tensor_name.replace(".language_model.", ".") + return tensor_name + + +def _transformer_backbone(model: Any) -> Any: + if hasattr(model, "model"): + inner = model.model + language_model = getattr(inner, "language_model", None) + if language_model is not None: + return language_model + return inner + if hasattr(model, "transformer"): + return model.transformer + raise ModelBackendError( + "unsupported HuggingFace model architecture: no transformer backbone found" + ) + + def _model_layers(model: Any) -> Any: - if hasattr(model, "model") and hasattr(model.model, "layers"): - return model.model.layers - if hasattr(model, "transformer") and hasattr(model.transformer, "h"): - return model.transformer.h + backbone = _transformer_backbone(model) + for attr in ("layers", "h", "blocks"): + layers = getattr(backbone, attr, None) + if layers is not None: + return layers raise ModelBackendError( "unsupported HuggingFace model architecture: no transformer layers found" ) def _embed_tokens(model: Any) -> Any: - if hasattr(model, "model") and hasattr(model.model, "embed_tokens"): - return model.model.embed_tokens - if hasattr(model, "transformer") and hasattr(model.transformer, "wte"): - return model.transformer.wte + backbone = _transformer_backbone(model) + for attr in ("embed_tokens", "wte"): + embed = getattr(backbone, attr, None) + if embed is not None: + return embed raise ModelBackendError( "unsupported HuggingFace model architecture: no token embeddings found" ) def _position_embeddings(model: Any) -> Any | None: - if hasattr(model, "transformer") and hasattr(model.transformer, "wpe"): - return model.transformer.wpe - return None + backbone = _transformer_backbone(model) + return getattr(backbone, "wpe", None) def _rotary_embedding_module(model: Any) -> Any | None: - if hasattr(model, "model") and hasattr(model.model, "rotary_emb"): - return model.model.rotary_emb - if hasattr(model, "transformer") and hasattr(model.transformer, "rotary_emb"): - return model.transformer.rotary_emb - return None + backbone = _transformer_backbone(model) + return getattr(backbone, "rotary_emb", None) def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> list[Any]: @@ -627,10 +664,11 @@ def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> l def _final_norm(model: Any) -> Any | None: - if hasattr(model, "model") and hasattr(model.model, "norm"): - return model.model.norm - if hasattr(model, "transformer") and hasattr(model.transformer, "ln_f"): - return model.transformer.ln_f + backbone = _transformer_backbone(model) + for attr in ("norm", "ln_f", "final_layer_norm"): + norm = getattr(backbone, attr, None) + if norm is not None: + return norm return None @@ -743,7 +781,12 @@ def _looks_like_oom(exc: BaseException) -> bool: current: BaseException | None = exc while current is not None: text = str(current).lower() - if "out of memory" in text or "cuda error: out of memory" in text: + if ( + "out of memory" in text + or "cuda error: out of memory" in text + or "paging file is too small" in text + or "os error 1455" in text + ): return True current = current.__cause__ or current.__context__ return False diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index f5a07ee..54fa139 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -995,6 +995,20 @@ def run_startup( ) if user_pinned_shard: shard_label = f"{shard_label} (pinned)" + if user_pinned_shard and assigned_total_layers and assignment_bytes_per_layer: + pinned_layers = shard_end - shard_start + 1 + max_layers = _max_assignable_layers( + memory_budget_mb, + assigned_total_layers, + assignment_bytes_per_layer, + ) + if pinned_layers > max_layers: + raise ValueError( + f"Pinned shard layers {shard_start}–{shard_end} ({pinned_layers} layers) exceed " + f"the {memory_budget_mb / 1024:.1f} GB {memory_budget_source} budget " + f"(fits up to {max_layers}/{assigned_total_layers} layers at bfloat16). " + "Drop --shard-start/--shard-end to let the tracker auto-assign, or pin a smaller range." + ) print(f" Shard: {shard_label}", flush=True) # 4. Download shard @@ -1020,7 +1034,77 @@ def run_startup( ) print(f" Cached at: {shard_path}", flush=True) - # 5. Start HTTP server + # 5. Start HTTP server — real HF weights use TorchNodeServer; stub-model stays stub. + _node_start_time = time.monotonic() + if hf_repo and assigned_model != "stub-model": + print("Loading real PyTorch model shard...", flush=True) + node = TorchNodeServer( + host=host, + port=port, + model_id=hf_repo, + shard_start=shard_start, + shard_end=shard_end, + quantization=quantization, + tracker_url=tracker_url, + route_timeout=route_timeout, + cache_dir=shard_path, + debug=debug, + max_loaded_shards=max_loaded_shards, + ) + actual_port = node.start() + total_layers = getattr(getattr(node, "backend", None), "total_layers", None) or assigned_total_layers + shard_label = _format_shard_label(shard_start, shard_end, total_layers, model_name=assigned_model) + if user_pinned_shard: + shard_label = f"{shard_label} (pinned)" + public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host) + endpoint = f"http://{public_host}:{actual_port}" + local_base_url = f"http://127.0.0.1:{actual_port}" + relay_bridge, relay_fields = _start_relay_bridge_if_available( + tracker_url, + address, + local_base_url, + endpoint, + relay_url=relay_url, + ) + _attach_relay_bridge(node, relay_bridge) + reg_payload = { + "endpoint": endpoint, + "model": assigned_model, + "hf_repo": hf_repo, + "num_layers": total_layers, + "shard_start": shard_start, + "shard_end": shard_end, + "downloaded_models": downloaded_models, + "hardware_profile": hw, + "wallet_address": address, + "quantization": quantization, + "score": 1.0, + "tracker_mode": (shard_start == 0), + "managed_assignment": not user_pinned_shard, + "model_metadata": model_metadata_for(hf_repo, total_layers, cache_dir=shard_path), + **registration_capabilities, + **relay_fields, + } + tracker_node_id = _register_with_tracker( + tracker_url, reg_payload, node, _node_start_time, + ) + print( + f"\n{'=' * 32}\n" + f"meshnet-node ready\n" + f" Wallet: {address}\n" + f" Model ID: {hf_repo}\n" + f" Shard: {shard_label}\n" + f" {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n" + f" Quantization: {quantization}\n" + f" Endpoint: {endpoint}\n" + f" Node ID: {tracker_node_id or 'unregistered'}\n" + f" Hardware: {_hardware_label(device, gpu_name)}\n" + f" Benchmark: {bench_tps:,.0f} (throughput index)\n" + f"{'=' * 32}", + flush=True, + ) + return node + is_last = shard_end >= assignment.get("model_layers_end", shard_end) node = StubNodeServer( host=host, @@ -1031,7 +1115,6 @@ def run_startup( model=assigned_model, shard_path=shard_path, ) - _node_start_time = time.monotonic() actual_port = node.start() public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host) endpoint = f"http://{public_host}:{actual_port}" diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 1ea498d..d28ed8a 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -4864,6 +4864,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): "model": resolved_name, "model_layers_end": required_end, "peers": peers, + "bytes_per_layer": _preset_bytes_per_layer(preset), "model_sources": self._model_sources( resolved_name, preset, diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index 3a22a51..fdf09a9 100644 --- a/tests/test_node_startup.py +++ b/tests/test_node_startup.py @@ -1646,6 +1646,106 @@ def test_preset_model_startup_honors_pinned_shard_range(tmp_path, monkeypatch): tracker.stop() +def test_preset_startup_rejects_pinned_shard_above_memory_budget(tmp_path, monkeypatch): + """Pinned layer ranges that exceed the node memory budget fail before model load.""" + import meshnet_node.startup as startup_mod + + monkeypatch.setattr( + startup_mod, + "detect_hardware", + lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 8 * 1024}, + ) + + tracker = TrackerServer(model_presets={ + "big-model": { + "layers_start": 0, + "layers_end": 39, + "hf_repo": "org/big-model", + "bytes_per_layer": {"bfloat16": 2 * 1024 * 1024 * 1024}, + }, + }) + tracker_port = tracker.start() + tracker_url = f"http://127.0.0.1:{tracker_port}" + try: + with pytest.raises(ValueError, match="Pinned shard layers 0–39"): + run_startup( + tracker_url=tracker_url, + model="big-model", + shard_start=0, + shard_end=39, + wallet_path=tmp_path / "wallet.json", + cache_dir=tmp_path / "shards", + ) + finally: + tracker.stop() + + +def test_preset_model_with_hf_repo_loads_torch_backend(tmp_path, monkeypatch, capsys): + """Named presets that advertise hf_repo must load TorchNodeServer, not the stub server.""" + import meshnet_node.startup as startup_mod + + class FakeBackend: + total_layers = 16 + + torch_calls: list[dict] = [] + + class FakeTorchNodeServer: + def __init__(self, **kwargs): + torch_calls.append(kwargs) + self.backend = FakeBackend() + self.port = None + self.chat_completion_count = 0 + self.tracker_node_id = None + + def start(self): + self.port = 7002 + return self.port + + def stop(self): + pass + + monkeypatch.setattr( + startup_mod, + "detect_hardware", + lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16 * 1024}, + ) + monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer) + monkeypatch.setattr(startup_mod, "StubNodeServer", lambda **_kw: (_ for _ in ()).throw(AssertionError("preset with hf_repo must not use StubNodeServer"))) + + model_dir = tmp_path / "node-shards" / "tiny-llama" + model_dir.mkdir(parents=True) + (model_dir / "config.json").write_text('{"num_hidden_layers": 16}') + monkeypatch.setattr(startup_mod, "download_shard", lambda *_a, **_kw: model_dir) + + tracker = TrackerServer(model_presets={ + "tiny-llama": {"layers_start": 0, "layers_end": 15, "hf_repo": "org/tiny-llama-shards"} + }) + tracker_port = tracker.start() + tracker_url = f"http://127.0.0.1:{tracker_port}" + try: + node = run_startup( + tracker_url=tracker_url, + model="tiny-llama", + wallet_path=tmp_path / "wallet.json", + cache_dir=tmp_path / "node-shards", + ) + try: + assert len(torch_calls) == 1 + assert torch_calls[0]["model_id"] == "org/tiny-llama-shards" + assert torch_calls[0]["cache_dir"] == model_dir + output = capsys.readouterr().out + assert "Loading real PyTorch model shard..." in output + assert "Model ID: org/tiny-llama-shards" in output + network_map = _get_json(f"{tracker_url}/v1/network/map") + registered = network_map["nodes"][0] + assert registered["hf_repo"] == "org/tiny-llama-shards" + assert registered["num_layers"] == 16 + finally: + node.stop() + finally: + tracker.stop() + + def test_torch_startup_retries_registration_when_tracker_unreachable( tmp_path, monkeypatch, diff --git a/tests/test_real_model_backend.py b/tests/test_real_model_backend.py index 656c55f..de8e5e5 100644 --- a/tests/test_real_model_backend.py +++ b/tests/test_real_model_backend.py @@ -17,6 +17,7 @@ from meshnet_node.model_backend import ( TensorPayload, TorchModelShard, _call_layer, + _checkpoint_tensor_name_for_model, _load_partial_model_from_snapshot, _should_partial_materialize_shard, _decoder_attention_mask, @@ -429,7 +430,7 @@ def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapsho 39, total_layers_hint=40, uses_quantized_weights=False, - ) is False + ) is True assert _should_partial_materialize_shard( str(snapshot_dir), 4, @@ -446,6 +447,118 @@ def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapsho ) is False +def test_checkpoint_tensor_name_remapped_for_text_only_causal_lm(): + class TextOnlyModel: + def __init__(self): + self.model = types.SimpleNamespace(layers=[]) + + model = TextOnlyModel() + assert _checkpoint_tensor_name_for_model( + model, + "model.language_model.layers.0.mlp.gate.weight", + ) == "model.layers.0.mlp.gate.weight" + assert _checkpoint_tensor_name_for_model( + model, + "model.language_model.embed_tokens.weight", + ) == "model.embed_tokens.weight" + + +def test_checkpoint_tensor_name_kept_for_multimodal_backbone(): + class MultimodalModel: + def __init__(self): + self.model = types.SimpleNamespace(language_model=types.SimpleNamespace()) + + model = MultimodalModel() + name = "model.language_model.layers.0.mlp.gate.weight" + assert _checkpoint_tensor_name_for_model(model, name) == name + + +def test_partial_snapshot_loader_remaps_language_model_checkpoint_keys(tmp_path): + snapshot_dir = tmp_path / "snapshot" + snapshot_dir.mkdir() + (snapshot_dir / "config.json").write_text(json.dumps({ + "text_config": {"num_hidden_layers": 3}, + })) + (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({ + "weight_map": { + "model.language_model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors", + } + })) + (snapshot_dir / "shard-2.safetensors").write_bytes(b"stub") + + class FakeModule: + def __init__(self): + self.to_calls = [] + + def to(self, device): + self.to_calls.append(device) + return self + + class FakeModel: + def __init__(self): + self.model = types.SimpleNamespace( + layers=[FakeModule(), FakeModule(), FakeModule()], + rotary_emb=FakeModule(), + ) + + def tie_weights(self): + pass + + class AutoConfigStub: + @staticmethod + def from_pretrained(model_id): + return types.SimpleNamespace( + text_config=types.SimpleNamespace(num_hidden_layers=3), + get_text_config=lambda: types.SimpleNamespace(num_hidden_layers=3), + ) + + class AutoModelStub: + @staticmethod + def from_config(cfg, torch_dtype=None): + return FakeModel() + + set_calls = [] + + def fake_set_tensor(module, tensor_name, device, value=None, dtype=None): + set_calls.append(tensor_name) + + class FakeSafeOpen: + def __init__(self, filename, framework, device): + self.filename = Path(filename).name + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def get_tensor(self, tensor_name): + return tensor_name + + class UnusedContext: + def __enter__(self): + return None + + def __exit__(self, exc_type, exc, tb): + return False + + _load_partial_model_from_snapshot( + AutoConfigStub, + AutoModelStub, + types.SimpleNamespace(), + str(snapshot_dir), + 1, + 1, + "bf16", + "cpu:0", + init_empty_weights_fn=UnusedContext, + set_tensor_fn=fake_set_tensor, + safe_open_fn=FakeSafeOpen, + ) + + assert set_calls == ["model.layers.1.self_attn.q_proj.weight"] + + def test_partial_snapshot_loader_materializes_only_assigned_tensors(tmp_path): snapshot_dir = tmp_path / "snapshot" snapshot_dir.mkdir()