fix(vl): strip model.visual.* in qwen2_vl/qwen3_vl/qwen3_vl_moe sanitize

HuggingFace ForConditionalGeneration VL checkpoints (Qwen3_5MoeForConditionalGeneration,
Qwen3VLForConditionalGeneration, Qwen2VLForConditionalGeneration) nest the
language model and vision tower under a top-level 'model' dict:
  - model.language_model.* (language model weights)
  - model.visual.*        (vision tower weights)

The existing sanitize() helpers only popped bare 'visual'/'vision_tower'
top-level keys, which matches the mlx-vlm conversion layout but not the HF
layout. On HF checkpoints:
  - qwen3_vl_moe.sanitize crashed with KeyError('model') because
    weights['language_model']['model'] doesn't exist (the language model
    is nested under model.language_model.model.*).
  - qwen3_vl and qwen2_vl leaked model.visual.* into the sanitized weights,
    re-prefixed as language_model.model.visual.* — causing 'parameters not
    in model' errors downstream.

Fix: after tree_unflatten, pop the top-level 'model' dict, re-root its
'language_model' to the top level, and drop its 'visual'/'vision_tower'
subtrees.

Verified on deepreinforce-ai/Ornith-1.0-35B (333 model.visual.* tensors,
full vision_config): the grafted checkpoint now loads without the
'768 parameters not in model' error that blocked it before.

Tests: TestVLSanitize covers all three model classes against the HF
ForConditionalGeneration layout — each must (a) not crash, (b) leave zero
visual/vision_tower keys, (c) preserve the language model tensors.

Author: jonathanharder

mlx_lm/models/qwen2_vl.py

@@ -44,6 +44,12 @@ def sanitize(self, weights):
         weights = tree_unflatten(list(weights.items()))
         weights.pop("visual", None)
         weights.pop("vision_tower", None)
+        # HF ForConditionalGeneration nests the language model under
+        # ``model.language_model.*`` and vision under ``model.visual.*``.
+        nested = weights.pop("model", None)
+        if isinstance(nested, dict):
+            if isinstance(nested.get("language_model"), dict):
+                weights["language_model"] = nested["language_model"]
         weights = dict(tree_flatten(weights))
 
         sanitized = {}

mlx_lm/models/qwen3_vl.py

@@ -43,6 +43,12 @@ def __call__(
     def sanitize(self, weights):
         weights = tree_unflatten(list(weights.items()))
         weights.pop("vision_tower", None)
+        # HF ForConditionalGeneration nests the language model under
+        # ``model.language_model.*`` and vision under ``model.visual.*``.
+        nested = weights.pop("model", None)
+        if isinstance(nested, dict):
+            if isinstance(nested.get("language_model"), dict):
+                weights["language_model"] = nested["language_model"]
         weights = dict(tree_flatten(weights))
 
         sanitized = {}

mlx_lm/models/qwen3_vl_moe.py

@@ -39,12 +39,20 @@ def __call__(
     def sanitize(self, weights):
         weights = tree_unflatten(list(weights.items()))
         weights.pop("visual", None)
+        # HF ForConditionalGeneration nests the language model under
+        # ``model.language_model.*`` and vision under ``model.visual.*``.
+        nested = weights.pop("model", None)
+        if isinstance(nested, dict):
+            if isinstance(nested.get("language_model"), dict):
+                weights["language_model"] = nested["language_model"]
+
+        language_model = weights["language_model"]
         weights = dict(
             tree_flatten(
                 {
                     "language_model": {
-                        "model": weights["language_model"]["model"],
-                        "lm_head": weights["language_model"]["lm_head"],
+                        "model": language_model["model"],
+                        "lm_head": language_model["lm_head"],
                     }
                 }
             )

tests/test_models.py

@@ -3221,5 +3221,142 @@ def test_gated_delta_masked(self):
                 self.assertTrue(mx.allclose(st, st_gt, rtol=1e-4, atol=1e-3))
 
 
+class TestVLSanitize(unittest.TestCase):
+    """Vision-language model sanitize() must strip the vision tower for the
+    HuggingFace ForConditionalGeneration checkpoint layout, where vision
+    weights ship under ``model.visual.*`` (or ``model.vision_tower.*``) and the
+    language model under ``model.language_model.*`` — not the bare
+    ``visual`` / ``vision_tower`` / ``language_model.*`` prefixes the sanitize
+    helpers were written against.
+    """
+
+    # HF Qwen3_5MoeForConditionalGeneration layout (Ornith-1.0, etc.).
+    HF_VL_WEIGHTS = {
+        "model.visual.pos_embed.weight": None,
+        "model.visual.blocks.0.attn.qkv.weight": None,
+        "model.visual.merger.norm.weight": None,
+        "model.language_model.model.embed_tokens.weight": None,
+        "model.language_model.model.layers.0.self_attn.q_proj.weight": None,
+        "model.language_model.model.layers.0.mlp.experts.gate_up_proj.weight": None,
+        "model.language_model.model.layers.0.mlp.experts.down_proj.weight": None,
+        "model.language_model.lm_head.weight": None,
+        "language_model.lm_head.weight": None,
+    }
+
+    def _qwen3_moe_text_config(self):
+        return {
+            "model_type": "qwen3_moe",
+            "num_hidden_layers": 1,
+            "num_experts": 4,
+            "num_experts_per_tok": 2,
+            "hidden_size": 32,
+            "intermediate_size": 64,
+            "moe_intermediate_size": 16,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
+            "head_dim": 8,
+            "vocab_size": 100,
+            "rms_norm_eps": 1e-6,
+            "full_attention_interval": 4,
+            "linear_num_value_heads": 4,
+            "linear_num_key_heads": 2,
+            "linear_key_head_dim": 8,
+            "linear_value_head_dim": 8,
+            "linear_conv_kernel_dim": 4,
+            "tie_word_embeddings": False,
+            "decoder_sparse_step": 1,
+            "mlp_only_layers": [],
+            "rope_theta": 100000.0,
+            "max_position_embeddings": 4096,
+            "norm_topk_prob": True,
+        }
+
+    def _qwen2_text_config(self):
+        return {
+            "model_type": "qwen2",
+            "num_hidden_layers": 1,
+            "hidden_size": 32,
+            "intermediate_size": 64,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
+            "head_dim": 8,
+            "vocab_size": 100,
+            "rms_norm_eps": 1e-6,
+            "tie_word_embeddings": False,
+            "rope_theta": 100000.0,
+            "max_position_embeddings": 4096,
+        }
+
+    def _qwen3_text_config(self):
+        return {
+            "model_type": "qwen3",
+            "num_hidden_layers": 1,
+            "hidden_size": 32,
+            "intermediate_size": 64,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
+            "head_dim": 8,
+            "vocab_size": 100,
+            "rms_norm_eps": 1e-6,
+            "tie_word_embeddings": False,
+            "rope_theta": 100000.0,
+            "max_position_embeddings": 4096,
+        }
+
+    def test_qwen3_vl_moe_strips_model_visual_prefix(self):
+        from mlx_lm.models.qwen3_vl_moe import Model, ModelArgs
+
+        args = ModelArgs(
+            model_type="qwen3_vl_moe",
+            text_config=self._qwen3_moe_text_config(),
+        )
+        model = Model(args)
+        result = model.sanitize(dict(self.HF_VL_WEIGHTS))
+
+        visual = [k for k in result if "visual" in k or "vision_tower" in k]
+        self.assertEqual(
+            visual, [], f"vision keys leaked into sanitized weights: {visual}"
+        )
+        # The language model tensors must survive, re-rooted under language_model.*.
+        self.assertIn(
+            "language_model.model.layers.0.self_attn.q_proj.weight",
+            result,
+        )
+
+    def test_qwen3_vl_strips_model_visual_prefix(self):
+        from mlx_lm.models.qwen3_vl import Model, ModelArgs
+
+        args = ModelArgs(
+            model_type="qwen3_vl",
+            text_config=self._qwen3_text_config(),
+        )
+        model = Model(args)
+        result = model.sanitize(dict(self.HF_VL_WEIGHTS))
+
+        visual = [k for k in result if "visual" in k or "vision_tower" in k]
+        self.assertEqual(
+            visual, [], f"vision keys leaked into sanitized weights: {visual}"
+        )
+        self.assertIn(
+            "language_model.model.layers.0.self_attn.q_proj.weight",
+            result,
+        )
+
+    def test_qwen2_vl_strips_model_visual_prefix(self):
+        from mlx_lm.models.qwen2_vl import Model, ModelArgs
+
+        args = ModelArgs(
+            model_type="qwen2_vl",
+            text_config=self._qwen2_text_config(),
+        )
+        model = Model(args)
+        result = model.sanitize(dict(self.HF_VL_WEIGHTS))
+
+        visual = [k for k in result if "visual" in k or "vision_tower" in k]
+        self.assertEqual(
+            visual, [], f"vision keys leaked into sanitized weights: {visual}"
+        )
+
+
 if __name__ == "__main__":
     unittest.main()