Phi-3.5 exporter support: LongRoPE su + real CLIP vision + golden coverage (reworks #430)

scripts/generate_golden.py

@@ -214,7 +214,9 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None:
         torch_forward,
     )
 
-    model, tokenizer = load_torch_model(case.model_id, device=device)
+    model, tokenizer = load_torch_model(
+        case.model_id, device=device, trust_remote_code=case.trust_remote_code
+    )
 
     encoded = tokenizer(case.prompts[0], return_tensors="np", padding=False)
     input_ids = encoded["input_ids"]

src/mobius/_config_resolver_test.py

@@ -118,7 +118,9 @@ def test_phi3_resolves(self):
         )
         result = _config_from_hf(hf)
         assert isinstance(result, ArchitectureConfig)
-        assert result.rope_type == "su"
+        # The legacy Phi-3 ``"su"`` rope_type is canonicalized to ``"longrope"``
+        # (they name the same LongRoPE algorithm); see _canonical_rope_type.
+        assert result.rope_type == "longrope"
 
     def test_phi3_partial_rotary(self):
         hf = _fake_hf_config("phi3", partial_rotary_factor=0.5)

src/mobius/_configs/_base.py

@@ -134,6 +134,28 @@ def _first_not_none(*values, default=None):
     return default
 
 
+# Legacy rope_type spellings that name the same algorithm as a canonical
+# rope_type understood by ``initialize_rope``. Phi-3/Phi-3.5 checkpoints
+# label LongRoPE as ``"su"`` (short/long-factor scaled rotary embeddings);
+# newer HuggingFace configs call the identical algorithm ``"longrope"``.
+# Both must resolve to the same code path, so we canonicalize the alias at
+# extraction time rather than special-casing ``"su"`` downstream.
+_ROPE_TYPE_ALIASES: dict[str, str] = {
+    "su": "longrope",
+}
+
+
+def _canonical_rope_type(rope_type: str | None) -> str | None:
+    """Map legacy rope_type aliases to their canonical spelling.
+
+    Returns the input unchanged when it is not a known alias (including
+    ``None`` and ``"default"``).
+    """
+    if rope_type is None:
+        return None
+    return _ROPE_TYPE_ALIASES.get(rope_type, rope_type)
+
+
 def _leading_layer_type_count(layer_types, target: str) -> int:
     count = 0
     for layer_type in layer_types or ():
@@ -247,12 +269,14 @@ def _extract_rope_config(config) -> RoPEConfig | None:
     rope_parameters = raw_rope_parameters or {}
 
     return RoPEConfig(
-        rope_type=_first_not_none(
-            rope_scaling.get("rope_type", None),
-            rope_scaling.get("type", None),
-            rope_parameters.get("rope_type", None),
-            _nested_rope_type(rope_scaling, "full_attention"),
-            default="default",
+        rope_type=_canonical_rope_type(
+            _first_not_none(
+                rope_scaling.get("rope_type", None),
+                rope_scaling.get("type", None),
+                rope_parameters.get("rope_type", None),
+                _nested_rope_type(rope_scaling, "full_attention"),
+                default="default",
+            )
         ),
         rope_theta=_first_not_none(
             getattr(config, "rope_theta", None),

src/mobius/_configs/_extractors_test.py

@@ -285,6 +285,44 @@ def test_hunyuan_vl_mot_filter_only_fires_for_that_model(loaded_vision_hooks):
     assert out == {}
 
 
+def test_phi3_v_vision_parses_img_processor(loaded_vision_hooks):
+    """Phi-3.5-Vision reads image_dim_out + layer_idx from the img_processor dict."""
+    cfg = _FakeHFConfig(
+        img_processor={
+            "name": "clip_vision_model",
+            "model_name": "openai/clip-vit-large-patch14-336",
+            "image_dim_out": 1024,
+            "num_img_tokens": 144,
+        }
+    )
+    out = _extractors.extract_vision_config(cfg, None, "phi3_v")
+    vision = out["vision"]
+    # CLIP ViT-L/14-336 geometry.
+    assert vision.hidden_size == 1024
+    assert vision.intermediate_size == 4096
+    assert vision.num_hidden_layers == 24
+    assert vision.num_attention_heads == 16
+    assert vision.image_size == 336
+    assert vision.patch_size == 14
+    assert vision.hidden_act == "quick_gelu"
+    # layer_idx defaults to -2 (feature extraction skips the last layer).
+    assert vision.feature_layer == -2
+
+
+def test_phi3_v_vision_respects_explicit_layer_idx(loaded_vision_hooks):
+    cfg = _FakeHFConfig(
+        img_processor={"name": "clip_vision_model", "image_dim_out": 1024, "layer_idx": -1}
+    )
+    out = _extractors.extract_vision_config(cfg, None, "phi3_v")
+    assert out["vision"].feature_layer == -1
+
+
+def test_phi3_v_vision_does_not_fire_for_unrelated_model(loaded_vision_hooks):
+    cfg = _FakeHFConfig(img_processor={"image_dim_out": 1024})
+    out = _extractors.extract_vision_config(cfg, None, "llama")
+    assert out == {}
+
+
 def test_phi4mm_image_token_id_survives_default_hook(loaded_vision_hooks):
     """Per-model image_token_id must not be clobbered by _vision_default.
 

src/mobius/_configs/_sub_configs.py

@@ -74,6 +74,10 @@ class VisionConfig:
     pooling_kernel_size: int | None = None
     # MLP activation for vision encoder layers (e.g. "gelu_pytorch_tanh" for Gemma4 SigLIP)
     hidden_act: str | None = None
+    # CLIP-style feature extraction: which ``hidden_states`` index to output
+    # (HuggingFace convention, e.g. -2 for Phi-3.5-Vision). ``None`` means use
+    # the final hidden state (all layers + post_layernorm).
+    feature_layer: int | None = None
 
 
 @dataclasses.dataclass

src/mobius/_configs/per_model/_phi_vision.py

@@ -10,12 +10,19 @@
 
 @register_vision_hook("phi3_v")
 def _phi3_v_vision(config, parent_config, model_type: str, fields: dict):
-    """Extract Phi-3 Vision's CLIP encoder fields from ``img_processor``."""
+    """Extract Phi-3 Vision's CLIP encoder fields from ``img_processor``.
+
+    Phi-3/3.5-Vision store the CLIP ViT-L/14-336 geometry in a non-standard
+    ``img_processor`` dict rather than a standard ``vision_config``. The encoder
+    itself is the fixed openai/clip-vit-large-patch14-336 architecture, so the
+    geometry is constant; only ``image_dim_out`` and the feature ``layer_idx``
+    are read from the checkpoint.
+    """
     vision_source = parent_config or config
     img_processor = getattr(vision_source, "img_processor", None)
-    image_dim_out = (
-        img_processor.get("image_dim_out", 1024) if isinstance(img_processor, dict) else 1024
-    )
+    if not isinstance(img_processor, dict):
+        img_processor = {}
+    image_dim_out = img_processor.get("image_dim_out", 1024)
     fields.update(
         hidden_size=image_dim_out,
         intermediate_size=4096,
@@ -24,6 +31,10 @@ def _phi3_v_vision(config, parent_config, model_type: str, fields: dict):
         image_size=336,
         patch_size=14,
         norm_eps=1e-5,
+        hidden_act="quick_gelu",
+        # HuggingFace Phi3V selects hidden_states[layer_idx] (default -2) and
+        # keeps only the patch tokens (type_feature="patch").
+        feature_layer=img_processor.get("layer_idx", -2),
     )
     fields.setdefault("image_token_id", 32044)
     return None

src/mobius/_testing/torch_reference.py

@@ -13,6 +13,68 @@
 logger = logging.getLogger(__name__)
 
 
+def _install_dynamic_cache_legacy_shims() -> None:
+    """Restore ``DynamicCache`` methods removed in transformers 5.x.
+
+    transformers 5.x removed ``DynamicCache.from_legacy_cache`` and
+    ``DynamicCache.get_usable_length``, but some ``trust_remote_code``
+    checkpoints (e.g. Phi-3.5 family) still call them from their bundled
+    modeling code. Re-add minimal, behaviour-preserving implementations so
+    reference generation continues to work on transformers 5.x. Both the
+    causal-LM and multimodal loaders rely on this shim.
+    """
+    import transformers
+
+    if not hasattr(transformers.DynamicCache, "from_legacy_cache"):
+
+        @classmethod  # type: ignore[misc]
+        def _from_legacy_cache(cls, past_key_values=None):  # type: ignore[misc]
+            cache = cls()
+            if past_key_values is not None:
+                for i, (k, v) in enumerate(past_key_values):
+                    cache.update(k, v, i)
+            return cache
+
+        transformers.DynamicCache.from_legacy_cache = _from_legacy_cache
+
+    if not hasattr(transformers.DynamicCache, "to_legacy_cache"):
+
+        def _to_legacy_cache(self):  # type: ignore[misc]
+            legacy_cache = []
+            for layer in getattr(self, "layers", []):
+                keys = getattr(layer, "keys", None)
+                values = getattr(layer, "values", None)
+                if keys is None or values is None:
+                    continue
+                legacy_cache.append((keys, values))
+            return tuple(legacy_cache)
+
+        transformers.DynamicCache.to_legacy_cache = _to_legacy_cache  # type: ignore[method-assign]
+
+    if not hasattr(transformers.DynamicCache, "seen_tokens"):
+
+        def _seen_tokens(self):  # type: ignore[misc]
+            # Legacy alias for the number of tokens the cache has seen so
+            # far, which for a full (non-windowed) cache equals the current
+            # sequence length of layer 0.
+            return self.get_seq_length(0)
+
+        transformers.DynamicCache.seen_tokens = property(_seen_tokens)  # type: ignore[assignment]
+
+    if not hasattr(transformers.DynamicCache, "get_usable_length"):
+
+        def _get_usable_length(self, new_seq_length: int, layer_idx: int = 0) -> int:  # type: ignore[misc]
+            # Mirrors the removed transformers implementation: the usable
+            # length is the number of tokens ALREADY cached for this layer
+            # (the past length), NOT including the incoming ``new_seq_length``.
+            # Returning ``previous + new`` double-counts the query tokens and
+            # makes callers (e.g. the Phi-3.5 bundled modeling code) build a
+            # causal mask of the wrong size.
+            return self.get_seq_length(layer_idx)
+
+        transformers.DynamicCache.get_usable_length = _get_usable_length  # type: ignore[method-assign]
+
+
 def _fix_nemotron_h_init_weights(model: torch.nn.Module, model_id: str) -> None:
     """Restore Mamba2 params from checkpoint after HF clobbers them.
 
@@ -97,25 +159,38 @@ def load_torch_model(
     model_id: str,
     dtype: torch.dtype = torch.float32,
     device: str = "cpu",
+    trust_remote_code: bool = True,
 ):
     """Load a HuggingFace causal LM model for reference inference.
 
     Args:
         model_id: HuggingFace model identifier.
         dtype: Model dtype (default float32 for numerical comparison).
         device: Device to load on.
+        trust_remote_code: Whether to load the checkpoint's bundled modeling
+            code. Defaults to ``True`` for backward compatibility. Pass
+            ``False`` for models (e.g. Phi-3.5-mini) whose ``model_type`` is
+            natively supported by the installed transformers, so the
+            transformers-5.x-compatible implementation is used instead of an
+            older bundled ``modeling_*.py`` that relies on removed cache APIs.
 
     Returns:
         Tuple of (model, tokenizer).
     """
     import transformers
 
-    tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    _install_dynamic_cache_legacy_shims()
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        model_id, trust_remote_code=trust_remote_code
+    )
 
     # NemotronH: disable rescale_prenorm_residual before loading to
     # prevent _init_weights from corrupting out_proj.weight with
     # random kaiming_uniform_ initialization after checkpoint loading.
-    config = transformers.AutoConfig.from_pretrained(model_id, trust_remote_code=True)
+    config = transformers.AutoConfig.from_pretrained(
+        model_id, trust_remote_code=trust_remote_code
+    )
     if getattr(config, "model_type", None) == "nemotron_h":
         config.rescale_prenorm_residual = False
 
@@ -124,7 +199,7 @@ def load_torch_model(
         config=config,
         dtype=dtype,
         device_map=device,
-        trust_remote_code=True,
+        trust_remote_code=trust_remote_code,
     )
     _fix_nemotron_h_init_weights(model, model_id)
     model.eval()
@@ -160,25 +235,7 @@ def load_torch_multimodal_model(
     # Shim: transformers 5.x removed DynamicCache.from_legacy_cache and
     # DynamicCache.get_usable_length, but some trust_remote_code models
     # (e.g. Phi-3.5-vision-instruct) still call them.
-    if not hasattr(transformers.DynamicCache, "from_legacy_cache"):
-
-        @classmethod  # type: ignore[misc]
-        def _from_legacy_cache(cls, past_key_values=None):  # type: ignore[misc]
-            cache = cls()
-            if past_key_values is not None:
-                for i, (k, v) in enumerate(past_key_values):
-                    cache.update(k, v, i)
-            return cache
-
-        transformers.DynamicCache.from_legacy_cache = _from_legacy_cache
-
-    if not hasattr(transformers.DynamicCache, "get_usable_length"):
-
-        def _get_usable_length(self, new_seq_length: int, layer_idx: int = 0) -> int:  # type: ignore[misc]
-            # Without a sliding window, the full cache is usable.
-            return self.get_seq_length(layer_idx) + new_seq_length
-
-        transformers.DynamicCache.get_usable_length = _get_usable_length  # type: ignore[method-assign]
+    _install_dynamic_cache_legacy_shims()
 
     # Load config and force eager attention so flash_attn is not required.
     # Some models (e.g. Phi-3.5-vision-instruct) hardcode flash_attention_2
@@ -190,13 +247,32 @@ def _get_usable_length(self, new_seq_length: int, layer_idx: int = 0) -> int:  #
     # Some trust_remote_code VLMs (e.g. Phi-3-Vision) are registered as
     # AutoModelForCausalLM, not AutoModelForImageTextToText.  Fall back
     # gracefully so golden generation works for both.
-    model_kwargs = dict(config=config, dtype=dtype, device_map=device, trust_remote_code=True)
+    #
+    # The weight-dtype keyword was renamed from ``torch_dtype`` to ``dtype`` in
+    # transformers 5.x; support both so offline golden generation also works on
+    # the older transformers (e.g. 4.43) required by 4.x-era remote-code models.
+    base_kwargs = dict(config=config, device_map=device, trust_remote_code=True)
+
+    def _load_from_pretrained(auto_cls):
+        try:
+            return auto_cls.from_pretrained(model_id, dtype=dtype, **base_kwargs)
+        except TypeError:
+            return auto_cls.from_pretrained(model_id, torch_dtype=dtype, **base_kwargs)
+
     try:
-        model = transformers.AutoModelForImageTextToText.from_pretrained(
-            model_id, **model_kwargs
-        )
-    except (ValueError, KeyError):
-        model = transformers.AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)
+        image_text_to_text_cls = transformers.AutoModelForImageTextToText
+    except AttributeError:
+        # Older transformers (e.g. 4.43, used for offline golden generation of
+        # 4.x-era remote-code checkpoints) predate AutoModelForImageTextToText.
+        image_text_to_text_cls = None
+    model = None
+    if image_text_to_text_cls is not None:
+        try:
+            model = _load_from_pretrained(image_text_to_text_cls)
+        except (ValueError, KeyError):
+            model = None
+    if model is None:
+        model = _load_from_pretrained(transformers.AutoModelForCausalLM)
 
     model.eval()