try fix model loading quen3.6-35b
This commit is contained in:
Dobromir Popov
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user