Add vlm_vision_weights helper and adopt it in VLM vision sub-models (#336)

## Summary

DRY refactor: dedupe the VLM vision-tower `preprocess_weights`
boilerplate.

The vision sub-models of **LLaVA**, **Gemma3**, and **Mllama**
hand-wrote the identical loop: filter to the vision-tower prefixes, then
rename the HuggingFace vision MLP projections `mlp.fc1` → `mlp.up_proj`
and `mlp.fc2` → `mlp.down_proj` to match our `FCMLP` component naming.

### Changes
- **Add `vlm_vision_weights(state_dict, prefixes)`** to
`_weight_utils.py`, alongside the existing `vlm_decoder_weights` /
`vlm_embedding_weights` helpers.
- Adopt it in gemma3, llava, mllama. The only per-model difference is
the prefix tuple: gemma3/llava use `("vision_tower.",
"multi_modal_projector.")`, mllama uses `("vision_model.",)`.
- Add 3 unit tests (filter+rename, single prefix, no-match).

Behaviour-preserving: same filtering and same renames, just centralised.

### Verification
- Fast suite: **2792 passed** (baseline 2789 + 3 new tests), 43 skipped.
- `ruff check` + `ruff format --check` clean.

Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: justinchuby

src/mobius/_weight_utils.py

@@ -521,6 +521,39 @@ def vlm_embedding_weights(
     return renamed
 
 
+def vlm_vision_weights(
+    state_dict: dict[str, torch.Tensor],
+    prefixes: tuple[str, ...],
+) -> dict[str, torch.Tensor]:
+    """Extract vision-tower weights for a VLM vision sub-model.
+
+    Keeps only keys starting with one of *prefixes* and renames the vision
+    MLP projections ``mlp.fc1`` → ``mlp.up_proj`` and ``mlp.fc2`` →
+    ``mlp.down_proj`` to match our ``FCMLP`` component naming.
+
+    This is the standard pattern for VLM vision sub-models (LLaVA, Gemma3,
+    Mllama) whose HuggingFace vision encoders use ``fc1``/``fc2`` MLP names.
+
+    Args:
+        state_dict: Full model state dict.
+        prefixes: Prefixes identifying vision-tower weights to keep, e.g.
+            ``("vision_tower.", "multi_modal_projector.")`` or
+            ``("vision_model.",)``.
+
+    Returns:
+        New dictionary with the kept vision weights and renamed MLP keys.
+    """
+    renamed: dict[str, torch.Tensor] = {}
+    for key, value in state_dict.items():
+        if not key.startswith(prefixes):
+            continue
+        new_key = key.replace(".mlp.fc1.", ".mlp.up_proj.").replace(
+            ".mlp.fc2.", ".mlp.down_proj."
+        )
+        renamed[new_key] = value
+    return renamed
+
+
 def _reshape_packed_qweight(value: torch.Tensor, blob_size: int) -> torch.Tensor:
     """Transpose and reshape a packed qweight tensor for MatMulNBits.
 

src/mobius/_weight_utils_test.py

@@ -21,6 +21,7 @@
     tie_word_embeddings,
     vlm_decoder_weights,
     vlm_embedding_weights,
+    vlm_vision_weights,
 )
 
 
@@ -396,6 +397,44 @@ def test_custom_keyword(self):
         assert list(result.keys()) == ["word_embedding.weight"]
 
 
+class TestVlmVisionWeights:
+    """Tests for vlm_vision_weights."""
+
+    def test_filters_and_renames(self):
+        """Keeps prefixed keys and renames fc1/fc2."""
+        fc1 = torch.randn(4, 8)
+        fc2 = torch.randn(8, 4)
+        sd = {
+            "vision_tower.encoder.layers.0.mlp.fc1.weight": fc1,
+            "vision_tower.encoder.layers.0.mlp.fc2.weight": fc2,
+            "multi_modal_projector.linear.weight": torch.randn(4),
+            "language_model.model.layers.0.weight": torch.randn(4),
+        }
+        result = vlm_vision_weights(sd, ("vision_tower.", "multi_modal_projector."))
+        assert set(result.keys()) == {
+            "vision_tower.encoder.layers.0.mlp.up_proj.weight",
+            "vision_tower.encoder.layers.0.mlp.down_proj.weight",
+            "multi_modal_projector.linear.weight",
+        }
+        assert result["vision_tower.encoder.layers.0.mlp.up_proj.weight"].data_ptr() == (
+            fc1.data_ptr()
+        )
+
+    def test_single_prefix(self):
+        """Works with a single-element prefix tuple."""
+        sd = {
+            "vision_model.layers.0.mlp.fc1.weight": torch.randn(2),
+            "other.weight": torch.randn(2),
+        }
+        result = vlm_vision_weights(sd, ("vision_model.",))
+        assert list(result.keys()) == ["vision_model.layers.0.mlp.up_proj.weight"]
+
+    def test_empty_when_no_match(self):
+        """Returns empty dict when no key matches the prefixes."""
+        sd = {"language_model.layers.0.weight": torch.randn(2)}
+        assert vlm_vision_weights(sd, ("vision_tower.",)) == {}
+
+
 class TestPreprocessGptqWeights:
     """Tests for GPTQ weight preprocessing.
 

src/mobius/models/gemma3.py

@@ -18,7 +18,11 @@
 from onnxscript import OpBuilder, nn
 
 from mobius._configs import ArchitectureConfig
-from mobius._weight_utils import vlm_decoder_weights, vlm_embedding_weights
+from mobius._weight_utils import (
+    vlm_decoder_weights,
+    vlm_embedding_weights,
+    vlm_vision_weights,
+)
 from mobius.components import (
     Gemma3MultiModalProjector,
     Linear,
@@ -93,15 +97,7 @@ def forward(self, op: OpBuilder, pixel_values: ir.Value):
     def preprocess_weights(
         self, state_dict: dict[str, torch.Tensor]
     ) -> dict[str, torch.Tensor]:
-        renamed: dict[str, torch.Tensor] = {}
-        for key, value in state_dict.items():
-            if not key.startswith(("vision_tower.", "multi_modal_projector.")):
-                continue
-            # VisionModel MLP uses up_proj/down_proj; HF uses fc1/fc2
-            key = key.replace(".mlp.fc1.", ".mlp.up_proj.")
-            key = key.replace(".mlp.fc2.", ".mlp.down_proj.")
-            renamed[key] = value
-        return renamed
+        return vlm_vision_weights(state_dict, ("vision_tower.", "multi_modal_projector."))
 
 
 class _Gemma3EmbeddingModel(nn.Module):

src/mobius/models/llava.py

@@ -27,7 +27,11 @@
 from onnxscript import OpBuilder, nn
 
 from mobius._configs import ArchitectureConfig
-from mobius._weight_utils import vlm_decoder_weights, vlm_embedding_weights
+from mobius._weight_utils import (
+    vlm_decoder_weights,
+    vlm_embedding_weights,
+    vlm_vision_weights,
+)
 from mobius.components import (
     Embedding,
     Linear,
@@ -92,15 +96,7 @@ def forward(self, op: OpBuilder, pixel_values: ir.Value):
     def preprocess_weights(
         self, state_dict: dict[str, torch.Tensor]
     ) -> dict[str, torch.Tensor]:
-        renamed: dict[str, torch.Tensor] = {}
-        for key, value in state_dict.items():
-            if not key.startswith(("vision_tower.", "multi_modal_projector.")):
-                continue
-            # VisionModel MLP uses up_proj/down_proj; HF uses fc1/fc2
-            key = key.replace(".mlp.fc1.", ".mlp.up_proj.")
-            key = key.replace(".mlp.fc2.", ".mlp.down_proj.")
-            renamed[key] = value
-        return renamed
+        return vlm_vision_weights(state_dict, ("vision_tower.", "multi_modal_projector."))
 
 
 class _LLaVAEmbeddingModel(nn.Module):

src/mobius/models/mllama.py

@@ -23,7 +23,11 @@
 from onnxscript import OpBuilder, nn
 
 from mobius._configs import ArchitectureConfig, MllamaConfig
-from mobius._weight_utils import vlm_decoder_weights, vlm_embedding_weights
+from mobius._weight_utils import (
+    vlm_decoder_weights,
+    vlm_embedding_weights,
+    vlm_vision_weights,
+)
 from mobius.components import (
     MLP,
     DecoderLayer,
@@ -317,15 +321,7 @@ def forward(self, op: OpBuilder, pixel_values: ir.Value):
     def preprocess_weights(
         self, state_dict: dict[str, torch.Tensor]
     ) -> dict[str, torch.Tensor]:
-        renamed: dict[str, torch.Tensor] = {}
-        for key, value in state_dict.items():
-            if not key.startswith("vision_model."):
-                continue
-            # VisionModel MLP uses up_proj/down_proj; HF uses fc1/fc2
-            key = key.replace(".mlp.fc1.", ".mlp.up_proj.")
-            key = key.replace(".mlp.fc2.", ".mlp.down_proj.")
-            renamed[key] = value
-        return renamed
+        return vlm_vision_weights(state_dict, ("vision_model.",))
 
 
 class _MllamaEmbeddingModel(nn.Module):