dash QOL

tests/test_real_model_backend.py

@@ -11,8 +11,12 @@ import pytest
 
 from meshnet_node.model_backend import (
     InsufficientVRAMError,
+    PartialModelLoadUnsupported,
     TensorPayload,
+    TorchModelShard,
     _call_layer,
+    _load_partial_model_from_snapshot,
+    _should_partial_materialize_shard,
     _decoder_attention_mask,
     _int_tensor_header,
     build_quantization_config,
@@ -334,6 +338,295 @@ def test_call_layer_passes_rotary_position_embeddings():
     ) == "hidden"
 
 
+def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapshot(tmp_path):
+    snapshot_dir = tmp_path / "snapshot"
+    snapshot_dir.mkdir()
+    (snapshot_dir / "config.json").write_text("{}")
+    (snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}')
+
+    assert _should_partial_materialize_shard(
+        str(snapshot_dir),
+        4,
+        7,
+        total_layers_hint=40,
+        uses_quantized_weights=False,
+    ) is True
+    assert _should_partial_materialize_shard(
+        str(snapshot_dir),
+        0,
+        39,
+        total_layers_hint=40,
+        uses_quantized_weights=False,
+    ) is False
+    assert _should_partial_materialize_shard(
+        str(snapshot_dir),
+        4,
+        7,
+        total_layers_hint=40,
+        uses_quantized_weights=True,
+    ) is False
+    assert _should_partial_materialize_shard(
+        "repo/model",
+        4,
+        7,
+        total_layers_hint=40,
+        uses_quantized_weights=False,
+    ) is False
+
+
+def test_partial_snapshot_loader_materializes_only_assigned_tensors(tmp_path):
+    snapshot_dir = tmp_path / "snapshot"
+    snapshot_dir.mkdir()
+    (snapshot_dir / "config.json").write_text("{}")
+    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
+        "weight_map": {
+            "model.embed_tokens.weight": "shard-1.safetensors",
+            "model.layers.0.self_attn.q_proj.weight": "shard-1.safetensors",
+            "model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors",
+            "model.layers.2.self_attn.q_proj.weight": "shard-3.safetensors",
+            "model.norm.weight": "shard-3.safetensors",
+            "lm_head.weight": "shard-3.safetensors",
+        }
+    }))
+    for rel in ("shard-1.safetensors", "shard-2.safetensors", "shard-3.safetensors"):
+        (snapshot_dir / rel).write_bytes(b"stub")
+
+    class FakeModule:
+        def __init__(self, name):
+            self.name = name
+            self.to_calls = []
+
+        def to(self, device):
+            self.to_calls.append(device)
+            return self
+
+    class FakeModel:
+        def __init__(self):
+            self.model = types.SimpleNamespace(
+                embed_tokens=FakeModule("embed"),
+                layers=[FakeModule("layer0"), FakeModule("layer1"), FakeModule("layer2")],
+                rotary_emb=FakeModule("rotary"),
+                norm=FakeModule("norm"),
+            )
+            self.lm_head = FakeModule("lm_head")
+            self.tie_weights_called = 0
+
+        def tie_weights(self):
+            self.tie_weights_called += 1
+
+    class AutoConfigStub:
+        @staticmethod
+        def from_pretrained(model_id):
+            assert model_id == str(snapshot_dir)
+            return types.SimpleNamespace(num_hidden_layers=3)
+
+    class AutoModelStub:
+        @staticmethod
+        def from_config(cfg, torch_dtype=None):
+            assert cfg.num_hidden_layers == 3
+            assert torch_dtype == "bf16"
+            return FakeModel()
+
+    class EmptyWeights:
+        def __init__(self):
+            self.entered = 0
+            self.exited = 0
+
+        def __call__(self):
+            return self
+
+        def __enter__(self):
+            self.entered += 1
+            return None
+
+        def __exit__(self, exc_type, exc, tb):
+            self.exited += 1
+            return False
+
+    init_empty_weights = EmptyWeights()
+    set_calls = []
+
+    def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
+        set_calls.append((tensor_name, device, value, dtype))
+
+    tensors = {
+        "shard-1.safetensors": {
+            "model.embed_tokens.weight": "embed",
+            "model.layers.0.self_attn.q_proj.weight": "layer0",
+        },
+        "shard-2.safetensors": {
+            "model.layers.1.self_attn.q_proj.weight": "layer1",
+        },
+        "shard-3.safetensors": {
+            "model.layers.2.self_attn.q_proj.weight": "layer2",
+            "model.norm.weight": "norm",
+            "lm_head.weight": "lm_head",
+        },
+    }
+
+    class FakeSafeOpen:
+        def __init__(self, filename, framework, device):
+            assert framework == "pt"
+            assert device == "cpu"
+            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 tensors[self.filename][tensor_name]
+
+    model = _load_partial_model_from_snapshot(
+        AutoConfigStub,
+        AutoModelStub,
+        types.SimpleNamespace(),
+        str(snapshot_dir),
+        1,
+        1,
+        "bf16",
+        "cpu:0",
+        init_empty_weights_fn=init_empty_weights,
+        set_tensor_fn=fake_set_tensor,
+        safe_open_fn=FakeSafeOpen,
+    )
+
+    assert init_empty_weights.entered == 1
+    assert init_empty_weights.exited == 1
+    assert model.tie_weights_called == 1
+    assert [call[0] for call in set_calls] == ["model.layers.1.self_attn.q_proj.weight"]
+    assert model.model.layers[1].to_calls == ["cpu:0"]
+    assert model.model.layers[0].to_calls == []
+    assert model.model.layers[2].to_calls == []
+    assert model.model.embed_tokens.to_calls == []
+    assert model.model.norm.to_calls == []
+    assert model.lm_head.to_calls == []
+    assert model.model.rotary_emb.to_calls == ["cpu:0"]
+
+
+def test_partial_snapshot_loader_requires_known_layer_count(tmp_path):
+    snapshot_dir = tmp_path / "snapshot"
+    snapshot_dir.mkdir()
+    (snapshot_dir / "config.json").write_text("{}")
+    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
+        "weight_map": {"model.layers.0.self_attn.q_proj.weight": "shard.safetensors"}
+    }))
+    (snapshot_dir / "shard.safetensors").write_bytes(b"stub")
+
+    class AutoConfigStub:
+        @staticmethod
+        def from_pretrained(model_id):
+            return types.SimpleNamespace()
+
+    class AutoModelStub:
+        @staticmethod
+        def from_config(cfg, torch_dtype=None):
+            raise AssertionError("from_config should not run without a known layer count")
+
+    class UnusedContext:
+        def __enter__(self):
+            return None
+
+        def __exit__(self, exc_type, exc, tb):
+            return False
+
+    with pytest.raises(PartialModelLoadUnsupported, match="num_hidden_layers"):
+        _load_partial_model_from_snapshot(
+            AutoConfigStub,
+            AutoModelStub,
+            types.SimpleNamespace(),
+            str(snapshot_dir),
+            0,
+            0,
+            "bf16",
+            "cpu:0",
+            init_empty_weights_fn=lambda: UnusedContext(),
+            set_tensor_fn=lambda *args, **kwargs: None,
+            safe_open_fn=lambda *args, **kwargs: None,
+        )
+
+
+def test_torch_model_shard_prefers_partial_loader_for_local_snapshot(tmp_path, monkeypatch):
+    import meshnet_node.model_backend as backend
+
+    snapshot_dir = tmp_path / "snapshot"
+    snapshot_dir.mkdir()
+    (snapshot_dir / "config.json").write_text("{}")
+    (snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}')
+
+    class FakeModel:
+        def __init__(self):
+            self.model = types.SimpleNamespace(
+                layers=[object(), object(), object()],
+                embed_tokens=object(),
+            )
+            self.config = types.SimpleNamespace(hidden_size=8)
+            self.eval_called = 0
+
+        def eval(self):
+            self.eval_called += 1
+
+    fake_model = FakeModel()
+    partial_calls = []
+
+    class AutoConfigStub:
+        @staticmethod
+        def from_pretrained(model_id, cache_dir=None):
+            return types.SimpleNamespace(num_hidden_layers=3, text_config=types.SimpleNamespace(dtype="torch.bfloat16"))
+
+    class AutoModelStub:
+        @staticmethod
+        def from_pretrained(*args, **kwargs):
+            raise AssertionError("full model load should not run for partial local shards")
+
+    class AutoTokenizerStub:
+        @staticmethod
+        def from_pretrained(model_id, cache_dir=None):
+            assert model_id == str(snapshot_dir)
+            return types.SimpleNamespace()
+
+    monkeypatch.setitem(
+        sys.modules,
+        "torch",
+        types.SimpleNamespace(
+            cuda=types.SimpleNamespace(is_available=lambda: False),
+            device=lambda value: value,
+            bfloat16="bf16",
+        ),
+    )
+    monkeypatch.setitem(
+        sys.modules,
+        "transformers",
+        types.SimpleNamespace(
+            AutoConfig=AutoConfigStub,
+            AutoModelForCausalLM=AutoModelStub,
+            AutoTokenizer=AutoTokenizerStub,
+        ),
+    )
+    monkeypatch.setattr(
+        backend,
+        "_load_partial_model_from_snapshot",
+        lambda *args, **kwargs: partial_calls.append((args, kwargs)) or fake_model,
+    )
+
+    shard = TorchModelShard(
+        "repo/model",
+        1,
+        1,
+        quantization="auto",
+        cache_dir=snapshot_dir,
+    )
+
+    assert len(partial_calls) == 1
+    assert shard.model is fake_model
+    assert fake_model.eval_called == 1
+    assert shard.total_layers == 3
+    assert shard.is_head is False
+    assert shard.is_tail is False
+
+
 @pytest.mark.integration
 def test_two_node_gpt2_completion_is_deterministic():
     if os.environ.get("CI"):