[EXPERIMENT][WIP][OpenVINO] Add support for Cohere2 MoE (cohere2_moe)

docs/source/openvino/models.mdx

@@ -38,6 +38,7 @@ Here is the list of the supported architectures :
 - CodeGen2
 - Cohere
 - Cohere2
+- Cohere2 MoE
 - ConvBERT
 - ConvNeXt
 - DBRX

optimum/exporters/openvino/model_configs.py

@@ -98,6 +98,7 @@
     BloomModelPatcher,
     ChatGLMModelPatcher,
     CodeGenModelPatcher,
+    Cohere2MoePatcher,
     CommonImageEmbeddingsModelPatcher,
     DBRXModelPatcher,
     DeciLMModelPatcher,
@@ -218,7 +219,6 @@
     NormalizedVisionConfig,
 )
 
-
 COMMON_TEXT_TASKS = [
     "feature-extraction",
     "fill-mask",
@@ -814,6 +814,16 @@ class Cohere2OpenVINOConfig(LlamaOpenVINOConfig):
     pass
 
 
+@register_in_tasks_manager(
+    "cohere2_moe",
+    *["text-generation", "text-generation-with-past"],
+    library_name="transformers",
+)
+class Cohere2MoeOpenVINOConfig(Cohere2OpenVINOConfig):
+    MIN_TRANSFORMERS_VERSION = "5.8.0"
+    _MODEL_PATCHER = Cohere2MoePatcher
+
+
 @register_in_tasks_manager("qwen", *["text-generation", "text-generation-with-past"])
 class QwenOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig):
     MIN_TRANSFORMERS_VERSION = "4.51.0"

optimum/exporters/openvino/model_patcher.py

@@ -9469,6 +9469,35 @@ def __exit__(self, exc_type, exc_value, traceback):
                 sparse_moe_block.forward = sparse_moe_block._orig_forward
 
 
+class Cohere2MoePatcher(OVDecoderModelPatcher):
+    """
+    Model patcher for cohere2_moe (Cohere2MoeForCausalLM).
+
+    Patches Cohere2MoeExperts.forward with the vectorized lfm2_moe_experts_forward
+    that uses torch.bmm over the fused gate_up_proj [E, 2*I, H] weight,
+    split via chunk(2, dim=-2), to make the MoE expert dispatch traceable
+    for OpenVINO export.
+
+    Requires transformers >= 5.8.0 (PR #46115: cohere2_moe model).
+    """
+
+    def __enter__(self):
+        super().__enter__()
+        if is_transformers_version(">=", "5.8.0"):
+            from transformers.models.cohere2_moe.modeling_cohere2_moe import Cohere2MoeExperts
+
+            self._orig_cohere2moe_experts_forward = Cohere2MoeExperts.forward
+            Cohere2MoeExperts.forward = lfm2_moe_experts_forward
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if is_transformers_version(">=", "5.8.0"):
+            from transformers.models.cohere2_moe.modeling_cohere2_moe import Cohere2MoeExperts
+
+            Cohere2MoeExperts.forward = self._orig_cohere2moe_experts_forward
+        super().__exit__(exc_type, exc_value, traceback)
+
+
 class Qwen3ASRModelPatcher(OVSeq2SeqModelPatcher):
     """
     Model patcher for Qwen3-ASR encoder-decoder export.

optimum/intel/openvino/configuration.py

@@ -40,7 +40,6 @@
     PREDEFINED_VISUAL_LM_DATASETS,
 )
 
-
 if is_nncf_available():
     import nncf
 
@@ -55,6 +54,20 @@ class OVQuantizationMethod(str, Enum):
 
 # Default configs for 4-bit weight quantization
 _DEFAULT_4BIT_WQ_CONFIGS = {
+    # cohere2_moe sparse MoE (128 experts, top-8 sigmoid router, tied embeddings).
+    # The tied embed_tokens/lm_head and the per-layer MoE routers are numerically
+    # sensitive; quantizing them collapses greedy generation quality, so they are
+    # kept in higher precision via ignored_scope. Note: 4-bit weight compression of
+    # the experts still perturbs the top-8 routing enough to noticeably degrade
+    # greedy similarity for this model; int8 is recommended when accuracy is critical.
+    "CohereLabs/North-Mini-Code-1.0": {
+        "bits": 4,
+        "sym": False,
+        "group_size": 64,
+        "ratio": 1.0,
+        "backup_precision": "int8_asym",
+        "ignored_scope": {"patterns": [".*embed_tokens.*", ".*mlp\\.gate/aten::linear.*"]},
+    },
     "databricks/dolly-v2-3b": {
         "bits": 4,
         "sym": False,
@@ -505,6 +518,14 @@ class OVQuantizationMethod(str, Enum):
 
 _DEFAULT_8BIT_WQ_CONFIGS = {
     "Qwen/Qwen2.5-Coder-3B-Instruct": {"bits": 8, "sym": False, "dq_group_size": 128},
+    # cohere2_moe: keep the tied embed_tokens/lm_head and MoE routers out of weight
+    # compression (ignored_scope). With this protection int8 weight compression
+    # recovers near-baseline greedy similarity for this sparse-MoE model.
+    "CohereLabs/North-Mini-Code-1.0": {
+        "bits": 8,
+        "sym": False,
+        "ignored_scope": {"patterns": [".*embed_tokens.*", ".*mlp\\.gate/aten::linear.*"]},
+    },
 }
 
 # Add configs for model id aliases

tests/openvino/test_decoder.py

@@ -37,7 +37,6 @@
 from optimum.intel.pipelines import pipeline as optimum_pipeline
 from optimum.intel.utils.import_utils import is_openvino_version, is_transformers_version
 
-
 if is_transformers_version(">=", "4.55"):
     from transformers import Mxfp4Config
 
@@ -100,6 +99,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
         "minicpm3",
         "arctic",
         "deepseek",
+        "cohere2_moe",
         ## not supporter after v5
         "llama4",
         "bitnet",
@@ -197,6 +197,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
         "dbrx": 2,
         "cohere": 2,
         "cohere2": 2,
+        "cohere2_moe": 2,
         "qwen2": 2,
         "qwen2_moe": 4,
         "arctic": 4,

tests/openvino/utils_tests.py

@@ -177,6 +177,7 @@ def _create_tiny_kokoro_model():
     "convbert": "optimum-intel-internal-testing/tiny-random-ConvBertForSequenceClassification",
     "cohere": "optimum-intel-internal-testing/tiny-random-CohereForCausalLM",
     "cohere2": "optimum-intel-internal-testing/tiny-random-aya-base",
+    "cohere2_moe": "optimum-internal-testing/tiny-random-Cohere2MoeForCausalLM",
     "chatglm": "optimum-intel-internal-testing/tiny-random-chatglm",
     "chatglm4": "optimum-intel-internal-testing/tiny-random-chatglm4",
     "codegen": "optimum-intel-internal-testing/tiny-random-CodeGenForCausalLM",