From effc8ed3a87cd2bb43d91dd7eed517959c6e117e Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 25 Jul 2026 15:12:16 +0000 Subject: [PATCH 01/10] Support Phi LongRoPE rope_type="su" and unblock Phi-3.5-mini golden Phi-3/Phi-3.5 checkpoints label LongRoPE as rope_scaling type "su" (short/ long-factor scaled rotary embeddings); newer HuggingFace configs spell the identical algorithm "longrope". The exporter recognized only "longrope" and raised `ValueError: Unsupported rope type: su` when loading e.g. the Phi-3.5-vision text tower. Config extraction: - Add a rope_type alias map in mobius/_configs/_base.py (_ROPE_TYPE_ALIASES = {"su": "longrope"} + _canonical_rope_type) so both spellings resolve to the existing LongRope code path. "su" is canonicalized at extraction time rather than special-cased downstream. Golden reference (transformers 5.x): - Factor the DynamicCache legacy shims into one shared helper _install_dynamic_cache_legacy_shims() used by both the causal-LM and multimodal loaders. Fix a pre-existing bug where get_usable_length returned past + new (must be past length only), which made the bundled Phi-3.5 modeling build a doubled causal mask; add to_legacy_cache and seen_tokens. - load_torch_model now honors a trust_remote_code flag (default True) and the golden generator passes the YAML case's value. Set phi-3_5-mini to trust_remote_code: false so the transformers-native phi3 implementation is used; the checkpoint's bundled modeling_phi3.py depends on cache APIs removed in transformers 5.x. - Un-skip testdata/cases/causal-lm/phi-3_5-mini.yaml and add regenerated L4/L5 golden files. L4 (prefill argmax) and L5 (generation) golden tests now pass. Tests: - Add TestLongRopeAliasExtraction: su and longrope produce identical RoPEConfig, su dispatches to LongRope, end-to-end phi3 graph build, missing original_max_position_embeddings, factor-length mismatch, and non-alias-unchanged cases. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- scripts/generate_golden.py | 4 +- src/mobius/_configs/_base.py | 36 ++++- src/mobius/_testing/torch_reference.py | 101 +++++++++--- testdata/cases/causal-lm/phi-3_5-mini.yaml | 8 +- testdata/golden/causal-lm/phi-3_5-mini.json | 41 +++++ .../causal-lm/phi-3_5-mini_generation.json | 27 ++++ tests/build_graph_test.py | 145 ++++++++++++++++++ 7 files changed, 330 insertions(+), 32 deletions(-) create mode 100644 testdata/golden/causal-lm/phi-3_5-mini.json create mode 100644 testdata/golden/causal-lm/phi-3_5-mini_generation.json diff --git a/scripts/generate_golden.py b/scripts/generate_golden.py index e603d536..374d7810 100644 --- a/scripts/generate_golden.py +++ b/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"] diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 7ebbf390..2ecbc655 100644 --- a/src/mobius/_configs/_base.py +++ b/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), diff --git a/src/mobius/_testing/torch_reference.py b/src/mobius/_testing/torch_reference.py index 00b89453..c1fb8479 100644 --- a/src/mobius/_testing/torch_reference.py +++ b/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: tuple = () + 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 += ((keys, values),) + return 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,6 +159,7 @@ 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. @@ -104,18 +167,30 @@ def load_torch_model( 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 diff --git a/testdata/cases/causal-lm/phi-3_5-mini.yaml b/testdata/cases/causal-lm/phi-3_5-mini.yaml index 6fccab4f..7ea732d1 100644 --- a/testdata/cases/causal-lm/phi-3_5-mini.yaml +++ b/testdata/cases/causal-lm/phi-3_5-mini.yaml @@ -14,6 +14,8 @@ generation: max_new_tokens: 20 do_sample: false -trust_remote_code: true -notes: "Phi-3.5 with LongRoPE positional embeddings." -skip_reason: "DynamicCache.from_legacy_cache removed in transformers 5.x." +trust_remote_code: false +notes: "Phi-3.5 with LongRoPE positional embeddings. trust_remote_code is + false so the transformers-native phi3 implementation is used; the + checkpoint's bundled modeling_phi3.py relies on cache APIs removed in + transformers 5.x." diff --git a/testdata/golden/causal-lm/phi-3_5-mini.json b/testdata/golden/causal-lm/phi-3_5-mini.json new file mode 100644 index 00000000..bbfc3690 --- /dev/null +++ b/testdata/golden/causal-lm/phi-3_5-mini.json @@ -0,0 +1,41 @@ +{ + "top1_id": 13, + "top2_id": 376, + "top10_ids": [ + 13, + 376, + 29871, + 450, + 306, + 525, + 259, + 1346, + 30004, + 512 + ], + "top10_logits": [ + "0x1.eb13bc0000000p+5", + "0x1.c87d580000000p+5", + "0x1.c6beee0000000p+5", + "0x1.c5d3bc0000000p+5", + "0x1.c11f500000000p+5", + "0x1.ba65aa0000000p+5", + "0x1.b900720000000p+5", + "0x1.b804680000000p+5", + "0x1.b72cf20000000p+5", + "0x1.b463800000000p+5" + ], + "logits_summary": [ + "0x1.eb13bc0000000p+5", + "0x1.b3930c0000000p+4", + "0x1.54b4d67e56ddcp+5", + "0x1.c51609082cd93p+1" + ], + "input_ids": [ + 2266, + 338, + 590, + 26576, + 29901 + ] +} diff --git a/testdata/golden/causal-lm/phi-3_5-mini_generation.json b/testdata/golden/causal-lm/phi-3_5-mini_generation.json new file mode 100644 index 00000000..1ccadd43 --- /dev/null +++ b/testdata/golden/causal-lm/phi-3_5-mini_generation.json @@ -0,0 +1,27 @@ +{ + "model_id": "microsoft/Phi-3.5-mini-instruct", + "prompt": "Here is my poem:", + "generated_tokens": [ + 13, + 13, + 1576, + 6575, + 6166, + 297, + 263, + 12995, + 911, + 310, + 26080, + 29892, + 13, + 29925, + 475, + 1259, + 278, + 14744, + 411, + 298 + ], + "generated_text": "\n\nThe sun sets in a blaze of glory,\nPainting the sky with h" +} diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index fa1b4bf4..a96caec2 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -5997,3 +5997,148 @@ def test_window_without_flag_is_kept(self): self._cfg(model_type="mistral", sliding_window=4096), "mistral" ) assert cfg.sliding_window == 4096 + + +class TestLongRopeAliasExtraction: + """``rope_type`` alias handling for Phi LongRoPE. + + Phi-3/Phi-3.5 checkpoints label LongRoPE as ``"su"`` (short/long-factor + scaled rotary embeddings); newer HuggingFace configs spell the identical + algorithm ``"longrope"``. ``_extract_rope_config`` must canonicalize the + legacy ``"su"`` alias to ``"longrope"`` so both configs resolve to the + same ``LongRope`` code path. + """ + + _SHORT_FACTOR = [1.0] * (TINY_HEAD_DIM // 2) + _LONG_FACTOR = [2.0] * (TINY_HEAD_DIM // 2) + + @staticmethod + def _cfg(rope_scaling, **attrs): + from types import SimpleNamespace + + defaults = dict( + model_type="phi3", + hidden_size=TINY_HIDDEN, + intermediate_size=TINY_INTERMEDIATE, + num_attention_heads=TINY_HEADS, + num_key_value_heads=TINY_KV_HEADS, + head_dim=TINY_HEAD_DIM, + num_hidden_layers=TINY_LAYERS, + vocab_size=TINY_VOCAB, + max_position_embeddings=1024, + original_max_position_embeddings=128, + hidden_act="silu", + rms_norm_eps=1e-6, + rope_theta=10000.0, + rope_scaling=rope_scaling, + ) + defaults.update(attrs) + return SimpleNamespace(**defaults) + + def _scaling(self, rope_type_key, rope_type_value): + return { + rope_type_key: rope_type_value, + "short_factor": self._SHORT_FACTOR, + "long_factor": self._LONG_FACTOR, + } + + def test_su_and_longrope_produce_identical_rope_config(self): + """``type="su"`` and ``rope_type="longrope"`` extract identically.""" + from mobius._configs._base import _extract_rope_config + + su_config = _extract_rope_config(self._cfg(self._scaling("type", "su"))) + longrope_config = _extract_rope_config( + self._cfg(self._scaling("rope_type", "longrope")) + ) + + assert su_config is not None + assert su_config.rope_type == "longrope" + assert longrope_config.rope_type == "longrope" + assert su_config.original_max_position_embeddings == 128 + assert su_config.rope_type == longrope_config.rope_type + assert ( + su_config.rope_scaling["short_factor"] + == longrope_config.rope_scaling["short_factor"] + ) + assert ( + su_config.rope_scaling["long_factor"] + == longrope_config.rope_scaling["long_factor"] + ) + assert ( + su_config.original_max_position_embeddings + == longrope_config.original_max_position_embeddings + ) + + def test_su_alias_dispatches_to_longrope_module(self): + """A ``"su"`` config resolves to the ``LongRope`` runtime module.""" + from mobius.components._rotary_embedding import LongRope, initialize_rope + + config = ArchitectureConfig.from_transformers(self._cfg(self._scaling("type", "su"))) + assert config.rope_type == "longrope" + rope = initialize_rope(config) + assert isinstance(rope, LongRope) + + def test_su_graph_builds_end_to_end(self): + """A phi3 ``"su"`` config builds a valid ONNX graph without weights.""" + config = ArchitectureConfig.from_transformers(self._cfg(self._scaling("type", "su"))) + module = registry.get("phi3")(config) + task = get_task(_default_task_for_model("phi3")) + pkg = task.build(module, config) + model = pkg["model"] + + assert model.graph is not None + input_names = {inp.name for inp in model.graph.inputs} + assert "input_ids" in input_names + assert "position_ids" in input_names + output_names = {out.name for out in model.graph.outputs} + assert "logits" in output_names + + def test_missing_original_max_position_embeddings_still_maps_to_longrope(self): + """A ``"su"`` config without ``original_max_position_embeddings``. + + The alias must still resolve to ``longrope`` and ``LongRope`` falls + back to ``max_position_embeddings`` for the short cache length rather + than crashing. + """ + from mobius._configs._base import _extract_rope_config + from mobius.components._rotary_embedding import LongRope, initialize_rope + + config_source = self._cfg( + self._scaling("type", "su"), original_max_position_embeddings=None + ) + rope_config = _extract_rope_config(config_source) + assert rope_config.rope_type == "longrope" + assert rope_config.original_max_position_embeddings is None + + arch_config = ArchitectureConfig.from_transformers(config_source) + rope = initialize_rope(arch_config) + assert isinstance(rope, LongRope) + + def test_factor_length_mismatch_is_rejected(self): + """Short/long factor arrays must match the rotary dimension. + + A factor list whose length does not equal ``head_dim / 2`` cannot be + broadcast against the inverse-frequency vector, so ``LongRope`` + construction raises rather than silently producing a wrong cache. + """ + from mobius.components._rotary_embedding import initialize_rope + + bad_scaling = { + "type": "su", + "short_factor": [1.0] * (TINY_HEAD_DIM // 2 + 1), + "long_factor": [2.0] * (TINY_HEAD_DIM // 2 + 1), + } + config = ArchitectureConfig.from_transformers(self._cfg(bad_scaling)) + assert config.rope_type == "longrope" + with pytest.raises(ValueError, match="broadcast"): + initialize_rope(config) + + def test_non_alias_rope_types_are_unchanged(self): + """Canonicalization only rewrites known aliases, not other types.""" + from mobius._configs._base import _canonical_rope_type + + assert _canonical_rope_type("su") == "longrope" + assert _canonical_rope_type("longrope") == "longrope" + assert _canonical_rope_type("yarn") == "yarn" + assert _canonical_rope_type("default") == "default" + assert _canonical_rope_type(None) is None From a247660a68112a0a2e07e68f576eafc71ba2f31d Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 25 Jul 2026 15:14:04 +0000 Subject: [PATCH 02/10] phi3_v: real CLIP ViT-L/14-336 vision tower export (Phi-3.5-Vision) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #430 only added a skip_reason. This makes the mobius exporter actually load and build the Phi-3.5-Vision (model_type `phi3_v`) checkpoint's vision path, replacing the previous SigLIP-style stub that could neither load the CLIP CLS-token weights nor match the projector shape. Two concrete gaps fixed: 1. Non-standard vision config: the CLIP geometry lives in the `img_processor` dict, not a standard `vision_config`. The phi3_v vision hook now also parses `layer_idx` (feature layer, default -2) and sets `hidden_act=quick_gelu`; a general `feature_layer` field was added to VisionConfig. 2. CLIP CLS-token weights the exporter did not support: the vision encoder now uses the real CLIP tower (CLS token + pre_layrnorm + 577-entry position embedding) instead of the SigLIP `VisionModel`. It emits the exact HuggingFace `get_img_features` output: run to `hidden_states[-2]` (23 of 24 layers, no post_layernorm) and drop the CLS token. Implementation (DRY — reuses the existing CLIP infra in clip.py): - clip.py: `CLIPVisionModel` gains `feature_layer` + `drop_class_token`; only the encoder layers needed to reach the selected hidden state are built. `_rename_clip_vision_weight` now maps HF `patch_embedding.{weight,bias}` onto the `.projection` Conv sub-module (previously unmapped → unfed initializer), and CLIP's patch conv is now correctly bias-free. Added `ClipVisionConfigView` to bridge nested VisionConfig → CLIP's field names and `resolve_clip_feature_num_layers` for the hidden-states math. - phi3_v.py: `_Phi3VVisionEncoderModel` builds the CLIP tower and outputs raw patch features `(num_crops, 576, 1024)`; `_rename_phi3v_vision_weight` maps `img_processor.vision_model.*` → `vision_tower.*` and drops the host-side `img_projection` / `sub_GN` / `glb_GN` (HD feature transform). The `img_projection` MLP (fc1 is 3072x4096) and the HD 2x2 patch-merge with learnable separators stay host-side: the projector's 4096-wide input only exists after the image-size-dependent HD transform, which cannot be a static graph. The raw patch-feature tensor is the clean, faithful boundary. Tests (writing-tests skill): - clip_test.py: CLIP weight rename (patch/CLS/pre-layrnorm/mlp) + feature-layer resolver edge cases. - phi3_v_test.py: img_processor CLIP weight mapping + projector/separator drop. - _extractors_test.py: img_processor vision-config parsing (dims + layer_idx). - vision_integration_test.py: numeric parity for feature_layer=-2 + CLS drop vs a PyTorch CLIP reference (and CLIP patch conv is now bias-free). Verified: real Phi-3.5-vision checkpoint builds the vision encoder standalone, all 373 weight initializers populate, and ONNX Runtime produces (num_crops, 576, 1024). End-to-end golden un-skip additionally needs the rope_type=su text tower fix and a host-side HD-transform+projector shim (tracked separately). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/_configs/_extractors_test.py | 38 +++++ src/mobius/_configs/_sub_configs.py | 4 + src/mobius/_configs/per_model/_phi_vision.py | 19 ++- src/mobius/models/clip.py | 134 +++++++++++++-- src/mobius/models/clip_test.py | 96 +++++++++++ src/mobius/models/phi3_v.py | 161 ++++++++++--------- src/mobius/models/phi3_v_test.py | 133 +++++++++++++++ tests/vision_integration_test.py | 107 +++++++++++- 8 files changed, 599 insertions(+), 93 deletions(-) create mode 100644 src/mobius/models/clip_test.py create mode 100644 src/mobius/models/phi3_v_test.py diff --git a/src/mobius/_configs/_extractors_test.py b/src/mobius/_configs/_extractors_test.py index 897c280f..3d315946 100644 --- a/src/mobius/_configs/_extractors_test.py +++ b/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. diff --git a/src/mobius/_configs/_sub_configs.py b/src/mobius/_configs/_sub_configs.py index f872d9a7..7a8b1fdb 100644 --- a/src/mobius/_configs/_sub_configs.py +++ b/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 diff --git a/src/mobius/_configs/per_model/_phi_vision.py b/src/mobius/_configs/per_model/_phi_vision.py index 4bbb35d2..74a73e92 100644 --- a/src/mobius/_configs/per_model/_phi_vision.py +++ b/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 diff --git a/src/mobius/models/clip.py b/src/mobius/models/clip.py index e195add2..85304a1f 100644 --- a/src/mobius/models/clip.py +++ b/src/mobius/models/clip.py @@ -13,12 +13,40 @@ from mobius._configs import ArchitectureConfig from mobius.components import FCMLP from mobius.components._common import Embedding, LayerNorm -from mobius.components._conv import Conv2d +from mobius.components._conv import Conv2d, Conv2dNoBias from mobius.components._encoder import EncoderAttention if TYPE_CHECKING: import onnx_ir as ir +# Largest int64 value; used as an open-ended ``Slice`` end (ONNX clamps it to +# the actual dimension size). +_INT64_MAX = 9223372036854775807 + + +class ClipVisionConfigView: + """Adapter exposing a :class:`VisionConfig` under CLIP's field names. + + The CLIP vision modules read the encoder geometry from top-level attributes + (``hidden_size``, ``num_hidden_layers``, ``num_channels``, ``rms_norm_eps``, + ``hidden_act`` ...). In a multimodal checkpoint those live inside the nested + :class:`~mobius._configs._sub_configs.VisionConfig` (e.g. Phi-3.5-Vision's + ``img_processor`` dict). This view bridges the two so the CLIP tower can be + reused without duplicating its module definitions. + """ + + def __init__(self, vision_config, *, default_hidden_act: str = "quick_gelu"): + self._vision_config = vision_config + self.hidden_size = vision_config.hidden_size + self.intermediate_size = vision_config.intermediate_size + self.num_hidden_layers = vision_config.num_hidden_layers + self.num_attention_heads = vision_config.num_attention_heads + self.image_size = vision_config.image_size + self.patch_size = vision_config.patch_size + self.num_channels = vision_config.in_channels + self.rms_norm_eps = vision_config.norm_eps + self.hidden_act = vision_config.hidden_act or default_hidden_act + class _CLIPVisionEmbeddings(nn.Module): """CLIP vision embeddings: Conv2d patch + CLS token + position embeddings.""" @@ -30,11 +58,13 @@ def __init__(self, config: ArchitectureConfig): image_size = config.image_size self.class_embedding = nn.Parameter((hidden_size,)) + # CLIP's patch-embedding Conv2d has no bias (HuggingFace uses bias=False). self.patch_embedding = _Conv2dPatchEmbed( config.num_channels, hidden_size, kernel_size=patch_size, stride=patch_size, + bias=False, ) num_patches = (image_size // patch_size) ** 2 self.position_embedding = Embedding(num_patches + 1, hidden_size) @@ -58,7 +88,10 @@ def forward(self, op: OpBuilder, pixel_values: ir.Value): class _Conv2dPatchEmbed(nn.Module): - """Conv2d-based patch embedding with reshape and transpose.""" + """Conv2d-based patch embedding with reshape and transpose. + + ``bias`` selects a biased (SigLIP) or bias-free (CLIP) patch convolution. + """ def __init__( self, @@ -66,9 +99,11 @@ def __init__( out_channels: int, kernel_size: int, stride: int, + bias: bool = True, ): super().__init__() - self.projection = Conv2d(in_channels, out_channels, kernel_size, stride) + conv_cls = Conv2d if bias else Conv2dNoBias + self.projection = conv_cls(in_channels, out_channels, kernel_size, stride) def forward(self, op: OpBuilder, x: ir.Value): # Conv2d: [batch, channels, H, W] -> [batch, out_channels, H', W'] @@ -110,23 +145,88 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): return hidden_states +def resolve_clip_feature_num_layers(num_hidden_layers: int, feature_layer: int) -> int: + """Number of encoder layers to run to reach ``hidden_states[feature_layer]``. + + HuggingFace CLIP exposes ``num_hidden_layers + 1`` hidden states: index 0 is + the pre-encoder embedding output and index ``k`` is the output after ``k`` + encoder layers. A model such as Phi-3.5-Vision extracts features from an + intermediate layer (``layer_idx = -2``), which corresponds to running one + fewer encoder layer and skipping the final ``post_layernorm``. + + Args: + num_hidden_layers: Total number of CLIP encoder layers in the checkpoint. + feature_layer: The ``hidden_states`` index to extract (may be negative, + using Python's from-the-end convention over the ``N + 1`` states). + + Returns: + The number of encoder layers to actually instantiate and run. + + Raises: + ValueError: If ``feature_layer`` is out of range for the encoder depth. + """ + num_hidden_states = num_hidden_layers + 1 + index = feature_layer if feature_layer >= 0 else num_hidden_states + feature_layer + if not 0 <= index <= num_hidden_layers: + raise ValueError( + f"feature_layer={feature_layer} is out of range for a CLIP encoder " + f"with {num_hidden_layers} layers (valid hidden-state indices span " + f"0..{num_hidden_layers})." + ) + return index + + class CLIPVisionModel(nn.Module): """CLIP vision model for standalone image feature extraction. - Outputs last_hidden_state from the vision encoder. + By default this outputs the final ``last_hidden_state`` (all encoder layers + followed by ``post_layernorm``). + + Two options support multimodal feature extraction used by models such as + Phi-3.5-Vision: + + * ``feature_layer`` selects an intermediate ``hidden_states`` index + (HuggingFace convention, e.g. ``-2``). When set, only the encoder layers + needed to reach that hidden state are instantiated and run, and the final + ``post_layernorm`` is skipped (it is only applied to the last hidden + state in HuggingFace). + * ``drop_class_token`` removes the leading CLS token from the output so that + only the patch features remain (HuggingFace ``img_feature[:, 1:]``). """ default_task = "image-classification" category = "vision" - def __init__(self, config: ArchitectureConfig): + def __init__( + self, + config: ArchitectureConfig, + *, + feature_layer: int | None = None, + drop_class_token: bool = False, + ): super().__init__() + self.feature_layer = feature_layer + self.drop_class_token = drop_class_token + + if feature_layer is None: + num_encoder_layers = config.num_hidden_layers + else: + num_encoder_layers = resolve_clip_feature_num_layers( + config.num_hidden_layers, feature_layer + ) + self.embeddings = _CLIPVisionEmbeddings(config) self.encoder = nn.ModuleList( - [_CLIPVisionEncoderLayer(config) for _ in range(config.num_hidden_layers)] + [_CLIPVisionEncoderLayer(config) for _ in range(num_encoder_layers)] ) self.pre_layrnorm = LayerNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_layernorm = LayerNorm(config.hidden_size, eps=config.rms_norm_eps) + # post_layernorm is only applied to the final hidden state; when we + # extract an intermediate feature layer it is unused (and its weights + # are intentionally left unmapped). + if feature_layer is None: + self.post_layernorm = LayerNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.post_layernorm = None def forward(self, op: OpBuilder, pixel_values: ir.Value): hidden_states = self.embeddings(op, pixel_values) @@ -135,7 +235,19 @@ def forward(self, op: OpBuilder, pixel_values: ir.Value): for layer in self.encoder: hidden_states = layer(op, hidden_states) - hidden_states = self.post_layernorm(op, hidden_states) + if self.post_layernorm is not None: + hidden_states = self.post_layernorm(op, hidden_states) + + if self.drop_class_token: + # Keep patch tokens only, dropping the leading CLS token: + # (batch, 1 + num_patches, hidden) -> (batch, num_patches, hidden). + hidden_states = op.Slice( + hidden_states, + [1], # starts + [_INT64_MAX], # ends (clamped to sequence length) + [1], # axes: sequence dimension + ) + return hidden_states def preprocess_weights( @@ -175,8 +287,12 @@ def _rename_clip_vision_weight(name: str) -> str | None: ): return None - # Embeddings + # Embeddings — HF stores a flat ``patch_embedding.{weight,bias}`` Conv, but + # our ``_Conv2dPatchEmbed`` wraps the Conv in a ``.projection`` sub-module. if name.startswith("embeddings."): + name = name.replace( + "embeddings.patch_embedding.", "embeddings.patch_embedding.projection." + ) return name # Pre/post layer norm diff --git a/src/mobius/models/clip_test.py b/src/mobius/models/clip_test.py new file mode 100644 index 00000000..d0fd0055 --- /dev/null +++ b/src/mobius/models/clip_test.py @@ -0,0 +1,96 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Unit tests for CLIP vision weight rename and feature-layer selection. + +Covers the pieces exercised by CLIP-based multimodal encoders such as +Phi-3.5-Vision: the HuggingFace → ONNX weight name mapping (including the +bias-free patch convolution wrapping) and the intermediate feature-layer +math. No HF model downloads needed. +""" + +from __future__ import annotations + +import pytest + +from mobius.models.clip import ( + _rename_clip_vision_weight, + resolve_clip_feature_num_layers, +) + + +class TestRenameClipVisionWeight: + """Test _rename_clip_vision_weight for HF CLIP vision weight patterns.""" + + def test_class_embedding(self): + assert ( + _rename_clip_vision_weight("vision_model.embeddings.class_embedding") + == "embeddings.class_embedding" + ) + + def test_position_embedding(self): + assert ( + _rename_clip_vision_weight("vision_model.embeddings.position_embedding.weight") + == "embeddings.position_embedding.weight" + ) + + def test_patch_embedding_wraps_projection(self): + # HF stores a flat patch_embedding.weight; our module nests it under + # a .projection Conv sub-module. + assert ( + _rename_clip_vision_weight("vision_model.embeddings.patch_embedding.weight") + == "embeddings.patch_embedding.projection.weight" + ) + + def test_pre_and_post_layernorm(self): + assert ( + _rename_clip_vision_weight("vision_model.pre_layrnorm.weight") + == "pre_layrnorm.weight" + ) + assert ( + _rename_clip_vision_weight("vision_model.post_layernorm.bias") + == "post_layernorm.bias" + ) + + def test_encoder_attention_projection(self): + assert ( + _rename_clip_vision_weight("vision_model.encoder.layers.5.self_attn.q_proj.weight") + == "encoder.5.self_attn.q_proj.weight" + ) + + def test_encoder_mlp_renamed(self): + assert ( + _rename_clip_vision_weight("vision_model.encoder.layers.0.mlp.fc1.weight") + == "encoder.0.mlp.up_proj.weight" + ) + assert ( + _rename_clip_vision_weight("vision_model.encoder.layers.0.mlp.fc2.bias") + == "encoder.0.mlp.down_proj.bias" + ) + + def test_text_weights_dropped(self): + assert _rename_clip_vision_weight("text_model.encoder.layers.0.mlp.fc1.weight") is None + assert _rename_clip_vision_weight("visual_projection.weight") is None + + +class TestResolveClipFeatureNumLayers: + """Test resolve_clip_feature_num_layers (HF hidden_states indexing).""" + + def test_negative_two_skips_last_layer(self): + # 24 layers -> 25 hidden states; index -2 == state 23 == run 23 layers. + assert resolve_clip_feature_num_layers(24, -2) == 23 + + def test_negative_one_is_last_layer(self): + assert resolve_clip_feature_num_layers(24, -1) == 24 + + def test_positive_index(self): + assert resolve_clip_feature_num_layers(24, 10) == 10 + + def test_zero_is_pre_encoder(self): + assert resolve_clip_feature_num_layers(24, 0) == 0 + + def test_out_of_range_raises(self): + with pytest.raises(ValueError): + resolve_clip_feature_num_layers(24, -30) + with pytest.raises(ValueError): + resolve_clip_feature_num_layers(24, 25) diff --git a/src/mobius/models/phi3_v.py b/src/mobius/models/phi3_v.py index 1cd24af2..836b4f0a 100644 --- a/src/mobius/models/phi3_v.py +++ b/src/mobius/models/phi3_v.py @@ -11,19 +11,22 @@ Architecture: pixel_values (num_crops, 3, 336, 336) - → CLIP ViT-L/14-336 → (num_crops, 576, 1024) - → MLP projector → (num_crops, 576, 3072) + → CLIP ViT-L/14-336 (hidden_states[-2], CLS dropped) → (num_crops, 576, 1024) + → [host-side HD transform: 2x2 merge + sub_GN/glb_GN + MLP projector] → concat + insert into input sequence → Phi3 text decoder The vision encoder uses CLIP ViT-L/14-336 (patch_size=14, image_size=336), -which produces 576 patches per crop. The MLP projector maps from the vision -hidden size (1024) to the text hidden size (3072). +which produces 576 patches per crop (after dropping the CLS token). Features +are taken from ``hidden_states[layer_idx]`` (default ``-2``): the last encoder +layer and ``post_layernorm`` are skipped, matching HuggingFace's +``get_img_features``. HD transform note: the HuggingFace phi3_v model supports high-definition -tiling (sub_glb ordering with learnable separators). This ONNX implementation -processes a batch of pre-tiled crops uniformly — the HD tiling and separator -insertion is expected to happen in the preprocessing step outside ONNX. +tiling (sub_glb ordering with learnable separators) and the ``img_projection`` +MLP whose input width is ``image_dim_out * 4``. Both are image-size dependent +and are expected to run host-side, outside ONNX — the ONNX vision encoder +emits the raw per-crop CLIP patch features that feed into them. HuggingFace weight name prefixes:: @@ -35,11 +38,10 @@ model.norm.* / lm_head.* → decoder.model.norm.* / decoder.lm_head.* model.vision_embed_tokens.img_processor.vision_model.* - → vision_encoder.vision_tower.vision_model.* - model.vision_embed_tokens.img_projection.0.* (fc1) - → vision_encoder.multi_modal_projector.fc1.* - model.vision_embed_tokens.img_projection.2.* (fc2) - → vision_encoder.multi_modal_projector.fc2.* + → vision_encoder.vision_tower.* (CLIP tower only) + + (model.vision_embed_tokens.img_projection.* and sub_GN/glb_GN are + host-side and intentionally not exported.) Reference: ``microsoft/Phi-3-vision-128k-instruct``, ``microsoft/Phi-3.5-vision-instruct``. @@ -58,10 +60,13 @@ from mobius.components import ( Embedding, Linear, - MLPMultiModalProjector, - VisionModel, ) from mobius.models.base import TextModel +from mobius.models.clip import ( + ClipVisionConfigView, + CLIPVisionModel, + _rename_clip_vision_weight, +) if TYPE_CHECKING: import onnx_ir as ir @@ -140,66 +145,82 @@ def preprocess_weights( return renamed -class _Phi3VVisionEncoderModel(nn.Module): - """Phi3-V vision encoder: CLIP ViT-L/14-336 + MLP projector. +# Prefix under which the CLIP ViT weights live in a Phi-3-V checkpoint. +_VISION_TOWER_PREFIX = "model.vision_embed_tokens.img_processor." + + +def _rename_phi3v_vision_weight(name: str) -> str | None: + """Map a Phi-3-V ``img_processor`` CLIP weight to ``vision_tower.*``. - Accepts a batch of image crops (pre-tiled), applies the CLIP encoder to - each, and projects from the vision hidden size to the text hidden size. + Strips the ``model.vision_embed_tokens.img_processor.`` prefix, then reuses + the shared CLIP vision renamer (which handles the ``vision_model.`` prefix, + the ``patch_embedding`` → ``patch_embedding.projection`` Conv wrapping, the + ``mlp.fc1/fc2`` → ``mlp.up_proj/down_proj`` naming, and the per-layer + ``encoder.layers.N`` → ``encoder.N`` flattening). - Input shape: (num_crops, 3, 336, 336) — NCHW - Output shape: (num_crops * num_patches, text_hidden_size) + Returns ``None`` for weights that are not part of the CLIP tower (e.g. the + ``img_projection`` MLP and the learnable ``sub_GN``/``glb_GN`` separators, + which are applied host-side as part of the HD feature transform). + """ + if not name.startswith(_VISION_TOWER_PREFIX): + return None + clip_name = _rename_clip_vision_weight(name[len(_VISION_TOWER_PREFIX) :]) + if clip_name is None: + return None + return "vision_tower." + clip_name + + +class _Phi3VVisionEncoderModel(nn.Module): + """Phi3-V vision encoder: CLIP ViT-L/14-336 patch-feature extractor. + + Faithfully reproduces HuggingFace ``Phi3ImageEmbedding.get_img_features``: + run the CLIP tower to ``hidden_states[layer_idx]`` (default ``-2``, i.e. all + but the last encoder layer and no ``post_layernorm``) and keep only the + patch tokens, dropping the leading CLS token. + + The ``img_projection`` MLP and the HD 2x2 patch-merge with learnable + ``sub_GN``/``glb_GN`` separators are intentionally *not* part of this ONNX + graph: the projector's input width (``image_dim_out * 4``) only exists after + the host-side HD feature transform, which is image-size dependent and cannot + be expressed as a static graph. They run host-side alongside the pre-tiling + of crops, consistent with the model's HD-transform design. + + Input shape: (num_crops, 3, image_size, image_size) — NCHW + Output shape: (num_crops, num_patches, image_dim_out) — raw patch features """ def __init__(self, config: ArchitectureConfig): super().__init__() - self.vision_tower = VisionModel(config) - self.multi_modal_projector = MLPMultiModalProjector( - vision_hidden_size=config.vision.hidden_size, - text_hidden_size=config.hidden_size, + assert config.vision is not None, "Phi3-V requires a vision config" + clip_config = ClipVisionConfigView(config.vision) + self.vision_tower = CLIPVisionModel( + clip_config, + feature_layer=config.vision.feature_layer, + drop_class_token=True, ) def forward(self, op: builder.OpBuilder, pixel_values: ir.Value): - # pixel_values: (num_crops, 3, H, W) - vision_features = self.vision_tower(op, pixel_values) - # vision_features: (num_crops, num_patches, vision_hidden) - # Flatten crops and patches into a single token sequence - num_crops = op.Shape(pixel_values, start=0, end=1) # scalar - num_patches = op.Shape(vision_features, start=1, end=2) - total = op.Mul(num_crops, num_patches) - vision_dim = op.Shape(vision_features, start=2, end=3) - flat_shape = op.Concat(total, vision_dim, axis=0) - vision_features = op.Reshape(vision_features, flat_shape) - # vision_features: (num_crops * num_patches, vision_hidden) - return self.multi_modal_projector(op, vision_features) + # pixel_values: (num_crops, 3, H, W) -> (num_crops, num_patches, image_dim_out) + return self.vision_tower(op, pixel_values) def preprocess_weights( self, state_dict: dict[str, torch.Tensor] ) -> dict[str, torch.Tensor]: - """Extract vision encoder and projector weights from the full checkpoint. + """Extract CLIP vision-tower weights from the full checkpoint. HF path → ONNX path:: model.vision_embed_tokens.img_processor.vision_model.* - → vision_tower.vision_model.* - model.vision_embed_tokens.img_projection.0.* - → multi_modal_projector.fc1.* - model.vision_embed_tokens.img_projection.2.* - → multi_modal_projector.fc2.* + → vision_tower.* (via the shared CLIP renamer) + + The ``img_projection`` and ``sub_GN``/``glb_GN`` tensors are dropped + (host-side HD transform). """ renamed: dict[str, torch.Tensor] = {} - - vision_pfx = "model.vision_embed_tokens.img_processor." - proj_0_pfx = "model.vision_embed_tokens.img_projection.0." - proj_2_pfx = "model.vision_embed_tokens.img_projection.2." - for key, value in state_dict.items(): - if key.startswith(vision_pfx): - renamed["vision_tower." + key[len(vision_pfx) :]] = value - elif key.startswith(proj_0_pfx): - renamed["multi_modal_projector.fc1." + key[len(proj_0_pfx) :]] = value - elif key.startswith(proj_2_pfx): - renamed["multi_modal_projector.fc2." + key[len(proj_2_pfx) :]] = value - + new_key = _rename_phi3v_vision_weight(key) + if new_key is not None: + renamed[new_key] = value return renamed @@ -295,32 +316,22 @@ def preprocess_weights( → decoder.model.layers.N.* (split) model.norm.* / lm_head.* → decoder.model.norm.* / decoder.lm_head.* - model.vision_embed_tokens.img_processor.* - → vision_encoder.vision_tower.* - model.vision_embed_tokens.img_projection.0.* - → vision_encoder.multi_modal_projector.fc1.* - model.vision_embed_tokens.img_projection.2.* - → vision_encoder.multi_modal_projector.fc2.* + model.vision_embed_tokens.img_processor.vision_model.* + → vision_encoder.vision_tower.* (CLIP tower only) + + The ``img_projection`` MLP and the learnable ``sub_GN``/``glb_GN`` + separators are dropped — they are applied host-side as part of the HD + feature transform (see :class:`_Phi3VVisionEncoderModel`). """ renamed: dict[str, torch.Tensor] = {} - vision_pfx = "model.vision_embed_tokens.img_processor." - proj_0_pfx = "model.vision_embed_tokens.img_projection.0." - proj_2_pfx = "model.vision_embed_tokens.img_projection.2." - for key, value in state_dict.items(): - if key.startswith(vision_pfx): - new_key = "vision_encoder.vision_tower." + key[len(vision_pfx) :] - renamed[new_key] = value - elif key.startswith(proj_0_pfx): - sfx = key[len(proj_0_pfx) :] - renamed[f"vision_encoder.multi_modal_projector.fc1.{sfx}"] = value - elif key.startswith(proj_2_pfx): - sfx = key[len(proj_2_pfx) :] - renamed[f"vision_encoder.multi_modal_projector.fc2.{sfx}"] = value + 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."): - # Learnable separators (sub_GN, glb_GN) and other VE state - # are not used in the ONNX decoder sub-model — skip them. + # img_projection / sub_GN / glb_GN and any other vision-embed + # state is host-side (HD transform) — skip it. pass elif key == "model.embed_tokens.weight": renamed["decoder.model.embed_tokens.weight"] = value diff --git a/src/mobius/models/phi3_v_test.py b/src/mobius/models/phi3_v_test.py new file mode 100644 index 00000000..16a6595d --- /dev/null +++ b/src/mobius/models/phi3_v_test.py @@ -0,0 +1,133 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Unit tests for Phi-3/3.5-Vision (phi3_v) vision-tower weight mapping. + +Verifies that the non-standard ``img_processor``-nested CLIP weights are +mapped onto the CLIP vision-tower initializer names, and that the host-side +img_projection / learnable-separator tensors are dropped. No HF downloads. +""" + +from __future__ import annotations + +import torch + +from mobius.models.phi3_v import ( + _Phi3VVisionEncoderModel, + _rename_phi3v_vision_weight, +) + + +class TestRenamePhi3VVisionWeight: + """Test _rename_phi3v_vision_weight for the img_processor CLIP prefix.""" + + _PFX = "model.vision_embed_tokens.img_processor.vision_model." + + def test_class_embedding(self): + assert ( + _rename_phi3v_vision_weight(self._PFX + "embeddings.class_embedding") + == "vision_tower.embeddings.class_embedding" + ) + + def test_patch_embedding_wraps_projection(self): + assert ( + _rename_phi3v_vision_weight(self._PFX + "embeddings.patch_embedding.weight") + == "vision_tower.embeddings.patch_embedding.projection.weight" + ) + + def test_position_embedding(self): + assert ( + _rename_phi3v_vision_weight(self._PFX + "embeddings.position_embedding.weight") + == "vision_tower.embeddings.position_embedding.weight" + ) + + def test_pre_layrnorm(self): + assert ( + _rename_phi3v_vision_weight(self._PFX + "pre_layrnorm.weight") + == "vision_tower.pre_layrnorm.weight" + ) + + def test_encoder_layer_mlp(self): + assert ( + _rename_phi3v_vision_weight(self._PFX + "encoder.layers.7.mlp.fc1.weight") + == "vision_tower.encoder.7.mlp.up_proj.weight" + ) + + def test_img_projection_dropped(self): + # Projector is applied host-side (HD transform) — not exported. + assert ( + _rename_phi3v_vision_weight("model.vision_embed_tokens.img_projection.0.weight") + is None + ) + assert ( + _rename_phi3v_vision_weight("model.vision_embed_tokens.img_projection.2.weight") + is None + ) + + def test_learnable_separators_dropped(self): + assert _rename_phi3v_vision_weight("model.vision_embed_tokens.sub_GN") is None + assert _rename_phi3v_vision_weight("model.vision_embed_tokens.glb_GN") is None + + def test_text_decoder_weights_dropped(self): + assert _rename_phi3v_vision_weight("model.layers.0.self_attn.qkv_proj.weight") is None + assert _rename_phi3v_vision_weight("lm_head.weight") is None + + +class TestPhi3VVisionPreprocessWeights: + """The vision encoder maps a realistic checkpoint slice with no leftovers.""" + + def _config(self, feature_layer): + from mobius._configs import ArchitectureConfig + from mobius._configs._sub_configs import VisionConfig + + vision = VisionConfig( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=3, + num_attention_heads=2, + image_size=28, + patch_size=14, + norm_eps=1e-5, + hidden_act="quick_gelu", + feature_layer=feature_layer, + ) + return ArchitectureConfig( + hidden_size=16, + intermediate_size=32, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=8, + num_hidden_layers=1, + vocab_size=32, + max_position_embeddings=64, + rms_norm_eps=1e-5, + rope_type="default", + rope_theta=10_000.0, + pad_token_id=0, + vision=vision, + image_token_id=32044, + ) + + def test_feature_layer_maps_only_needed_layers(self): + pfx = "model.vision_embed_tokens.img_processor.vision_model." + state_dict = { + pfx + "embeddings.class_embedding": torch.zeros(8), + pfx + "embeddings.patch_embedding.weight": torch.zeros(8, 3, 14, 14), + pfx + "embeddings.position_embedding.weight": torch.zeros(5, 8), + pfx + "pre_layrnorm.weight": torch.zeros(8), + pfx + "encoder.layers.0.mlp.fc1.weight": torch.zeros(16, 8), + pfx + "encoder.layers.2.mlp.fc1.weight": torch.zeros(16, 8), + pfx + "post_layernorm.weight": torch.zeros(8), + "model.vision_embed_tokens.img_projection.0.weight": torch.zeros(16, 32), + "model.vision_embed_tokens.sub_GN": torch.zeros(1, 1, 1, 32), + } + module = _Phi3VVisionEncoderModel(self._config(feature_layer=-2)) + renamed = module.preprocess_weights(state_dict) + + 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 + # 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. + assert all(k.startswith("vision_tower.") for k in renamed) diff --git a/tests/vision_integration_test.py b/tests/vision_integration_test.py index c6453aaa..5042059a 100644 --- a/tests/vision_integration_test.py +++ b/tests/vision_integration_test.py @@ -143,7 +143,7 @@ class _TorchCLIPVisionModel(nn.Module): post_layernorm.{weight,bias} """ - def __init__(self): + def __init__(self, num_layers: int = _LAYERS): super().__init__() num_patches = (_IMAGE_SIZE // _PATCH_SIZE) ** 2 # Embeddings @@ -151,16 +151,31 @@ def __init__(self): self.embeddings.class_embedding = nn.Parameter(torch.zeros(_HIDDEN)) self.embeddings.patch_embedding = nn.Module() self.embeddings.patch_embedding.projection = nn.Conv2d( - _CHANNELS, _HIDDEN, kernel_size=_PATCH_SIZE, stride=_PATCH_SIZE + _CHANNELS, _HIDDEN, kernel_size=_PATCH_SIZE, stride=_PATCH_SIZE, bias=False ) self.embeddings.position_embedding = nn.Embedding(num_patches + 1, _HIDDEN) # Pre/post norms self.pre_layrnorm = nn.LayerNorm(_HIDDEN, eps=_NORM_EPS) self.post_layernorm = nn.LayerNorm(_HIDDEN, eps=_NORM_EPS) # Encoder as indexed ModuleList (keys: encoder.0.*, encoder.1.*, ...) - self.encoder = nn.ModuleList([_TorchCLIPEncoderLayer() for _ in range(_LAYERS)]) + self.encoder = nn.ModuleList([_TorchCLIPEncoderLayer() for _ in range(num_layers)]) def forward(self, pixel_values): + return self._encode(pixel_values, feature_layer=None, drop_class_token=False) + + def forward_features(self, pixel_values, feature_layer, drop_class_token): + """Match CLIPVisionModel(feature_layer=..., drop_class_token=...). + + Returns ``hidden_states[feature_layer]`` (HuggingFace convention, no + ``post_layernorm``), optionally with the leading CLS token dropped. + """ + return self._encode( + pixel_values, + feature_layer=feature_layer, + drop_class_token=drop_class_token, + ) + + def _encode(self, pixel_values, feature_layer, drop_class_token): # Patch embed x = self.embeddings.patch_embedding.projection(pixel_values) x = x.flatten(2).transpose(1, 2) # (B, num_patches, hidden) @@ -174,10 +189,18 @@ def forward(self, pixel_values): x = x + self.embeddings.position_embedding(position_ids) # Pre-norm x = self.pre_layrnorm(x) - # Encoder + # Encoder — collect HuggingFace-style hidden states (index 0 == input). + hidden_states = [x] for layer in self.encoder: x = layer(x) - return self.post_layernorm(x) + hidden_states.append(x) + if feature_layer is None: + out = self.post_layernorm(x) + else: + out = hidden_states[feature_layer] + if drop_class_token: + out = out[:, 1:] + return out class _TorchCLIPEncoderLayer(nn.Module): @@ -353,3 +376,77 @@ def test_clip_vision_hidden_states_match(self): rtol=1e-3, atol=1e-3, ) + + +@pytest.mark.integration +@pytest.mark.integration_fast +class TestCLIPVisionFeatureLayerParity: + """CLIP feature extraction (Phi-3.5-Vision style): intermediate layer + CLS drop.""" + + def test_feature_layer_and_cls_drop_match(self): + """hidden_states[-2] with the CLS token dropped matches PyTorch.""" + import onnx_ir as ir + + num_layers = 3 + feature_layer = -2 + config = ArchitectureConfig( + hidden_size=_HIDDEN, + intermediate_size=_INTERMEDIATE, + num_attention_heads=_HEADS, + num_key_value_heads=_HEADS, + head_dim=_HIDDEN // _HEADS, + num_hidden_layers=num_layers, + vocab_size=256, + max_position_embeddings=128, + hidden_act="quick_gelu", + rms_norm_eps=_NORM_EPS, + rope_type="default", + rope_theta=10_000.0, + pad_token_id=0, + image_size=_IMAGE_SIZE, + patch_size=_PATCH_SIZE, + num_channels=_CHANNELS, + ) + config.dtype = ir.DataType.FLOAT + + onnx_module = models.CLIPVisionModel( + config, feature_layer=feature_layer, drop_class_token=True + ) + pkg = build_from_module(onnx_module, config, task="image-classification") + + # Only the layers needed to reach hidden_states[-2] are instantiated, + # and post_layernorm is intentionally absent. + init_names = set(pkg["model"].graph.initializers) + assert not any("post_layernorm" in n for n in init_names) + assert not any("encoder.2." in n for n in init_names) # last layer skipped + + ref_model = _TorchCLIPVisionModel(num_layers=num_layers).float().eval() + # Only the weights that ended up in the graph are applied; extras (the + # skipped last layer + post_layernorm) are simply not present. + apply_weights(pkg["model"], dict(ref_model.state_dict())) + + rng = np.random.default_rng(7) + pixel_values = rng.standard_normal((2, _CHANNELS, _IMAGE_SIZE, _IMAGE_SIZE)).astype( + np.float32 + ) + with torch.no_grad(): + ref_out = ref_model.forward_features( + torch.from_numpy(pixel_values), + feature_layer=feature_layer, + drop_class_token=True, + ).numpy() + + num_patches = (_IMAGE_SIZE // _PATCH_SIZE) ** 2 + assert ref_out.shape == (2, num_patches, _HIDDEN) # CLS dropped + + session = OnnxModelSession(pkg["model"]) + onnx_out = session.run({"pixel_values": pixel_values}) + session.close() + + assert onnx_out["last_hidden_state"].shape == (2, num_patches, _HIDDEN) + assert_logits_close( + onnx_out["last_hidden_state"], + ref_out, + rtol=1e-3, + atol=1e-3, + ) From f111938786590b561c64d2327089f3f2105001f8 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 25 Jul 2026 16:24:12 +0000 Subject: [PATCH 03/10] phi3_v: add host-side HD-transform/projector shim; un-skip e2e golden (L4) Implement the real Phi-3.5-Vision host-side HD feature transform + img_projection in src/mobius/models/_phi3_vision_projector.py (pure NumPy, mirrors modeling_phi3_v.py get_img_features -> hd_feature_transform -> reshape_hd_patches_2x2merge -> add_image_newline -> img_projection): 2x2 spatial patch merge (N,576,1024)->(N,144,4096), learnable sub_GN/glb_GN separator insertion in the fixed sub_glb order, then Linear(4096->3072)->GELU->Linear. load_phi3_vision_projector_weights reads the projector + separator tensors (dropped from the ONNX graphs) directly from the checkpoint safetensors. Wire it into tests/e2e_golden_test.py via _run_vl_vision_to_image_features (flattens HF 5-D pixel_values, runs the ONNX vision encoder, applies the host-side transform) for both L4 prefill and L5 generation. Fix _build_mm_prompt to use Phi3V's inner-tokenizer chat template + img_tokens (<|image_1|>) placeholders. Un-skip testdata/cases/vision-language/phi3_5-vision-instruct.yaml and add the golden JSON regenerated from the real checkpoint + cat image. L4 golden passes (top-1 argmax matches HF over all 775 tokens); the host-side projector matches HF hd_feature_transform at cos=1.0000, max-abs 6.9e-5. L5 greedy generation is xfail (strict=False) with a precise reason: it drifts from HF after the first few tokens due to the same decoder-side LongRoPE parity gap already xfailed for phi4mm-multi-image-audio (shared Phi-3.5 decoder), not a vision/projector defect. Make load_torch_multimodal_model tolerant of older transformers (no AutoModelForImageTextToText; torch_dtype vs dtype) so the golden can be generated offline under transformers 4.43 (the bundled 4.x remote code cannot generate on transformers 5.x). Align a stale su-rope assertion in _config_resolver_test.py to Bishop's deliberate su->longrope canonicalization. Add 11 NumPy unit tests for the projector (2x2 merge ordering, separator insertion, GELU-MLP math, sub_glb token count/ordering, weight loader + error paths). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/_config_resolver_test.py | 4 +- src/mobius/_testing/torch_reference.py | 31 +- src/mobius/models/_phi3_vision_projector.py | 293 +++++++ .../models/_phi3_vision_projector_test.py | 223 +++++ .../phi3_5-vision-instruct.yaml | 3 +- .../phi3_5-vision-instruct.json | 811 ++++++++++++++++++ .../phi3_5-vision-instruct_generation.json | 37 + tests/e2e_golden_test.py | 157 +++- 8 files changed, 1506 insertions(+), 53 deletions(-) create mode 100644 src/mobius/models/_phi3_vision_projector.py create mode 100644 src/mobius/models/_phi3_vision_projector_test.py create mode 100644 testdata/golden/vision-language/phi3_5-vision-instruct.json create mode 100644 testdata/golden/vision-language/phi3_5-vision-instruct_generation.json diff --git a/src/mobius/_config_resolver_test.py b/src/mobius/_config_resolver_test.py index 7e1b3cf4..b7a0a694 100644 --- a/src/mobius/_config_resolver_test.py +++ b/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) diff --git a/src/mobius/_testing/torch_reference.py b/src/mobius/_testing/torch_reference.py index c1fb8479..9ab54f33 100644 --- a/src/mobius/_testing/torch_reference.py +++ b/src/mobius/_testing/torch_reference.py @@ -247,13 +247,32 @@ def load_torch_multimodal_model( # 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() diff --git a/src/mobius/models/_phi3_vision_projector.py b/src/mobius/models/_phi3_vision_projector.py new file mode 100644 index 00000000..a40662f8 --- /dev/null +++ b/src/mobius/models/_phi3_vision_projector.py @@ -0,0 +1,293 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Host-side HD feature transform and projector for Phi-3.5-Vision. + +The mobius ``vision_encoder`` ONNX graph for ``phi3_v`` emits the raw CLIP +patch features ``(num_crops, num_patches, image_dim_out)`` — exactly the output +of HuggingFace ``Phi3ImageEmbedding.get_img_features``. Everything that follows +(the 2x2 spatial "HD" patch merge, the learnable ``sub_GN``/``glb_GN`` line +separators, and the ``img_projection`` MLP) is image-size dependent: the crop +grid, the separator interleaving, and the final token count all vary with the +input image resolution, so they cannot be expressed as a static ONNX graph. + +This module reproduces that host-side tail faithfully in NumPy, mirroring +``modeling_phi3_v.py`` (``hd_feature_transform`` → ``reshape_hd_patches_2x2merge`` +→ ``add_image_newline`` → ``img_projection``). The output is the flat sequence +of projected image embeddings ``(total_image_tokens, hidden_size)`` that the +mobius ``embedding`` ONNX model scatters into the ``<|image|>`` token positions. + +The projector weights (``img_projection`` MLP and the ``sub_GN``/``glb_GN`` +separators) are intentionally dropped from the ONNX graphs (see +``phi3_v._rename_phi3v_vision_weight``); :func:`load_phi3_vision_projector_weights` +loads them directly from the checkpoint so the transform can run host-side. +""" + +from __future__ import annotations + +import dataclasses +import glob +import os + +import numpy as np + +# HF checkpoint key prefix for the projector / separators (outside the CLIP +# ``img_processor`` vision tower, which is exported to ONNX separately). +_PROJECTOR_PREFIX = "model.vision_embed_tokens." + +_GLOBAL_SEPARATOR_KEY = _PROJECTOR_PREFIX + "glb_GN" +_SUBLAYER_SEPARATOR_KEY = _PROJECTOR_PREFIX + "sub_GN" +_PROJECTION_FIRST_WEIGHT_KEY = _PROJECTOR_PREFIX + "img_projection.0.weight" +_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" + +# CLIP ViT-L/14-336 produces a 24x24 patch grid (576 patches). +_PATCH_GRID_SIDE = 24 +# Each crop is a 336x336 tile. +_CROP_PIXEL_SIZE = 336 + + +@dataclasses.dataclass(frozen=True) +class Phi3VisionProjectorWeights: + """Host-side projector weights for the Phi-3.5-Vision HD feature transform. + + All arrays are stored in ``float32`` for numerically faithful comparison + against the HuggingFace reference. + + Attributes: + global_separator: ``glb_GN`` separator, shape ``(1, 1, image_dim_out * 4)``. + sublayer_separator: ``sub_GN`` separator, shape ``(1, 1, 1, image_dim_out * 4)``. + projection_first_weight: First ``img_projection`` Linear weight, + shape ``(hidden_size, image_dim_out * 4)``. + projection_first_bias: First Linear bias, shape ``(hidden_size,)``. + projection_second_weight: Second Linear weight, + shape ``(hidden_size, hidden_size)``. + projection_second_bias: Second Linear bias, shape ``(hidden_size,)``. + """ + + global_separator: np.ndarray + sublayer_separator: np.ndarray + projection_first_weight: np.ndarray + projection_first_bias: np.ndarray + projection_second_weight: np.ndarray + projection_second_bias: np.ndarray + + +def _resolve_checkpoint_directory(model_id_or_directory: str) -> str: + """Return a local directory that contains the checkpoint safetensors. + + 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 + + from huggingface_hub import snapshot_download + + return snapshot_download(model_id_or_directory, local_files_only=True) + + +def load_phi3_vision_projector_weights( + model_id_or_directory: str, +) -> Phi3VisionProjectorWeights: + """Load the Phi-3.5-Vision projector + separator weights from a checkpoint. + + These tensors are omitted from the exported ONNX graphs (they belong to the + host-side HD feature transform), so they are read directly from the + checkpoint's safetensors shards here. + + Args: + model_id_or_directory: Local checkpoint directory or a HuggingFace hub + model id whose weights are already present in the local cache. + + Returns: + The projector weights as ``float32`` NumPy arrays. + + Raises: + FileNotFoundError: If no safetensors shards are found. + KeyError: If an expected projector tensor is missing. + """ + from safetensors import safe_open + + directory = _resolve_checkpoint_directory(model_id_or_directory) + shard_paths = sorted(glob.glob(os.path.join(directory, "*.safetensors"))) + if not shard_paths: + raise FileNotFoundError( + f"No .safetensors shards found in checkpoint directory: {directory}" + ) + + wanted_keys = { + _GLOBAL_SEPARATOR_KEY, + _SUBLAYER_SEPARATOR_KEY, + _PROJECTION_FIRST_WEIGHT_KEY, + _PROJECTION_FIRST_BIAS_KEY, + _PROJECTION_SECOND_WEIGHT_KEY, + _PROJECTION_SECOND_BIAS_KEY, + } + collected: dict[str, np.ndarray] = {} + for shard_path in shard_paths: + with safe_open(shard_path, framework="numpy") as shard: + shard_keys = set(shard.keys()) + for key in wanted_keys & shard_keys: + collected[key] = shard.get_tensor(key).astype(np.float32) + + missing = wanted_keys - collected.keys() + if missing: + raise KeyError( + f"Missing Phi-3.5-Vision projector weights in {directory}: {sorted(missing)}" + ) + + return Phi3VisionProjectorWeights( + global_separator=collected[_GLOBAL_SEPARATOR_KEY], + sublayer_separator=collected[_SUBLAYER_SEPARATOR_KEY], + projection_first_weight=collected[_PROJECTION_FIRST_WEIGHT_KEY], + projection_first_bias=collected[_PROJECTION_FIRST_BIAS_KEY], + projection_second_weight=collected[_PROJECTION_SECOND_WEIGHT_KEY], + projection_second_bias=collected[_PROJECTION_SECOND_BIAS_KEY], + ) + + +def _reshape_hd_patches_2x2_merge( + patch_features: np.ndarray, height_crops: int, width_crops: int +) -> np.ndarray: + """Merge each 2x2 block of patches into one channel-concatenated token. + + Mirrors ``modeling_phi3_v.reshape_hd_patches_2x2merge``. + + Args: + patch_features: ``(num_images * num_crops, 576, image_dim_out)``. + height_crops: Number of crop tiles stacked vertically. + width_crops: Number of crop tiles stacked horizontally. + + Returns: + ``(num_images, height_crops * 12, width_crops * 12, image_dim_out * 4)``. + """ + total_crops, patch_count, channels = patch_features.shape + if patch_count != _PATCH_GRID_SIDE * _PATCH_GRID_SIDE: + raise ValueError( + f"Expected {_PATCH_GRID_SIDE * _PATCH_GRID_SIDE} patches, got {patch_count}" + ) + if total_crops % (height_crops * width_crops) != 0: + raise ValueError( + f"Crop count {total_crops} is not divisible by {height_crops} * {width_crops}" + ) + num_images = total_crops // (height_crops * width_crops) + grid_side = _PATCH_GRID_SIDE + half_side = grid_side // 2 + merged = ( + patch_features.reshape(total_crops, grid_side, grid_side, channels) + .reshape(total_crops, half_side, 2, half_side, 2, channels) + .transpose(0, 1, 3, 2, 4, 5) + .reshape(total_crops, half_side * half_side, 4 * channels) + .reshape(num_images, height_crops, width_crops, half_side, half_side, 4 * channels) + .transpose(0, 1, 3, 2, 4, 5) + .reshape(num_images, height_crops * half_side, width_crops * half_side, 4 * channels) + ) + return merged + + +def _add_image_newline( + merged_features: np.ndarray, sublayer_separator: np.ndarray +) -> np.ndarray: + """Append the learnable ``sub_GN`` separator to each patch row. + + Mirrors ``modeling_phi3_v.add_image_newline``. + + Args: + merged_features: ``(num_images, height, width, hidden_dim)``. + sublayer_separator: ``sub_GN`` of shape ``(1, 1, 1, hidden_dim)``. + + Returns: + ``(num_images, height * (width + 1), hidden_dim)``. + """ + num_images, height, _width, hidden_dim = merged_features.shape + separator_column = np.broadcast_to(sublayer_separator, (num_images, height, 1, hidden_dim)) + with_newline = np.concatenate([merged_features, separator_column], axis=2) + return with_newline.reshape(num_images, -1, hidden_dim) + + +def _apply_image_projection( + tokens: np.ndarray, weights: Phi3VisionProjectorWeights +) -> np.ndarray: + """Apply the ``img_projection`` MLP: Linear -> GELU -> Linear. + + Uses the exact (erf-based) GELU to match ``nn.GELU()``. + + Args: + tokens: ``(num_tokens, image_dim_out * 4)``. + weights: Projector weights. + + Returns: + ``(num_tokens, hidden_size)``. + """ + from math import erf, sqrt + + hidden = tokens @ weights.projection_first_weight.T + weights.projection_first_bias + # Exact GELU (erf form), matching torch.nn.GELU()'s default approximate="none". + gelu = np.vectorize(lambda value: 0.5 * value * (1.0 + erf(value / sqrt(2.0)))) + activated = gelu(hidden).astype(np.float32) + projected = activated @ weights.projection_second_weight.T + weights.projection_second_bias + return projected.astype(np.float32) + + +def phi3_vision_hd_feature_transform( + image_features: np.ndarray, + image_sizes: np.ndarray, + weights: Phi3VisionProjectorWeights, +) -> np.ndarray: + """Reproduce Phi-3.5-Vision ``hd_feature_transform`` + ``img_projection``. + + Mirrors ``modeling_phi3_v.Phi3ImageEmbedding.hd_feature_transform`` with the + fixed ``hd_transform_order == 'sub_glb'`` ordering used by Phi-3.5-Vision: + for each image the token layout is ``[sub-crop features + newlines, + glb_GN separator, global features + newlines]``. + + Args: + image_features: Raw CLIP patch features of shape + ``(num_images, num_crops + 1, 576, image_dim_out)`` — index 0 along + axis 1 is the global (thumbnail) crop, the rest are the HD sub-crops. + image_sizes: ``(num_images, 2)`` padded ``(height, width)`` in pixels + per image, as produced by the HuggingFace image processor. + weights: Host-side projector + separator weights. + + Returns: + Projected image embeddings of shape ``(total_image_tokens, hidden_size)`` + ready to be scattered into ``<|image|>`` token positions. + """ + image_features = np.asarray(image_features, dtype=np.float32) + image_sizes = np.asarray(image_sizes) + if image_features.ndim != 4: + raise ValueError( + "image_features must be (num_images, num_crops+1, 576, image_dim_out); " + f"got shape {image_features.shape}" + ) + + global_separator = weights.global_separator.reshape(1, -1) # (1, image_dim_out*4) + sublayer_separator = weights.sublayer_separator # (1, 1, 1, image_dim_out*4) + + # The global (thumbnail) crop is a special 1x1 HD case, shared across images. + global_features = image_features[:, 0] # (num_images, 576, image_dim_out) + global_merged = _reshape_hd_patches_2x2_merge(global_features, 1, 1) + global_with_newline = _add_image_newline(global_merged, sublayer_separator) + + per_image_token_blocks: list[np.ndarray] = [] + for image_index, (height, width) in enumerate(image_sizes): + height_crops = int(height) // _CROP_PIXEL_SIZE + width_crops = int(width) // _CROP_PIXEL_SIZE + crop_count = height_crops * width_crops + + sub_features = image_features[image_index, 1 : 1 + crop_count] + sub_merged = _reshape_hd_patches_2x2_merge(sub_features, height_crops, width_crops) + sub_with_newline = _add_image_newline(sub_merged, sublayer_separator) + + per_image_token_blocks.extend( + [ + sub_with_newline[0], + global_separator, + global_with_newline[image_index], + ] + ) + + all_tokens = np.concatenate(per_image_token_blocks, axis=0) + return _apply_image_projection(all_tokens, weights) diff --git a/src/mobius/models/_phi3_vision_projector_test.py b/src/mobius/models/_phi3_vision_projector_test.py new file mode 100644 index 00000000..411e68e6 --- /dev/null +++ b/src/mobius/models/_phi3_vision_projector_test.py @@ -0,0 +1,223 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Unit tests for the Phi-3.5-Vision host-side HD feature transform + projector. + +These are L1 (pure NumPy) tests: they verify the 2x2 spatial patch merge, the +learnable separator insertion, the ``img_projection`` MLP math, the ``sub_glb`` +token ordering / token counts, and the checkpoint-weight loader — all without +downloading the real checkpoint. Correctness of the transform against the +HuggingFace reference is additionally covered end-to-end by the L4 golden test +(``testdata/cases/vision-language/phi3_5-vision-instruct.yaml``). +""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest + +from mobius.models._phi3_vision_projector import ( + Phi3VisionProjectorWeights, + _add_image_newline, + _apply_image_projection, + _reshape_hd_patches_2x2_merge, + load_phi3_vision_projector_weights, + phi3_vision_hd_feature_transform, +) + +# CLIP ViT-L/14-336 produces a 24x24 = 576 patch grid. +_PATCH_COUNT = 576 +_GRID_SIDE = 24 + + +def _make_projector_weights( + image_dim_out: int, hidden_size: int, seed: int = 0 +) -> Phi3VisionProjectorWeights: + """Build small random projector weights with the real tensor ranks/shapes.""" + rng = np.random.default_rng(seed) + merged_dim = image_dim_out * 4 + return Phi3VisionProjectorWeights( + global_separator=rng.standard_normal((1, 1, merged_dim)).astype(np.float32), + sublayer_separator=rng.standard_normal((1, 1, 1, merged_dim)).astype(np.float32), + projection_first_weight=rng.standard_normal((hidden_size, merged_dim)).astype( + np.float32 + ), + projection_first_bias=rng.standard_normal((hidden_size,)).astype(np.float32), + projection_second_weight=rng.standard_normal((hidden_size, hidden_size)).astype( + np.float32 + ), + projection_second_bias=rng.standard_normal((hidden_size,)).astype(np.float32), + ) + + +def test_reshape_hd_patches_2x2_merge_orders_the_four_subpatches() -> None: + """Each merged token concatenates its 2x2 block in (row, col) raster order.""" + # One image, one crop, single channel. Patch value encodes its grid position + # so we can assert exactly which four patches land in each merged token. + values = np.arange(_PATCH_COUNT, dtype=np.float32).reshape(1, _PATCH_COUNT, 1) + merged = _reshape_hd_patches_2x2_merge(values, height_crops=1, width_crops=1) + + # (num_images, h_crop*12, w_crop*12, 4*C) = (1, 12, 12, 4) + assert merged.shape == (1, 12, 12, 4) + # Top-left merged token = patches (0,0),(0,1),(1,0),(1,1) = [0, 1, 24, 25]. + np.testing.assert_array_equal(merged[0, 0, 0], [0.0, 1.0, 24.0, 25.0]) + # Merged token (i=1, j=2) = patches at rows 2/3, cols 4/5. + expected = [ + 2 * _GRID_SIDE + 4, + 2 * _GRID_SIDE + 5, + 3 * _GRID_SIDE + 4, + 3 * _GRID_SIDE + 5, + ] + np.testing.assert_array_equal(merged[0, 1, 2], expected) + + +def test_reshape_hd_patches_rejects_wrong_patch_count() -> None: + bad = np.zeros((1, _PATCH_COUNT - 1, 4), dtype=np.float32) + with pytest.raises(ValueError, match="patches"): + _reshape_hd_patches_2x2_merge(bad, 1, 1) + + +def test_reshape_hd_patches_rejects_indivisible_crop_grid() -> None: + # 3 crops cannot be arranged as a 2x2 grid. + three_crops = np.zeros((3, _PATCH_COUNT, 4), dtype=np.float32) + with pytest.raises(ValueError, match="divisible"): + _reshape_hd_patches_2x2_merge(three_crops, height_crops=2, width_crops=2) + + +def test_add_image_newline_appends_one_separator_per_row() -> None: + merged = np.zeros((1, 12, 12, 4), dtype=np.float32) + separator = np.full((1, 1, 1, 4), 7.0, dtype=np.float32) + with_newline = _add_image_newline(merged, separator) + + # Row width grows by one (the separator column): 12 -> 13, flattened. + assert with_newline.shape == (1, 12 * 13, 4) + # Every 13th token (end of each row) is the separator value. + separators = with_newline[0].reshape(12, 13, 4)[:, 12, :] + np.testing.assert_array_equal(separators, np.full((12, 4), 7.0)) + + +def test_apply_image_projection_matches_manual_linear_gelu_linear() -> None: + weights = _make_projector_weights(image_dim_out=8, hidden_size=16, seed=1) + tokens = np.random.default_rng(2).standard_normal((5, 32)).astype(np.float32) + + result = _apply_image_projection(tokens, weights) + + hidden = tokens @ weights.projection_first_weight.T + weights.projection_first_bias + gelu = np.array( + [0.5 * v * (1.0 + math.erf(v / math.sqrt(2.0))) for v in hidden.reshape(-1)] + ).reshape(hidden.shape) + expected = gelu @ weights.projection_second_weight.T + weights.projection_second_bias + + assert result.shape == (5, 16) + np.testing.assert_allclose(result, expected, rtol=1e-5, atol=1e-5) + + +def test_hd_feature_transform_token_count_and_shape_for_2x2_crops() -> None: + """A 672x672 image (2x2 HD crops) yields 757 projected image tokens. + + Layout per image = [sub features + newlines, glb_GN, global features + + newlines] = 600 + 1 + 156 = 757, matching HuggingFace and the golden case. + """ + image_dim_out, hidden_size = 8, 16 + weights = _make_projector_weights(image_dim_out, hidden_size, seed=3) + num_crops_including_global = 5 # 1 global + 2x2 sub-crops + image_features = np.random.default_rng(4).standard_normal( + (1, num_crops_including_global, _PATCH_COUNT, image_dim_out) + ) + image_sizes = np.array([[672, 672]]) + + projected = phi3_vision_hd_feature_transform(image_features, image_sizes, weights) + + sub_tokens = 24 * (24 + 1) # 2x2 crops -> 24x24 merged grid + newline column + global_tokens = 12 * (12 + 1) # 1x1 global crop -> 12x12 grid + newline column + expected_tokens = sub_tokens + 1 + global_tokens + assert expected_tokens == 757 + assert projected.shape == (expected_tokens, hidden_size) + + +def test_hd_feature_transform_places_global_separator_between_sub_and_global() -> None: + """The single glb_GN separator token sits exactly after the sub features.""" + image_dim_out, hidden_size = 4, 6 + weights = _make_projector_weights(image_dim_out, hidden_size, seed=5) + image_features = np.random.default_rng(6).standard_normal( + (1, 5, _PATCH_COUNT, image_dim_out) + ) + image_sizes = np.array([[672, 672]]) + + # Project only the raw (pre-projection) token stream by using an identity-ish + # projector is overkill; instead verify the separator lands at the right + # index by re-deriving the un-projected concatenation length. + sub_tokens = 24 * (24 + 1) + projected = phi3_vision_hd_feature_transform(image_features, image_sizes, weights) + + # The projector is deterministic, so the separator token (index sub_tokens) + # must equal the projection of the flattened glb_GN separator. + separator_projected = _apply_image_projection( + weights.global_separator.reshape(1, -1), weights + )[0] + np.testing.assert_allclose( + projected[sub_tokens], separator_projected, rtol=1e-5, atol=1e-5 + ) + + +def test_hd_feature_transform_rejects_non_4d_features() -> None: + weights = _make_projector_weights(4, 6) + with pytest.raises(ValueError, match="num_images"): + phi3_vision_hd_feature_transform( + np.zeros((5, _PATCH_COUNT, 4)), np.array([[336, 336]]), weights + ) + + +def test_load_projector_weights_reads_and_validates_safetensors(tmp_path) -> None: + safetensors_numpy = pytest.importorskip("safetensors.numpy") + image_dim_out, hidden_size = 8, 16 + merged_dim = image_dim_out * 4 + tensors = { + "model.vision_embed_tokens.glb_GN": np.ones((1, 1, merged_dim), dtype=np.float32), + "model.vision_embed_tokens.sub_GN": np.full( + (1, 1, 1, merged_dim), 2.0, dtype=np.float32 + ), + "model.vision_embed_tokens.img_projection.0.weight": np.zeros( + (hidden_size, merged_dim), dtype=np.float32 + ), + "model.vision_embed_tokens.img_projection.0.bias": np.zeros( + (hidden_size,), dtype=np.float32 + ), + "model.vision_embed_tokens.img_projection.2.weight": np.zeros( + (hidden_size, hidden_size), dtype=np.float32 + ), + "model.vision_embed_tokens.img_projection.2.bias": np.zeros( + (hidden_size,), dtype=np.float32 + ), + # An unrelated tensor that must be ignored by the loader. + "model.layers.0.mlp.weight": np.zeros((2, 2), dtype=np.float32), + } + shard_path = tmp_path / "model-00001-of-00001.safetensors" + safetensors_numpy.save_file(tensors, str(shard_path)) + + weights = load_phi3_vision_projector_weights(str(tmp_path)) + + assert weights.global_separator.shape == (1, 1, merged_dim) + assert weights.sublayer_separator.shape == (1, 1, 1, merged_dim) + assert weights.projection_first_weight.shape == (hidden_size, merged_dim) + np.testing.assert_array_equal( + weights.sublayer_separator, tensors["model.vision_embed_tokens.sub_GN"] + ) + + +def test_load_projector_weights_raises_on_missing_shards(tmp_path) -> None: + with pytest.raises(FileNotFoundError, match="safetensors"): + load_phi3_vision_projector_weights(str(tmp_path)) + + +def test_load_projector_weights_raises_on_missing_projector_tensor(tmp_path) -> None: + safetensors_numpy = pytest.importorskip("safetensors.numpy") + # Shard exists but lacks the projector tensors. + safetensors_numpy.save_file( + {"model.layers.0.mlp.weight": np.zeros((2, 2), dtype=np.float32)}, + str(tmp_path / "model.safetensors"), + ) + with pytest.raises(KeyError, match="projector"): + load_phi3_vision_projector_weights(str(tmp_path)) diff --git a/testdata/cases/vision-language/phi3_5-vision-instruct.yaml b/testdata/cases/vision-language/phi3_5-vision-instruct.yaml index 383075e6..418542fe 100644 --- a/testdata/cases/vision-language/phi3_5-vision-instruct.yaml +++ b/testdata/cases/vision-language/phi3_5-vision-instruct.yaml @@ -17,5 +17,4 @@ generation: do_sample: false trust_remote_code: true -skip_reason: "Custom trust_remote_code model code written for transformers 4.x is incompatible with transformers 5.x (DynamicCache API changes, attention mask shape changes)." -notes: "Phi-3.5 Vision. CLIP ViT-L/14-336 vision encoder + MLP projector + Phi-3.5 text decoder. Vision config stored in img_processor dict (non-standard)." +notes: "Phi-3.5 Vision. CLIP ViT-L/14-336 vision encoder + host-side HD feature transform + img_projection MLP + Phi-3.5 (rope_type=su) text decoder. Vision config stored in img_processor dict (non-standard). Vision encoder ONNX emits raw CLIP patch features (num_crops, 576, 1024); the HD 2x2 patch merge, learnable sub_GN/glb_GN separators and img_projection run host-side (image-size dependent)." diff --git a/testdata/golden/vision-language/phi3_5-vision-instruct.json b/testdata/golden/vision-language/phi3_5-vision-instruct.json new file mode 100644 index 00000000..1ecfb081 --- /dev/null +++ b/testdata/golden/vision-language/phi3_5-vision-instruct.json @@ -0,0 +1,811 @@ +{ + "top1_id": 450, + "top2_id": 910, + "top10_ids": [ + 450, + 910, + 319, + 530, + 29871, + 32007, + 8221, + 306, + 349, + 1670 + ], + "top10_logits": [ + "0x1.04b3ba0000000p+5", + "0x1.cb0f360000000p+4", + "0x1.c1716c0000000p+4", + "0x1.b422440000000p+4", + "0x1.adda8e0000000p+4", + "0x1.97ecc80000000p+4", + "0x1.8dbcec0000000p+4", + "0x1.8632f20000000p+4", + "0x1.8321ea0000000p+4", + "0x1.7565ae0000000p+4" + ], + "logits_summary": [ + "0x1.04b3ba0000000p+5", + "-0x1.890f1c0000000p+0", + "0x1.6ebfb30e24062p+3", + "0x1.41c185e346574p+1" + ], + "input_ids": [ + 1, + 32010, + 29871, + 13, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + -1, + 1, + 29871, + 13, + 5618, + 338, + 4318, + 297, + 445, + 1967, + 29973, + 32007, + 29871, + 13, + 32001 + ] +} diff --git a/testdata/golden/vision-language/phi3_5-vision-instruct_generation.json b/testdata/golden/vision-language/phi3_5-vision-instruct_generation.json new file mode 100644 index 00000000..6db29136 --- /dev/null +++ b/testdata/golden/vision-language/phi3_5-vision-instruct_generation.json @@ -0,0 +1,37 @@ +{ + "model_id": "microsoft/Phi-3.5-vision-instruct", + "prompt": "What is shown in this image?", + "generated_tokens": [ + 450, + 1967, + 3697, + 263, + 349, + 15892, + 29915, + 29879, + 6635, + 29892, + 263, + 2319, + 8775, + 6635, + 7531, + 304, + 278, + 17455, + 5252, + 322, + 1886, + 407, + 267, + 310, + 8068, + 14325, + 29889, + 739, + 338, + 2998 + ], + "generated_text": "The image shows a Pallas's cat, a small wild cat native to the grasslands and steppes of Central Asia. It is known" +} diff --git a/tests/e2e_golden_test.py b/tests/e2e_golden_test.py index 713fd05e..a5adcd38 100644 --- a/tests/e2e_golden_test.py +++ b/tests/e2e_golden_test.py @@ -148,6 +148,28 @@ def _build_mm_prompt( return processor.apply_chat_template( # type: ignore[attr-defined] messages, tokenize=False, add_generation_prompt=True ) + # Some processors (e.g. Phi-3.5-Vision's ``Phi3VProcessor``) expose the chat + # template on the *inner tokenizer* rather than the processor, and use + # numbered placeholders (``<|image_1|>``) enumerated in ``img_tokens``. + inner_tokenizer = getattr(processor, "tokenizer", None) + numbered_token_attribute = {"image": "img_tokens", "audio": "audio_tokens"}.get(media_kind) + numbered_media_tokens = ( + getattr(processor, numbered_token_attribute, None) + if numbered_token_attribute is not None + else None + ) + if ( + inner_tokenizer is not None + and numbered_media_tokens is not None + and getattr(inner_tokenizer, "chat_template", None) + ): + media_prefix = "".join( + f"{numbered_media_tokens[index]}\n" for index in range(len(media_paths)) + ) + messages = [{"role": "user", "content": media_prefix + base_prompt}] + return inner_tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) placeholder = getattr(processor, f"{media_kind}_token", None) if placeholder: return placeholder * len(media_paths) + base_prompt @@ -272,6 +294,24 @@ def _use_temp_hf_cache(tmp_path): ), # MLA compressed KV cache dimensions not yet handled by OnnxGenerator "text-generation/youtu-2b": "Youtu MLA KV cache dims differ from standard attention (v_head_dim != head_dim)", + # Phi-3.5-Vision L5 greedy generation drifts from HF after the first few + # tokens due to the SAME decoder-side LongRoPE parity gap already xfailed for + # the sibling phi4mm case (both share the Phi-3.5 text decoder): + # * L4 prefill PASSES — top-1 argmax matches golden over all 775 tokens, so + # the vision tower + host-side HD-transform/img_projection + embedding + # fusion are correct (the host-side projector matches HF at cos=1.0000, + # max-abs 6.9e-5 vs modeling_phi3_v.hd_feature_transform). + # * The first generated tokens match HF exactly ("The image shows a ..."); + # divergence begins mid-generation, deep in the text decoder. + # * Root cause is decoder-side logit drift (LongRoPE short-path freqs, + # rotary_dim=96, attention_scaling), NOT a vision/projector defect — + # identical to phi4mm-multi-image-audio above. + "image-text-to-text/phi3_5-vision-instruct": ( + "Phi-3.5-Vision decode loop drifts from HF after the first few tokens " + "(L4 prefill passes; host-side HD-transform/projector match HF at " + "cos=1.0; first tokens match). Decoder-side LongRoPE parity gap, same " + "as phi4mm-multi-image-audio — not a vision/projector defect." + ), } @@ -709,6 +749,71 @@ def _compute_mrope_position_ids( return position_ids +def _run_vl_vision_to_image_features( + pkg: ModelPackage, + case: GoldenTestCase, + processed: dict[str, np.ndarray], +) -> tuple[dict[str, np.ndarray], np.ndarray]: + """Run the ONNX vision encoder and return the embedding model's image features. + + Returns ``(vision_outputs, image_features)`` where ``image_features`` is the + tensor to feed into the embedding model's ``image_features`` input. + + For most vision-language models the vision encoder output is used directly. + Phi-3.5-Vision (``model_type == "phi3_v"``) additionally requires the + host-side HD feature transform + ``img_projection`` MLP — an image-size + dependent step that cannot be a static ONNX graph — to map the raw CLIP + patch features ``(num_crops, 576, image_dim_out)`` into projected image + embeddings ``(total_image_tokens, hidden_size)``. + """ + vis_session = OnnxModelSession(pkg["vision_encoder"], **_get_test_device_kwargs()) + try: + if case.model_type == "phi3_v": + from mobius.models._phi3_vision_projector import ( + load_phi3_vision_projector_weights, + phi3_vision_hd_feature_transform, + ) + + # The HF processor emits pixel_values shaped + # (num_images, num_crops + 1, 3, 336, 336). The ONNX vision encoder + # takes 4-D NCHW input, so flatten the leading crop dimension. + pixel_values = np.asarray(processed["pixel_values"], dtype=np.float32) + num_images, num_crops_including_global = pixel_values.shape[:2] + flat_pixel_values = pixel_values.reshape( + num_images * num_crops_including_global, *pixel_values.shape[2:] + ) + vis_out = vis_session.run({"pixel_values": flat_pixel_values}) + raw_patch_features = vis_out[next(iter(vis_out))] + image_dim_out = raw_patch_features.shape[-1] + image_features_per_crop = raw_patch_features.reshape( + num_images, num_crops_including_global, -1, image_dim_out + ) + projector_weights = load_phi3_vision_projector_weights(case.model_id) + projected = phi3_vision_hd_feature_transform( + image_features_per_crop, + np.asarray(processed["image_sizes"]), + projector_weights, + ) + return vis_out, projected + + vis_feeds: dict[str, np.ndarray] = {} + for name in vis_session.input_names: + if name in processed: + val = processed[name] + vis_feeds[name] = val if isinstance(val, np.ndarray) else np.array(val) + else: + # Handle HF↔ONNX name mismatches (e.g. HF "image_position_ids" + # vs ONNX "pixel_position_ids"). + for hf_key, val in processed.items(): + if hf_key.replace("image_", "pixel_") == name: + vis_feeds[name] = val if isinstance(val, np.ndarray) else np.array(val) + break + vis_out = vis_session.run(vis_feeds) + return vis_out, vis_out[next(iter(vis_out))] + finally: + vis_session.close() + + def _run_vision_language_prefill( pkg: ModelPackage, case: GoldenTestCase, @@ -742,28 +847,8 @@ def _run_vision_language_prefill( k: v.numpy() if hasattr(v, "numpy") else np.array(v) for k, v in processed_pt.items() } - # --- Step 1: Run vision encoder --- - vis_session = OnnxModelSession(pkg["vision_encoder"], **_get_test_device_kwargs()) - try: - vis_feeds: dict[str, np.ndarray] = {} - for name in vis_session.input_names: - if name in processed: - val = processed[name] - vis_feeds[name] = val if isinstance(val, np.ndarray) else np.array(val) - else: - # Handle HF↔ONNX name mismatches (e.g. HF "image_position_ids" - # vs ONNX "pixel_position_ids"). - for hf_key, val in processed.items(): - if hf_key.replace("image_", "pixel_") == name: - vis_feeds[name] = val if isinstance(val, np.ndarray) else np.array(val) - break - vis_out = vis_session.run(vis_feeds) - finally: - vis_session.close() - - # Extract the image hidden states (first output) - vis_hidden_key = next(iter(vis_out)) - vis_hidden = vis_out[vis_hidden_key] + # --- Step 1: Run vision encoder (+ host-side projector where required) --- + vis_out, image_features = _run_vl_vision_to_image_features(pkg, case, processed) # --- Step 2: Run embedding model --- emb_session = OnnxModelSession(pkg["embedding"], **_get_test_device_kwargs()) @@ -773,10 +858,10 @@ def _run_vision_language_prefill( } # Pass vision hidden states for name in emb_session.input_names: - if name not in emb_feeds and name in vis_out: + if name == "image_features": + emb_feeds[name] = image_features + elif name not in emb_feeds and name in vis_out: emb_feeds[name] = vis_out[name] - elif name == "image_features": - emb_feeds[name] = vis_hidden elif name not in emb_feeds: # Provide empty tensor for unused modalities (e.g. audio_features) shape = emb_session.get_input_shape(name) or [] @@ -926,24 +1011,8 @@ def _run_vl_generation( k: v.numpy() if hasattr(v, "numpy") else np.array(v) for k, v in processed_pt.items() } - # --- Step 1: vision encoder --- - vis_session = OnnxModelSession(pkg["vision_encoder"], **_get_test_device_kwargs()) - try: - vis_feeds: dict[str, np.ndarray] = {} - for name in vis_session.input_names: - if name in processed: - vis_feeds[name] = processed[name] - else: - # HF↔ONNX name mismatch (e.g. image_position_ids → pixel_position_ids) - for hf_key, val in processed.items(): - if hf_key.replace("image_", "pixel_") == name: - vis_feeds[name] = val if isinstance(val, np.ndarray) else np.array(val) - break - vis_out = vis_session.run(vis_feeds) - finally: - vis_session.close() - - vis_hidden = vis_out[next(iter(vis_out))] # image feature tensor + # --- Step 1: vision encoder (+ host-side projector where required) --- + _vis_out, image_features = _run_vl_vision_to_image_features(pkg, case, processed) # --- Step 2: embedding (prefill) --- # VL packages use "decoder" as the decoder key @@ -962,7 +1031,7 @@ def _run_vl_generation( "input_ids": processed["input_ids"].astype(np.int64), } if image_feat_input is not None: - emb_feeds[image_feat_input] = vis_hidden + emb_feeds[image_feat_input] = image_features # Provide empty tensors for unused modalities (e.g. audio_features) for name in emb_session.input_names: if name not in emb_feeds: From 2b251dd228711ef674d1bb5d4303d52f818f264f Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 25 Jul 2026 19:22:09 +0000 Subject: [PATCH 04/10] phi3_v: fix image-feature fusion (negative placeholder ids); clear L5 xfail MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root-cause the Phi-3.5-Vision L5 greedy-generation drift that was xfailed as a "decoder-side LongRoPE parity gap". It is NOT LongRoPE. Evidence (fp32/fp64 oracle + instrumented decode): * The LongRoPE cos/sin cache matches HF to ~3.6e-5 (a 1-2 ULP inv_freq ordering difference amplified by position). An A/B rebuild that switched the mobius inv_freq to HF's exact `1/(factor*theta**x)` ordering changed the step-4 logits by <1e-4 and did NOT change any generated token — so rope precision is not the cause. * Instrumented decode shows the divergence is a HARD miss at the first image-grounded token (step 4): the golden token ("P" of "Pallas's cat") is ranked 481 in mobius with a 7.28-logit deficit, while mobius emits "The image shows a mathematical expression written in LaTeX..." — the classic Phi-3.5-Vision *blind* (no-image) hallucination. Actual root cause: the embedding graph detected image slots with `input_ids == image_token_id` (32044), but the HF Phi-3/3.5-Vision processor marks image positions with *negative* placeholder ids (-1, -2, ...). The mask therefore never fired, image features were never fused into the prompt embeddings, and the decoder ran blind. The first few generic tokens ("The image shows a") match by text prior, then it diverges at the first token that needs the image. L4 prefill "passed" only because its argmax ("The") is image-independent. Fix (mirrors modeling_phi3_v exactly): * Detect image positions as `(input_ids < 0) & (input_ids > -MAX_INPUT_ID)`. * Pad image_features with one trailing zero row before the Gather so decode steps (which pass an empty 0-row image tensor) don't index an empty tensor and crash — the padding is discarded by the fusion `Where`. Result: phi3_5-vision L5 now matches the HF golden token-for-token (30/30); L4 still passes; phi-3_5-mini L4+L5 unchanged; build_graph + cli fast lane green. Cleared the phi3_5-vision L5 xfail. The sibling phi4mm-multi-image-audio xfail is a genuinely different issue and is left in place: its processor emits POSITIVE image/audio ids (200010/200011) that the InputMixer already matches correctly, so its fusion is not affected. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/models/phi3_v.py | 39 ++++++++++++++++++++++++++++++++----- tests/e2e_golden_test.py | 18 ----------------- 2 files changed, 34 insertions(+), 23 deletions(-) diff --git a/src/mobius/models/phi3_v.py b/src/mobius/models/phi3_v.py index 836b4f0a..d972c712 100644 --- a/src/mobius/models/phi3_v.py +++ b/src/mobius/models/phi3_v.py @@ -71,9 +71,16 @@ if TYPE_CHECKING: import onnx_ir as ir -# Phi-3-Vision: <|image|> token id +# 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``. +_MAX_INPUT_ID = 1_000_000_000 + class _Phi3VDecoderModel(nn.Module): """Phi3-V text decoder taking inputs_embeds. @@ -227,8 +234,10 @@ def preprocess_weights( class _Phi3VEmbeddingModel(nn.Module): """Phi3-V embedding: token lookup + image feature fusion. - Replaces <|image|> token positions with projected vision features - from the vision encoder. + Replaces image placeholder positions with projected vision features from + the vision encoder. Phi-3/3.5-Vision processors mark those positions with + *negative* placeholder ids (-1, -2, ...), matched here the same way HF's + ``modeling_phi3_v`` does — see ``forward``. """ def __init__(self, config: ArchitectureConfig): @@ -236,13 +245,23 @@ def __init__(self, config: ArchitectureConfig): self.embed_tokens = Embedding( config.vocab_size, config.hidden_size, config.pad_token_id ) + self.hidden_size = config.hidden_size self.image_token_id = ( config.image_token_id if config.image_token_id is not None else _IMAGE_TOKEN_ID ) def forward(self, op: builder.OpBuilder, input_ids: ir.Value, image_features: ir.Value): # (batch, seq_len, hidden_size) - image_mask = op.Equal(input_ids, op.Constant(value_int=self.image_token_id)) + # Phi-3/3.5-Vision processors mark image placeholder positions with + # *negative* token ids (-1 for the first image, -2 for the second, ...), + # NOT with a positive ``image_token_id``. Detect them exactly as HF's + # ``modeling_phi3_v`` does: ``(input_ids < 0) & (input_ids > -MAX_INPUT_ID)``. + # Matching a positive id here would never fire for real processor output, + # leaving image features unfused and the decoder running "blind". + image_mask = op.And( + op.Less(input_ids, op.Constant(value_int=0)), + op.Greater(input_ids, op.Constant(value_int=-_MAX_INPUT_ID)), + ) # Replace image token positions with index 0 before the Gather so that # negative or out-of-range image token IDs don't cause undefined behavior. # These positions will be overwritten by image features via op.Where below. @@ -256,7 +275,17 @@ def forward(self, op: builder.OpBuilder, input_ids: ir.Value, image_features: ir indices = op.Sub(cumsum, op.Constant(value_int=1)) indices = op.Clip(indices, op.Constant(value_int=0)) - gathered = op.Gather(image_features, indices, axis=0) + # Pad ``image_features`` with a single trailing zero row so the Gather is + # always valid — including at decode time, when there are no image + # placeholders and the harness/runtime passes an empty (0-row) tensor. + # Without the pad, ``Gather`` would index row 0 of a 0-row tensor and + # raise an out-of-bounds error. The gathered padding is discarded by the + # ``Where`` below (image_mask is all-false when there are no image rows). + zero_row = op.ConstantOfShape(op.Constant(value_ints=[1, self.hidden_size])) + zero_row = op.CastLike(zero_row, image_features) + padded_features = op.Concat(image_features, zero_row, axis=0) + + gathered = op.Gather(padded_features, indices, axis=0) return op.Where(image_mask_3d, gathered, text_embeds) def preprocess_weights( diff --git a/tests/e2e_golden_test.py b/tests/e2e_golden_test.py index a5adcd38..31ddf38a 100644 --- a/tests/e2e_golden_test.py +++ b/tests/e2e_golden_test.py @@ -294,24 +294,6 @@ def _use_temp_hf_cache(tmp_path): ), # MLA compressed KV cache dimensions not yet handled by OnnxGenerator "text-generation/youtu-2b": "Youtu MLA KV cache dims differ from standard attention (v_head_dim != head_dim)", - # Phi-3.5-Vision L5 greedy generation drifts from HF after the first few - # tokens due to the SAME decoder-side LongRoPE parity gap already xfailed for - # the sibling phi4mm case (both share the Phi-3.5 text decoder): - # * L4 prefill PASSES — top-1 argmax matches golden over all 775 tokens, so - # the vision tower + host-side HD-transform/img_projection + embedding - # fusion are correct (the host-side projector matches HF at cos=1.0000, - # max-abs 6.9e-5 vs modeling_phi3_v.hd_feature_transform). - # * The first generated tokens match HF exactly ("The image shows a ..."); - # divergence begins mid-generation, deep in the text decoder. - # * Root cause is decoder-side logit drift (LongRoPE short-path freqs, - # rotary_dim=96, attention_scaling), NOT a vision/projector defect — - # identical to phi4mm-multi-image-audio above. - "image-text-to-text/phi3_5-vision-instruct": ( - "Phi-3.5-Vision decode loop drifts from HF after the first few tokens " - "(L4 prefill passes; host-side HD-transform/projector match HF at " - "cos=1.0; first tokens match). Decoder-side LongRoPE parity gap, same " - "as phi4mm-multi-image-audio — not a vision/projector defect." - ), } From de328f287720a3a72359d77ecf01ba32a30a5863 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 25 Jul 2026 21:54:45 +0000 Subject: [PATCH 05/10] phi3_v: address projector review nits Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/models/_phi3_vision_projector.py | 24 +++++++++++++++---- .../models/_phi3_vision_projector_test.py | 17 +++++++++++++ src/mobius/models/clip.py | 8 ++----- testdata/golden/README.md | 5 ++++ tests/e2e_golden_test.py | 5 +++- 5 files changed, 48 insertions(+), 11 deletions(-) create mode 100644 testdata/golden/README.md diff --git a/src/mobius/models/_phi3_vision_projector.py b/src/mobius/models/_phi3_vision_projector.py index a40662f8..f88a3248 100644 --- a/src/mobius/models/_phi3_vision_projector.py +++ b/src/mobius/models/_phi3_vision_projector.py @@ -30,6 +30,7 @@ import os import numpy as np +import torch # HF checkpoint key prefix for the projector / separators (outside the CLIP # ``img_processor`` vision tower, which is exported to ONNX separately). @@ -221,12 +222,10 @@ def _apply_image_projection( Returns: ``(num_tokens, hidden_size)``. """ - from math import erf, sqrt - hidden = tokens @ weights.projection_first_weight.T + weights.projection_first_bias # Exact GELU (erf form), matching torch.nn.GELU()'s default approximate="none". - gelu = np.vectorize(lambda value: 0.5 * value * (1.0 + erf(value / sqrt(2.0)))) - activated = gelu(hidden).astype(np.float32) + hidden_tensor = torch.from_numpy(hidden) + activated = (0.5 * hidden_tensor * (1.0 + torch.erf(hidden_tensor / np.sqrt(2.0)))).numpy() projected = activated @ weights.projection_second_weight.T + weights.projection_second_bias return projected.astype(np.float32) @@ -262,6 +261,12 @@ def phi3_vision_hd_feature_transform( "image_features must be (num_images, num_crops+1, 576, image_dim_out); " f"got shape {image_features.shape}" ) + num_images, num_crops_including_global, _, _ = image_features.shape + if image_sizes.shape != (num_images, 2): + raise ValueError( + "image_sizes must be (num_images, 2) matching image_features; " + f"got shape {image_sizes.shape} for {num_images} images" + ) global_separator = weights.global_separator.reshape(1, -1) # (1, image_dim_out*4) sublayer_separator = weights.sublayer_separator # (1, 1, 1, image_dim_out*4) @@ -273,9 +278,20 @@ def phi3_vision_hd_feature_transform( per_image_token_blocks: list[np.ndarray] = [] for image_index, (height, width) in enumerate(image_sizes): + if height <= 0 or width <= 0 or height % _CROP_PIXEL_SIZE or width % _CROP_PIXEL_SIZE: + raise ValueError( + f"image_sizes[{image_index}] must contain positive multiples of " + f"{_CROP_PIXEL_SIZE}; got ({height}, {width})" + ) height_crops = int(height) // _CROP_PIXEL_SIZE width_crops = int(width) // _CROP_PIXEL_SIZE crop_count = height_crops * width_crops + available_sub_crops = num_crops_including_global - 1 + if crop_count > available_sub_crops: + raise ValueError( + f"image_sizes[{image_index}] requires {crop_count} HD crops, but " + f"image_features provides only {available_sub_crops} sub-crops" + ) sub_features = image_features[image_index, 1 : 1 + crop_count] sub_merged = _reshape_hd_patches_2x2_merge(sub_features, height_crops, width_crops) diff --git a/src/mobius/models/_phi3_vision_projector_test.py b/src/mobius/models/_phi3_vision_projector_test.py index 411e68e6..8ec0e1ea 100644 --- a/src/mobius/models/_phi3_vision_projector_test.py +++ b/src/mobius/models/_phi3_vision_projector_test.py @@ -170,6 +170,23 @@ def test_hd_feature_transform_rejects_non_4d_features() -> None: ) +@pytest.mark.parametrize("image_size", [(335, 336), (336, 337), (0, 336)]) +def test_hd_feature_transform_rejects_invalid_image_sizes( + image_size: tuple[int, int], +) -> None: + weights = _make_projector_weights(4, 6) + image_features = np.zeros((1, 2, _PATCH_COUNT, 4), dtype=np.float32) + with pytest.raises(ValueError, match="positive multiples"): + phi3_vision_hd_feature_transform(image_features, np.array([image_size]), weights) + + +def test_hd_feature_transform_rejects_insufficient_sub_crops() -> None: + weights = _make_projector_weights(4, 6) + image_features = np.zeros((1, 4, _PATCH_COUNT, 4), dtype=np.float32) + with pytest.raises(ValueError, match="requires 4 HD crops"): + phi3_vision_hd_feature_transform(image_features, np.array([[672, 672]]), weights) + + def test_load_projector_weights_reads_and_validates_safetensors(tmp_path) -> None: safetensors_numpy = pytest.importorskip("safetensors.numpy") image_dim_out, hidden_size = 8, 16 diff --git a/src/mobius/models/clip.py b/src/mobius/models/clip.py index 85304a1f..34a1b3d1 100644 --- a/src/mobius/models/clip.py +++ b/src/mobius/models/clip.py @@ -12,17 +12,13 @@ from mobius._configs import ArchitectureConfig from mobius.components import FCMLP -from mobius.components._common import Embedding, LayerNorm +from mobius.components._common import INT64_MAX, Embedding, LayerNorm from mobius.components._conv import Conv2d, Conv2dNoBias from mobius.components._encoder import EncoderAttention if TYPE_CHECKING: import onnx_ir as ir -# Largest int64 value; used as an open-ended ``Slice`` end (ONNX clamps it to -# the actual dimension size). -_INT64_MAX = 9223372036854775807 - class ClipVisionConfigView: """Adapter exposing a :class:`VisionConfig` under CLIP's field names. @@ -244,7 +240,7 @@ def forward(self, op: OpBuilder, pixel_values: ir.Value): hidden_states = op.Slice( hidden_states, [1], # starts - [_INT64_MAX], # ends (clamped to sequence length) + [INT64_MAX], # ends (clamped to sequence length) [1], # axes: sequence dimension ) diff --git a/testdata/golden/README.md b/testdata/golden/README.md new file mode 100644 index 00000000..10b40e31 --- /dev/null +++ b/testdata/golden/README.md @@ -0,0 +1,5 @@ +# Golden reference provenance + +The Phi-3.5 JSON goldens were generated offline with `transformers==4.43`. +They are replayed by the L4/L5 tests with the supported current Transformers +release; do not regenerate them as part of those tests. diff --git a/tests/e2e_golden_test.py b/tests/e2e_golden_test.py index 31ddf38a..c9283975 100644 --- a/tests/e2e_golden_test.py +++ b/tests/e2e_golden_test.py @@ -759,7 +759,10 @@ def _run_vl_vision_to_image_features( # The HF processor emits pixel_values shaped # (num_images, num_crops + 1, 3, 336, 336). The ONNX vision encoder # takes 4-D NCHW input, so flatten the leading crop dimension. - pixel_values = np.asarray(processed["pixel_values"], dtype=np.float32) + pixel_values_dtype = vis_session.get_input_dtype("pixel_values") + if pixel_values_dtype is None: + raise ValueError("Phi-3.5 vision encoder must have a pixel_values input dtype") + pixel_values = np.asarray(processed["pixel_values"], dtype=pixel_values_dtype) num_images, num_crops_including_global = pixel_values.shape[:2] flat_pixel_values = pixel_values.reshape( num_images * num_crops_including_global, *pixel_values.shape[2:] From 3aa714232b10b22c257d6d64da3f06788ca63e0f Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sun, 26 Jul 2026 23:27:47 +0000 Subject: [PATCH 06/10] Address Phi-3.5 exporter review feedback Remove the stale positive image-token field, validate numbered media placeholders, cache projector weights across L4/L5, and correct the host-projector dependency documentation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/models/_phi3_vision_projector.py | 2 +- src/mobius/models/phi3_v.py | 3 -- tests/e2e_golden_test.py | 39 +++++++++++++++++++-- 3 files changed, 38 insertions(+), 6 deletions(-) diff --git a/src/mobius/models/_phi3_vision_projector.py b/src/mobius/models/_phi3_vision_projector.py index f88a3248..8371728c 100644 --- a/src/mobius/models/_phi3_vision_projector.py +++ b/src/mobius/models/_phi3_vision_projector.py @@ -11,7 +11,7 @@ grid, the separator interleaving, and the final token count all vary with the input image resolution, so they cannot be expressed as a static ONNX graph. -This module reproduces that host-side tail faithfully in NumPy, mirroring +This module reproduces that host-side tail faithfully with NumPy and PyTorch, mirroring ``modeling_phi3_v.py`` (``hd_feature_transform`` → ``reshape_hd_patches_2x2merge`` → ``add_image_newline`` → ``img_projection``). The output is the flat sequence of projected image embeddings ``(total_image_tokens, hidden_size)`` that the diff --git a/src/mobius/models/phi3_v.py b/src/mobius/models/phi3_v.py index d972c712..b16698de 100644 --- a/src/mobius/models/phi3_v.py +++ b/src/mobius/models/phi3_v.py @@ -246,9 +246,6 @@ def __init__(self, config: ArchitectureConfig): config.vocab_size, config.hidden_size, config.pad_token_id ) self.hidden_size = config.hidden_size - self.image_token_id = ( - config.image_token_id if config.image_token_id is not None else _IMAGE_TOKEN_ID - ) def forward(self, op: builder.OpBuilder, input_ids: ir.Value, image_features: ir.Value): # (batch, seq_len, hidden_size) diff --git a/tests/e2e_golden_test.py b/tests/e2e_golden_test.py index c9283975..aa2ab6c2 100644 --- a/tests/e2e_golden_test.py +++ b/tests/e2e_golden_test.py @@ -25,6 +25,7 @@ from __future__ import annotations import dataclasses +import functools import os import warnings from pathlib import Path @@ -50,6 +51,16 @@ from mobius._testing.ort_inference import OnnxModelSession from mobius._testing.parity import ParityResult, compare_golden + +@functools.cache +def _load_phi3_vision_projector_weights_cached(model_id: str): + from mobius.models._phi3_vision_projector import ( + load_phi3_vision_projector_weights, + ) + + return load_phi3_vision_projector_weights(model_id) + + # Root of test data (images, audio, etc.) _TESTDATA_DIR = Path(__file__).resolve().parent.parent / "testdata" @@ -163,6 +174,12 @@ def _build_mm_prompt( and numbered_media_tokens is not None and getattr(inner_tokenizer, "chat_template", None) ): + if len(media_paths) > len(numbered_media_tokens): + raise ValueError( + f"{type(processor).__name__} exposes {len(numbered_media_tokens)} " + f"numbered {media_kind} placeholders, but the case requested " + f"{len(media_paths)} media items" + ) media_prefix = "".join( f"{numbered_media_tokens[index]}\n" for index in range(len(media_paths)) ) @@ -176,6 +193,25 @@ def _build_mm_prompt( return base_prompt +def test_build_mm_prompt_rejects_more_media_than_numbered_tokens(): + tokenizer = type( + "_Tokenizer", + (), + { + "chat_template": "template", + "apply_chat_template": lambda *args, **kwargs: "", + }, + )() + processor = type( + "_Processor", + (), + {"tokenizer": tokenizer, "img_tokens": ["<|image_1|>"]}, + )() + + with pytest.raises(ValueError, match="exposes 1 numbered image placeholders"): + _build_mm_prompt(processor, "describe", ["a.png", "b.png"], "image") + + def _make_empty_kv_cache( session: OnnxModelSession, config: object, @@ -752,7 +788,6 @@ def _run_vl_vision_to_image_features( try: if case.model_type == "phi3_v": from mobius.models._phi3_vision_projector import ( - load_phi3_vision_projector_weights, phi3_vision_hd_feature_transform, ) @@ -773,7 +808,7 @@ def _run_vl_vision_to_image_features( image_features_per_crop = raw_patch_features.reshape( num_images, num_crops_including_global, -1, image_dim_out ) - projector_weights = load_phi3_vision_projector_weights(case.model_id) + projector_weights = _load_phi3_vision_projector_weights_cached(case.model_id) projected = phi3_vision_hd_feature_transform( image_features_per_crop, np.asarray(processed["image_sizes"]), From e64f4be9185702d349b185e251c9423e67566987 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sun, 26 Jul 2026 23:37:52 +0000 Subject: [PATCH 07/10] Cover Phi-3.5 projector cache Verify projector checkpoint weights are loaded once per model across golden phases. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tests/e2e_golden_test.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/e2e_golden_test.py b/tests/e2e_golden_test.py index aa2ab6c2..37ea7ca2 100644 --- a/tests/e2e_golden_test.py +++ b/tests/e2e_golden_test.py @@ -29,6 +29,7 @@ import os import warnings from pathlib import Path +from unittest import mock import numpy as np import pytest @@ -212,6 +213,19 @@ def test_build_mm_prompt_rejects_more_media_than_numbered_tokens(): _build_mm_prompt(processor, "describe", ["a.png", "b.png"], "image") +def test_phi3_vision_projector_weights_are_cached(): + _load_phi3_vision_projector_weights_cached.cache_clear() + weights = object() + with mock.patch( + "mobius.models._phi3_vision_projector.load_phi3_vision_projector_weights", + return_value=weights, + ) as load: + assert _load_phi3_vision_projector_weights_cached("model-id") is weights + assert _load_phi3_vision_projector_weights_cached("model-id") is weights + load.assert_called_once_with("model-id") + _load_phi3_vision_projector_weights_cached.cache_clear() + + def _make_empty_kv_cache( session: OnnxModelSession, config: object, From 8e38cc76bf00e54309cc322cb2f0849f0e453856 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sun, 26 Jul 2026 23:44:42 +0000 Subject: [PATCH 08/10] Keep Phi projector GELU in float32 Use a dtype-matched sqrt scalar and build legacy caches linearly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/_testing/torch_reference.py | 6 +++--- src/mobius/models/_phi3_vision_projector.py | 3 ++- src/mobius/models/_phi3_vision_projector_test.py | 1 + 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/mobius/_testing/torch_reference.py b/src/mobius/_testing/torch_reference.py index 9ab54f33..edcaa5e0 100644 --- a/src/mobius/_testing/torch_reference.py +++ b/src/mobius/_testing/torch_reference.py @@ -40,14 +40,14 @@ def _from_legacy_cache(cls, past_key_values=None): # type: ignore[misc] if not hasattr(transformers.DynamicCache, "to_legacy_cache"): def _to_legacy_cache(self): # type: ignore[misc] - legacy_cache: tuple = () + 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 += ((keys, values),) - return legacy_cache + legacy_cache.append((keys, values)) + return tuple(legacy_cache) transformers.DynamicCache.to_legacy_cache = _to_legacy_cache # type: ignore[method-assign] diff --git a/src/mobius/models/_phi3_vision_projector.py b/src/mobius/models/_phi3_vision_projector.py index 8371728c..ea142951 100644 --- a/src/mobius/models/_phi3_vision_projector.py +++ b/src/mobius/models/_phi3_vision_projector.py @@ -225,7 +225,8 @@ def _apply_image_projection( hidden = tokens @ weights.projection_first_weight.T + weights.projection_first_bias # Exact GELU (erf form), matching torch.nn.GELU()'s default approximate="none". hidden_tensor = torch.from_numpy(hidden) - activated = (0.5 * hidden_tensor * (1.0 + torch.erf(hidden_tensor / np.sqrt(2.0)))).numpy() + sqrt_two = hidden_tensor.new_tensor(2.0).sqrt() + activated = (0.5 * hidden_tensor * (1.0 + torch.erf(hidden_tensor / sqrt_two))).numpy() projected = activated @ weights.projection_second_weight.T + weights.projection_second_bias return projected.astype(np.float32) diff --git a/src/mobius/models/_phi3_vision_projector_test.py b/src/mobius/models/_phi3_vision_projector_test.py index 8ec0e1ea..8994726b 100644 --- a/src/mobius/models/_phi3_vision_projector_test.py +++ b/src/mobius/models/_phi3_vision_projector_test.py @@ -111,6 +111,7 @@ def test_apply_image_projection_matches_manual_linear_gelu_linear() -> None: expected = gelu @ weights.projection_second_weight.T + weights.projection_second_bias assert result.shape == (5, 16) + assert result.dtype == np.float32 np.testing.assert_allclose(result, expected, rtol=1e-5, atol=1e-5) From d1d235ea59c77b9e587e5a02fc27663cdf5b297d Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sun, 26 Jul 2026 23:47:10 +0000 Subject: [PATCH 09/10] Type the Phi CLIP config adapter Cast the structural config view at the CLIP constructor boundary for mypy. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/models/phi3_v.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/mobius/models/phi3_v.py b/src/mobius/models/phi3_v.py index b16698de..798b5f40 100644 --- a/src/mobius/models/phi3_v.py +++ b/src/mobius/models/phi3_v.py @@ -49,7 +49,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import torch from onnxscript import nn @@ -201,7 +201,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( - clip_config, + cast(ArchitectureConfig, clip_config), feature_layer=config.vision.feature_layer, drop_class_token=True, ) From 8fa912e8871770a623d08eefa013f9ef96abee00 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 29 Jul 2026 17:31:38 +0000 Subject: [PATCH 10/10] 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 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/mobius/models/_phi3_vision_projector.py | 49 +++++++++++++++++---- src/mobius/models/clip.py | 43 ++++++++++++++---- src/mobius/models/phi3_v.py | 28 ++++++------ src/mobius/models/phi3_v_test.py | 2 + 4 files changed, 89 insertions(+), 33 deletions(-) diff --git a/src/mobius/models/_phi3_vision_projector.py b/src/mobius/models/_phi3_vision_projector.py index ea142951..713000b5 100644 --- a/src/mobius/models/_phi3_vision_projector.py +++ b/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( diff --git a/src/mobius/models/clip.py b/src/mobius/models/clip.py index 34a1b3d1..7e096386 100644 --- a/src/mobius/models/clip.py +++ b/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 diff --git a/src/mobius/models/phi3_v.py b/src/mobius/models/phi3_v.py index 798b5f40..a687dbfe 100644 --- a/src/mobius/models/phi3_v.py +++ b/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 diff --git a/src/mobius/models/phi3_v_test.py b/src/mobius/models/phi3_v_test.py index 16a6595d..99e5afce 100644 --- a/src/mobius/models/phi3_v_test.py +++ b/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.