Skip multimodal/MTP checkpoint tensors absent from the text-only causal LM

Qwen3.5/3.6-MoE checkpoints ship vision (model.visual.*) and multi-token-
prediction (mtp.*) weights; the partial shard loader assigned them into the
text-only Qwen3_5MoeForCausalLM and crashed with AttributeError 'mtp'.
Filter selected tensors against the built model's state_dict keys, matching
transformers' _keys_to_ignore_on_load_unexpected behavior.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

packages/node/meshnet_node/model_backend.py

@@ -482,12 +482,36 @@ def _load_partial_model_from_snapshot(
     if callable(tie_weights):
         tie_weights()
 
+    # Multimodal/MTP checkpoints (e.g. Qwen3.5/3.6-MoE) carry vision and
+    # multi-token-prediction tensors the text-only CausalLM never builds;
+    # transformers' from_pretrained drops them via _keys_to_ignore_on_load_unexpected,
+    # so the manual loader must skip them too.
+    expected_keys = _model_state_dict_keys(model)
     tensors_by_file: dict[str, list[str]] = {}
+    skipped: list[str] = []
     for tensor_name in sorted(tensor_names):
         rel_file = weight_map.get(tensor_name)
         if not isinstance(rel_file, str):
             continue
+        if (
+            expected_keys is not None
+            and _checkpoint_tensor_name_for_model(model, tensor_name) not in expected_keys
+        ):
+            skipped.append(tensor_name)
+            continue
         tensors_by_file.setdefault(rel_file, []).append(tensor_name)
+    if skipped:
+        preview = ", ".join(skipped[:3])
+        print(
+            f"  Skipping {len(skipped)} checkpoint tensors absent from the causal LM "
+            f"(e.g. {preview})",
+            flush=True,
+        )
+    if not tensors_by_file:
+        raise PartialModelLoadUnsupported(
+            f"no checkpoint tensors for layers {shard_start}-{shard_end} match the "
+            f"causal LM built from {snapshot_dir}"
+        )
 
     for rel_file, names in tensors_by_file.items():
         checkpoint_file = snapshot_dir / rel_file
@@ -584,6 +608,17 @@ def _causal_lm_config(cfg: Any) -> Any:
     return cfg
 
 
+def _model_state_dict_keys(model: Any) -> set[str] | None:
+    """Expected parameter/buffer names, or None when the model can't report them."""
+    state_dict = getattr(model, "state_dict", None)
+    if not callable(state_dict):
+        return None
+    try:
+        return set(state_dict().keys())
+    except Exception:
+        return None
+
+
 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)

tests/test_real_model_backend.py

@@ -559,6 +559,96 @@ def test_partial_snapshot_loader_remaps_language_model_checkpoint_keys(tmp_path)
     assert set_calls == ["model.layers.1.self_attn.q_proj.weight"]
 
 
+def test_partial_snapshot_loader_skips_tensors_absent_from_causal_lm(tmp_path):
+    # Multimodal/MTP checkpoints (Qwen3.5/3.6-MoE) carry mtp.* and model.visual.*
+    # tensors that the text-only CausalLM never builds — they must be skipped,
+    # not assigned (assignment raises AttributeError: 'mtp' / 'visual').
+    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",
+            "mtp.layers.1.input_layernorm.weight": "shard-2.safetensors",
+            "model.visual.blocks.1.attn.qkv.weight": "shard-2.safetensors",
+        }
+    }))
+    (snapshot_dir / "shard-2.safetensors").write_bytes(b"stub")
+
+    class FakeModule:
+        def to(self, device):
+            return self
+
+    class FakeModel:
+        def __init__(self):
+            self.model = types.SimpleNamespace(
+                layers=[FakeModule(), FakeModule(), FakeModule()],
+                rotary_emb=FakeModule(),
+            )
+
+        def tie_weights(self):
+            pass
+
+        def state_dict(self):
+            return {"model.layers.1.self_attn.q_proj.weight": None}
+
+    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):
+            pass
+
+        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()