phi3moe support (#1215)

* phi3moe support

* add tests

* use transformers code

Author: eaidova

docs/source/openvino/models.mdx

@@ -103,6 +103,7 @@ Here is the list of the supported architectures :
 - Persimmon
 - Phi
 - Phi3
+- Phi3.5-MoE
 - Phi3Vision
 - Pix2Struct
 - PoolFormer

optimum/exporters/openvino/model_configs.py

@@ -109,6 +109,7 @@
     PersimmonModelPatcher,
     Phi3ModelPatcher,
     Phi3VisionImageEmbeddingsPatcher,
+    PhiMoEModelPatcher,
     Qwen2_5_VLVisionEmbMergerPatcher,
     Qwen2VLLanguageModelPatcher,
     Qwen2VLVisionEmbMergerPatcher,
@@ -737,6 +738,26 @@ def patch_model_for_export(
         return Phi3ModelPatcher(self, model, model_kwargs=model_kwargs)
 
 
+@register_in_tasks_manager(
+    "phimoe",
+    *[
+        "feature-extraction",
+        "feature-extraction-with-past",
+        "text-generation",
+        "text-generation-with-past",
+        "text-classification",
+    ],
+    library_name="transformers",
+)
+class PhiMoEOpenVINOConfig(Phi3OpenVINOConfig):
+    MIN_TRANSFORMERS_VERSION = "4.46.0"
+
+    def patch_model_for_export(
+        self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
+    ) -> "ModelPatcher":
+        return PhiMoEModelPatcher(self, model, model_kwargs=model_kwargs)
+
+
 @register_in_tasks_manager(
     "phi",
     *[

optimum/exporters/openvino/model_patcher.py

@@ -1598,7 +1598,10 @@ def __enter__(self):
                 layer.self_attn.forward = types.MethodType(_phi3_self_attn_sdpa_forward, layer.self_attn)
                 layer.self_attn._orig_forward = orig_self_attn_fwd
 
-            if hasattr(layer.self_attn, "rotary_emb") and layer.self_attn.rotary_emb.inv_freq is None:
+            if (
+                hasattr(layer.self_attn, "rotary_emb")
+                and getattr(layer.self_attn.rotary_emb, "inv_freq", None) is None
+            ):
                 rotary_emb = layer.self_attn.rotary_emb
                 layer.self_attn.rotary_emb.inv_freq = 1.0 / (
                     rotary_emb.base ** (torch.arange(0, rotary_emb.dim, 2, dtype=torch.int64).float() / rotary_emb.dim)
@@ -1615,6 +1618,69 @@ def __exit__(self, exc_type, exc_value, traceback):
                 layer.self_attn.forward = layer.self_attn._orig_forward
 
 
+# Modified from https://github.com/huggingface/transformers/blob/v4.50.2/src/transformers/models/phimoe/modeling_phimoe.py#L756
+# removed usage nonfriendly for tracing operation continue
+def _phi_moe_sparse_moe_block_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    from transformers.models.phimoe.modeling_phimoe import sparsemixer
+
+    batch_size, sequence_length, hidden_dim = hidden_states.shape
+    if self.training and self.input_jitter_noise > 0:
+        hidden_states *= torch.empty_like(hidden_states).uniform_(
+            1.0 - self.input_jitter_noise, 1.0 + self.input_jitter_noise
+        )
+    hidden_states = hidden_states.view(-1, hidden_dim)
+    router_logits = self.gate(hidden_states)
+
+    routing_weights, selected_experts = sparsemixer(
+        router_logits,
+        jitter_eps=self.router_jitter_noise,
+        training=self.training,
+    )
+
+    final_hidden_states = torch.zeros(
+        (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
+    )
+
+    # One hot encode the selected experts to create an expert mask
+    # this will be used to easily index which expert is going to be sollicitated
+    expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
+
+    # Loop over all available experts in the model and perform the computation on each expert
+    for expert_idx in range(self.num_experts):
+        expert_layer = self.experts[expert_idx]
+        idx, top_x = torch.where(expert_mask[expert_idx])
+
+        # if top_x.shape[0] == 0:
+        #     continue
+
+        # Index the correct hidden states and compute the expert hidden state for
+        # the current expert. We need to make sure to multiply the output hidden
+        # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
+        current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
+        current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]
+
+        # However `index_add_` only support torch tensors for indexing so we'll use
+        # the `top_x` tensor here.
+        final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+    final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+    return final_hidden_states, router_logits
+
+
+class PhiMoEModelPatcher(Phi3ModelPatcher):
+    def __enter__(self):
+        super().__enter__()
+        for layer in self._model.model.layers:
+            layer.block_sparse_moe._orig_forward = layer.block_sparse_moe.forward
+            layer.block_sparse_moe.forward = types.MethodType(
+                _phi_moe_sparse_moe_block_forward, layer.block_sparse_moe
+            )
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        super().__exit__(exc_type, exc_value, traceback)
+        for layer in self._model.model.layers:
+            layer.block_sparse_moe.forward = layer.block_sparse_moe._orig_forward
+
+
 def _aquila_self_attn_sdpa_forward(
     self,
     hidden_states: torch.Tensor,

tests/openvino/test_modeling.py

@@ -1032,7 +1032,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
         SUPPORTED_ARCHITECTURES += ("granite", "granite-moe")
 
     if is_transformers_version(">=", "4.46.0"):
-        SUPPORTED_ARCHITECTURES += ("glm", "mistral-nemo", "minicpm3")
+        SUPPORTED_ARCHITECTURES += ("glm", "mistral-nemo", "minicpm3", "phi3-moe")
         # openvino 2025.0 required for disabling check_trace
         if is_openvino_version(">=", "2025.0"):
             SUPPORTED_ARCHITECTURES += ("deepseek",)

tests/openvino/utils_tests.py

@@ -126,6 +126,7 @@
     "pix2struct": "fxmarty/pix2struct-tiny-random",
     "phi": "echarlaix/tiny-random-PhiForCausalLM",
     "phi3": "Xenova/tiny-random-Phi3ForCausalLM",
+    "phi3-moe": "katuni4ka/phi-3.5-moe-tiny-random",
     "phi3_v": "katuni4ka/tiny-random-phi3-vision",
     "poolformer": "hf-internal-testing/tiny-random-PoolFormerModel",
     "qwen": "katuni4ka/tiny-random-qwen",