try fix model loading quen3.6-35b

This commit is contained in:
Dobromir Popov
2026-07-07 18:36:29 +02:00
parent f1eea5b6d4
commit cdd2699e63
5 changed files with 368 additions and 28 deletions

View File

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

View File

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

View File

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

View File

@@ -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 039"):
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,

View File

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