improve disabling trace check

Author: eaidova

optimum/exporters/openvino/convert.py

@@ -67,6 +67,7 @@
     OV_XML_FILE_NAME,
     _get_input_info,
     _get_open_clip_submodels_fn_and_export_configs,
+    allow_skip_tracing_check,
     clear_class_registry,
     remove_none_from_dummy_inputs,
     save_config,
@@ -436,7 +437,7 @@ def ts_patched_forward(*args, **kwargs):
             patcher.patched_forward = ts_patched_forward
 
             ts_decoder_kwargs = {}
-            if is_openvino_version(">=", "2025.0"):
+            if library_name == "diffusers" or allow_skip_tracing_check(model) and is_openvino_version(">=", "2025.0"):
                 ts_decoder_kwargs["trace_kwargs"] = {"check_trace": False}
 
             with patcher:

optimum/exporters/openvino/model_patcher.py

@@ -3631,24 +3631,18 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
     else:
         q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
     q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)
-    q_nope, q_pe = torch.split(
-        q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
-    )
+    q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
 
     compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
-    compressed_kv, k_pe = torch.split(
-        compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
-    )
+    compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
     k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)
     kv = (
         self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
         .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
         .transpose(1, 2)
     )
 
-    k_nope, value_states = torch.split(
-        kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1
-    )
+    k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
     kv_seq_len = value_states.shape[-2]
     if past_key_value is not None:
         if self.layer_idx is None:
@@ -3704,7 +3698,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
 
     attn_output = self.o_proj(attn_output)
 
-
     return attn_output, None, past_key_value
 
 

optimum/exporters/openvino/utils.py

@@ -297,3 +297,12 @@ def save_preprocessors(
                 logger.error(f"Saving {type(processor)} failed with {ex}")
     else:
         maybe_save_preprocessors(model_name_or_path, output, trust_remote_code=trust_remote_code)
+
+
+SKIP_CHECK_TRACE_MODELS = ("deepseek", "deepseek-v2", "deepseek-v3")
+
+
+def allow_skip_tracing_check(model):
+    config = getattr(model, "config", {})
+    model_type = getattr(config, "model_type", "").replace("_", "-")
+    return model_type in SKIP_CHECK_TRACE_MODELS

tests/openvino/test_modeling.py

@@ -995,10 +995,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
 
         # openvino 2025.0 required for disabling check_trace
         if is_openvino_version(">=", "2025.0"):
-            SUPPORTED_ARCHITECTURES += (
-                "deepseek-v2",
-                "deepseek-v3"
-            )
+            SUPPORTED_ARCHITECTURES += ("deepseek-v2", "deepseek-v3")
 
         # gptq and awq install disabled for windows test environment
         if platform.system() != "Windows":
@@ -1029,7 +1026,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
         "decilm",
         "minicpm3",
         "deepseek-v2",
-        "deepseek-v3"
+        "deepseek-v3",
     )
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES)