try fix model loading quen3.6-35b

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