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/_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/_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/_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/_testing/torch_reference.py b/src/mobius/_testing/torch_reference.py index 00b89453..edcaa5e0 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 = [] + for layer in getattr(self, "layers", []): + keys = getattr(layer, "keys", None) + values = getattr(layer, "values", None) + if keys is None or values is None: + continue + legacy_cache.append((keys, values)) + return tuple(legacy_cache) + + transformers.DynamicCache.to_legacy_cache = _to_legacy_cache # type: ignore[method-assign] + + if not hasattr(transformers.DynamicCache, "seen_tokens"): + + def _seen_tokens(self): # type: ignore[misc] + # Legacy alias for the number of tokens the cache has seen so + # far, which for a full (non-windowed) cache equals the current + # sequence length of layer 0. + return self.get_seq_length(0) + + transformers.DynamicCache.seen_tokens = property(_seen_tokens) # type: ignore[assignment] + + if not hasattr(transformers.DynamicCache, "get_usable_length"): + + def _get_usable_length(self, new_seq_length: int, layer_idx: int = 0) -> int: # type: ignore[misc] + # Mirrors the removed transformers implementation: the usable + # length is the number of tokens ALREADY cached for this layer + # (the past length), NOT including the incoming ``new_seq_length``. + # Returning ``previous + new`` double-counts the query tokens and + # makes callers (e.g. the Phi-3.5 bundled modeling code) build a + # causal mask of the wrong size. + return self.get_seq_length(layer_idx) + + transformers.DynamicCache.get_usable_length = _get_usable_length # type: ignore[method-assign] + + def _fix_nemotron_h_init_weights(model: torch.nn.Module, model_id: str) -> None: """Restore Mamba2 params from checkpoint after HF clobbers them. @@ -97,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 @@ -190,13 +247,32 @@ def _get_usable_length(self, new_seq_length: int, layer_idx: int = 0) -> int: # # Some trust_remote_code VLMs (e.g. Phi-3-Vision) are registered as # AutoModelForCausalLM, not AutoModelForImageTextToText. Fall back # gracefully so golden generation works for both. - model_kwargs = dict(config=config, dtype=dtype, device_map=device, trust_remote_code=True) + # + # The weight-dtype keyword was renamed from ``torch_dtype`` to ``dtype`` in + # transformers 5.x; support both so offline golden generation also works on + # the older transformers (e.g. 4.43) required by 4.x-era remote-code models. + base_kwargs = dict(config=config, device_map=device, trust_remote_code=True) + + def _load_from_pretrained(auto_cls): + try: + return auto_cls.from_pretrained(model_id, dtype=dtype, **base_kwargs) + except TypeError: + return auto_cls.from_pretrained(model_id, torch_dtype=dtype, **base_kwargs) + try: - model = transformers.AutoModelForImageTextToText.from_pretrained( - model_id, **model_kwargs - ) - except (ValueError, KeyError): - model = transformers.AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs) + image_text_to_text_cls = transformers.AutoModelForImageTextToText + except AttributeError: + # Older transformers (e.g. 4.43, used for offline golden generation of + # 4.x-era remote-code checkpoints) predate AutoModelForImageTextToText. + image_text_to_text_cls = None + model = None + if image_text_to_text_cls is not None: + try: + model = _load_from_pretrained(image_text_to_text_cls) + except (ValueError, KeyError): + model = None + if model is None: + model = _load_from_pretrained(transformers.AutoModelForCausalLM) model.eval() diff --git a/src/mobius/models/_phi3_vision_projector.py b/src/mobius/models/_phi3_vision_projector.py new file mode 100644 index 00000000..713000b5 --- /dev/null +++ b/src/mobius/models/_phi3_vision_projector.py @@ -0,0 +1,341 @@ +# 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 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 +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 json +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). +_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" +_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 +# 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_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 sorted(glob.glob(os.path.join(model_id_or_directory, "*.safetensors"))) + + from huggingface_hub import hf_hub_download + from huggingface_hub.errors import EntryNotFoundError, LocalEntryNotFoundError + + 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( + 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 + + shard_paths = _resolve_checkpoint_shard_paths(model_id_or_directory) + if not shard_paths: + raise FileNotFoundError( + f"No .safetensors shards found for checkpoint: {model_id_or_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( + "Missing Phi-3.5-Vision projector weights in " + f"{model_id_or_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)``. + """ + 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) + 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) + + +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}" + ) + 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) + + # 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): + 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) + 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..8994726b --- /dev/null +++ b/src/mobius/models/_phi3_vision_projector_test.py @@ -0,0 +1,241 @@ +# 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) + assert result.dtype == np.float32 + 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 + ) + + +@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 + 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/src/mobius/models/clip.py b/src/mobius/models/clip.py index e195add2..7e096386 100644 --- a/src/mobius/models/clip.py +++ b/src/mobius/models/clip.py @@ -5,36 +5,79 @@ 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 Embedding, LayerNorm -from mobius.components._conv import Conv2d -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. + + 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.""" - def __init__(self, config: ArchitectureConfig): + def __init__(self, config: _CLIPVisionConfig): super().__init__() hidden_size = config.hidden_size patch_size = config.patch_size 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 +101,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 +112,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'] @@ -86,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) @@ -110,23 +158,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: _CLIPVisionConfig, + *, + 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,16 +248,36 @@ 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( 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 @@ -175,8 +308,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 @@ -214,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/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..a687dbfe 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,16 +60,21 @@ 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 -# Phi-3-Vision: <|image|> token id -_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): @@ -140,74 +147,95 @@ 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.* - """ - 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." + → vision_tower.* (via the shared CLIP renamer) + The ``img_projection`` and ``sub_GN``/``glb_GN`` tensors are dropped + (host-side HD transform). + """ + clip_state_dict: dict[str, torch.Tensor] = {} 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 - + 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 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): @@ -215,13 +243,20 @@ def __init__(self, config: ArchitectureConfig): self.embed_tokens = Embedding( config.vocab_size, config.hidden_size, config.pad_token_id ) - self.image_token_id = ( - config.image_token_id if config.image_token_id is not None else _IMAGE_TOKEN_ID - ) + self.hidden_size = config.hidden_size 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. @@ -235,7 +270,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( @@ -295,32 +340,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 self.vision_encoder.preprocess_weights(state_dict).items(): + renamed["vision_encoder." + key] = value 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 - 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. + 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 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..99e5afce --- /dev/null +++ b/src/mobius/models/phi3_v_test.py @@ -0,0 +1,135 @@ +# 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 + 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. + assert all(k.startswith("vision_tower.") for k in renamed) 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/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/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/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/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/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 diff --git a/tests/e2e_golden_test.py b/tests/e2e_golden_test.py index 713fd05e..37ea7ca2 100644 --- a/tests/e2e_golden_test.py +++ b/tests/e2e_golden_test.py @@ -25,9 +25,11 @@ from __future__ import annotations import dataclasses +import functools import os import warnings from pathlib import Path +from unittest import mock import numpy as np import pytest @@ -50,6 +52,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" @@ -148,12 +160,72 @@ 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) + ): + 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)) + ) + 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 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 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, @@ -709,6 +781,73 @@ 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 ( + 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_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:] + ) + 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_cached(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 +881,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 +892,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 +1045,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 +1065,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: 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, + )