diff --git a/docs/source/openvino/models.mdx b/docs/source/openvino/models.mdx index db2701749f..1c8fb475cd 100644 --- a/docs/source/openvino/models.mdx +++ b/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 diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index fcacba0534..67fa40adcb 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/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" diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 30f4903922..6ddeee8707 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/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. diff --git a/optimum/intel/openvino/configuration.py b/optimum/intel/openvino/configuration.py index aa46ca44ee..d8a9a9aea0 100644 --- a/optimum/intel/openvino/configuration.py +++ b/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 diff --git a/tests/openvino/test_decoder.py b/tests/openvino/test_decoder.py index 9a17d0bdac..7de723e585 100644 --- a/tests/openvino/test_decoder.py +++ b/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, diff --git a/tests/openvino/utils_tests.py b/tests/openvino/utils_tests.py index f0e32b265b..5521559fae 100644 --- a/tests/openvino/utils_tests.py +++ b/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",