Address Phi-3.5 vision review feedback

Use the public components API for CLIP, broaden the CLIP vision config contract, and filter unused CLIP layer weights when Phi-3.5-Vision exports an intermediate feature layer. Resolve cached Phi-3.5 projector shards via the HuggingFace cache index instead of requiring a full local snapshot, and remove the stale positive image-token constant. Add assertions covering skipped vision tower weights.

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

Author: justinchuby

src/mobius/models/_phi3_vision_projector.py

@@ -27,6 +27,7 @@
 
 import dataclasses
 import glob
+import json
 import os
 
 import numpy as np
@@ -42,6 +43,8 @@
 _PROJECTION_FIRST_BIAS_KEY = _PROJECTOR_PREFIX + "img_projection.0.bias"
 _PROJECTION_SECOND_WEIGHT_KEY = _PROJECTOR_PREFIX + "img_projection.2.weight"
 _PROJECTION_SECOND_BIAS_KEY = _PROJECTOR_PREFIX + "img_projection.2.bias"
+_WEIGHT_INDEX_NAME = "model.safetensors.index.json"
+_SINGLE_WEIGHT_NAME = "model.safetensors"
 
 # CLIP ViT-L/14-336 produces a 24x24 patch grid (576 patches).
 _PATCH_GRID_SIDE = 24
@@ -75,18 +78,46 @@ class Phi3VisionProjectorWeights:
     projection_second_bias: np.ndarray
 
 
-def _resolve_checkpoint_directory(model_id_or_directory: str) -> str:
-    """Return a local directory that contains the checkpoint safetensors.
+def _resolve_checkpoint_shard_paths(model_id_or_directory: str) -> list[str]:
+    """Return local safetensors shard paths for a checkpoint.
 
     Accepts either a local path or a HuggingFace hub model id (which is
     resolved via the local cache; the weights must already be downloaded).
     """
     if os.path.isdir(model_id_or_directory):
-        return model_id_or_directory
+        return sorted(glob.glob(os.path.join(model_id_or_directory, "*.safetensors")))
 
-    from huggingface_hub import snapshot_download
+    from huggingface_hub import hf_hub_download
+    from huggingface_hub.errors import EntryNotFoundError, LocalEntryNotFoundError
 
-    return snapshot_download(model_id_or_directory, local_files_only=True)
+    try:
+        index_path = hf_hub_download(
+            repo_id=model_id_or_directory,
+            filename=_WEIGHT_INDEX_NAME,
+            local_files_only=True,
+        )
+        with open(index_path, encoding="utf-8") as f:
+            index = json.load(f)
+        filenames = sorted(set(index["weight_map"].values()))
+    except (EntryNotFoundError, LocalEntryNotFoundError):
+        filenames = [_SINGLE_WEIGHT_NAME]
+
+    shard_paths: list[str] = []
+    for filename in filenames:
+        try:
+            shard_paths.append(
+                hf_hub_download(
+                    repo_id=model_id_or_directory,
+                    filename=filename,
+                    local_files_only=True,
+                )
+            )
+        except (EntryNotFoundError, LocalEntryNotFoundError) as exc:
+            raise FileNotFoundError(
+                f"Could not find cached safetensors shard {filename!r} for "
+                f"{model_id_or_directory!r}. Build/load the model weights first."
+            ) from exc
+    return shard_paths
 
 
 def load_phi3_vision_projector_weights(
@@ -111,11 +142,10 @@ def load_phi3_vision_projector_weights(
     """
     from safetensors import safe_open
 
-    directory = _resolve_checkpoint_directory(model_id_or_directory)
-    shard_paths = sorted(glob.glob(os.path.join(directory, "*.safetensors")))
+    shard_paths = _resolve_checkpoint_shard_paths(model_id_or_directory)
     if not shard_paths:
         raise FileNotFoundError(
-            f"No .safetensors shards found in checkpoint directory: {directory}"
+            f"No .safetensors shards found for checkpoint: {model_id_or_directory}"
         )
 
     wanted_keys = {
@@ -136,7 +166,8 @@ def load_phi3_vision_projector_weights(
     missing = wanted_keys - collected.keys()
     if missing:
         raise KeyError(
-            f"Missing Phi-3.5-Vision projector weights in {directory}: {sorted(missing)}"
+            "Missing Phi-3.5-Vision projector weights in "
+            f"{model_id_or_directory}: {sorted(missing)}"
         )
 
     return Phi3VisionProjectorWeights(

src/mobius/models/clip.py

@@ -5,21 +5,38 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol
 
 import torch
 from onnxscript import OpBuilder, nn
 
 from mobius._configs import ArchitectureConfig
-from mobius.components import FCMLP
-from mobius.components._common import INT64_MAX, Embedding, LayerNorm
-from mobius.components._conv import Conv2d, Conv2dNoBias
-from mobius.components._encoder import EncoderAttention
+from mobius.components import (
+    FCMLP,
+    INT64_MAX,
+    Conv2d,
+    Conv2dNoBias,
+    Embedding,
+    EncoderAttention,
+    LayerNorm,
+)
 
 if TYPE_CHECKING:
     import onnx_ir as ir
 
 
+class _CLIPVisionConfig(Protocol):
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    image_size: int
+    patch_size: int
+    num_channels: int
+    rms_norm_eps: float
+    hidden_act: str | None
+
+
 class ClipVisionConfigView:
     """Adapter exposing a :class:`VisionConfig` under CLIP's field names.
 
@@ -47,7 +64,7 @@ def __init__(self, vision_config, *, default_hidden_act: str = "quick_gelu"):
 class _CLIPVisionEmbeddings(nn.Module):
     """CLIP vision embeddings: Conv2d patch + CLS token + position embeddings."""
 
-    def __init__(self, config: ArchitectureConfig):
+    def __init__(self, config: _CLIPVisionConfig):
         super().__init__()
         hidden_size = config.hidden_size
         patch_size = config.patch_size
@@ -117,7 +134,7 @@ def forward(self, op: OpBuilder, x: ir.Value):
 class _CLIPVisionEncoderLayer(nn.Module):
     """CLIP vision encoder layer: pre-norm with LayerNorm."""
 
-    def __init__(self, config: ArchitectureConfig):
+    def __init__(self, config: _CLIPVisionConfig):
         super().__init__()
         self.self_attn = EncoderAttention(config.hidden_size, config.num_attention_heads)
         self.layer_norm1 = LayerNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -195,7 +212,7 @@ class CLIPVisionModel(nn.Module):
 
     def __init__(
         self,
-        config: ArchitectureConfig,
+        config: _CLIPVisionConfig,
         *,
         feature_layer: int | None = None,
         drop_class_token: bool = False,
@@ -250,9 +267,17 @@ def preprocess_weights(
         self, state_dict: dict[str, torch.Tensor]
     ) -> dict[str, torch.Tensor]:
         new_state_dict = {}
+        num_encoder_layers = len(self.encoder)
         for name, tensor in state_dict.items():
             new_name = _rename_clip_vision_weight(name)
             if new_name is not None:
+                if self.post_layernorm is None and new_name.startswith("post_layernorm."):
+                    continue
+                if new_name.startswith("encoder."):
+                    parts = new_name.split(".", 2)
+                    if len(parts) >= 3 and parts[1].isdigit():
+                        if int(parts[1]) >= num_encoder_layers:
+                            continue
                 new_state_dict[new_name] = tensor
         return new_state_dict
 
@@ -326,7 +351,7 @@ def _rename_clip_vision_weight(name: str) -> str | None:
 class _SigLIPVisionEmbeddings(nn.Module):
     """SigLIP vision embeddings: Conv2d patch + position embeddings (no CLS token)."""
 
-    def __init__(self, config: ArchitectureConfig):
+    def __init__(self, config: _CLIPVisionConfig):
         super().__init__()
         hidden_size = config.hidden_size
         patch_size = config.patch_size

src/mobius/models/phi3_v.py

@@ -49,7 +49,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING
 
 import torch
 from onnxscript import nn
@@ -71,11 +71,6 @@
 if TYPE_CHECKING:
     import onnx_ir as ir
 
-# Phi-3-Vision: <|image|> token id (kept for reference; the HF processor marks
-# image slots with negative placeholder ids, not this positive id — see
-# ``_Phi3VEmbeddingModel.forward``).
-_IMAGE_TOKEN_ID = 32044
-
 # Upper bound (magnitude) for negative image placeholder ids, mirroring
 # ``modeling_phi3_v.MAX_INPUT_ID = int(1e9)``. Image positions satisfy
 # ``-_MAX_INPUT_ID < input_ids < 0``.
@@ -201,7 +196,7 @@ def __init__(self, config: ArchitectureConfig):
         assert config.vision is not None, "Phi3-V requires a vision config"
         clip_config = ClipVisionConfigView(config.vision)
         self.vision_tower = CLIPVisionModel(
-            cast(ArchitectureConfig, clip_config),
+            clip_config,
             feature_layer=config.vision.feature_layer,
             drop_class_token=True,
         )
@@ -223,11 +218,14 @@ def preprocess_weights(
         The ``img_projection`` and ``sub_GN``/``glb_GN`` tensors are dropped
         (host-side HD transform).
         """
-        renamed: dict[str, torch.Tensor] = {}
+        clip_state_dict: dict[str, torch.Tensor] = {}
         for key, value in state_dict.items():
-            new_key = _rename_phi3v_vision_weight(key)
-            if new_key is not None:
-                renamed[new_key] = value
+            if key.startswith(_VISION_TOWER_PREFIX):
+                clip_state_dict[key[len(_VISION_TOWER_PREFIX) :]] = value
+        renamed = {
+            "vision_tower." + key: value
+            for key, value in self.vision_tower.preprocess_weights(clip_state_dict).items()
+        }
         return renamed
 
 
@@ -351,11 +349,11 @@ def preprocess_weights(
         """
         renamed: dict[str, torch.Tensor] = {}
 
+        for key, value in self.vision_encoder.preprocess_weights(state_dict).items():
+            renamed["vision_encoder." + key] = value
+
         for key, value in state_dict.items():
-            vision_key = _rename_phi3v_vision_weight(key)
-            if vision_key is not None:
-                renamed["vision_encoder." + vision_key] = value
-            elif key.startswith("model.vision_embed_tokens."):
+            if key.startswith("model.vision_embed_tokens."):
                 # img_projection / sub_GN / glb_GN and any other vision-embed
                 # state is host-side (HD transform) — skip it.
                 pass

src/mobius/models/phi3_v_test.py

@@ -127,6 +127,8 @@ def test_feature_layer_maps_only_needed_layers(self):
         assert "vision_tower.embeddings.class_embedding" in renamed
         assert "vision_tower.embeddings.patch_embedding.projection.weight" in renamed
         assert "vision_tower.encoder.0.mlp.up_proj.weight" in renamed
+        assert "vision_tower.encoder.2.mlp.up_proj.weight" not in renamed
+        assert "vision_tower.post_layernorm.weight" not in renamed
         # Projector + separator are host-side.
         assert not any("projection.0" in k or "sub_GN" in k for k in renamed)
         # Only vision_tower.* names are produced.