Fix and test stateless encoder decoders (#1423)

* fix

* test other seq2seq stateless models

Author: IlyasMoutawwakil

optimum/exporters/openvino/model_configs.py

@@ -83,8 +83,6 @@
     BaichuanModelPatcher,
     BlenderbotModelPatcher,
     BlenderbotSmallModelPatcher,
-    BlenderbotSmallStatefulSeq2SeqDecoderPatcher,
-    BlenderbotStatefulSeq2SeqDecoderPatcher,
     BloomModelPatcher,
     ChatGLMModelPatcher,
     CodeGenModelPatcher,
@@ -117,7 +115,6 @@
     MairaImageEmbeddingModelPatcher,
     MambaPatcher,
     MarianModelPatcher,
-    MarianStatefulSeq2SeqDecoderPatcher,
     MiniCPM3Patcher,
     MiniCPMModelPatcher,
     MiniCPMVImageEmbeddingsModelPatcher,
@@ -126,9 +123,9 @@
     MixtralModelPatcher,
     MPTModelPatcher,
     OVDecoderModelPatcher,
+    OVSeq2SeqModelPatcher,
     OVSpeechT5ModelPatcher,
     PegasusModelPatcher,
-    PegasusStatefulSeq2SeqDecoderPatcher,
     PersimmonModelPatcher,
     Phi3ModelPatcher,
     Phi3VisionImageEmbeddingsPatcher,
@@ -144,7 +141,6 @@
     Qwen3MoeModelPatcher,
     QwenModelPatcher,
     SanaTextEncoderModelPatcher,
-    StatefulSeq2SeqDecoderPatcher,
     XverseModelPatcher,
 )
 
@@ -3738,9 +3734,7 @@ class WhisperOpenVINOConfig(WhisperOnnxConfig):
     def patch_model_for_export(
         self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
     ) -> ModelPatcher:
-        if getattr(self, "stateful", False) and self._behavior == ConfigBehavior.DECODER:
-            return StatefulSeq2SeqDecoderPatcher(self, model, model_kwargs)
-        return super().patch_model_for_export(model, model_kwargs)
+        return OVSeq2SeqModelPatcher(self, model, model_kwargs=model_kwargs)
 
     @property
     def inputs(self):
@@ -3764,9 +3758,7 @@ class T5OpenVINOConfig(T5OnnxConfig):
     def patch_model_for_export(
         self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
     ) -> ModelPatcher:
-        if getattr(self, "stateful", False) and self._behavior == ConfigBehavior.DECODER:
-            return StatefulSeq2SeqDecoderPatcher(self, model, model_kwargs)
-        return super().patch_model_for_export(model, model_kwargs)
+        return OVSeq2SeqModelPatcher(self, model, model_kwargs)
 
     @property
     def inputs(self):
@@ -3812,9 +3804,7 @@ class BartOpenVINOConfig(BartOnnxConfig):
     def patch_model_for_export(
         self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
     ) -> ModelPatcher:
-        if getattr(self, "stateful", False) and self._behavior == ConfigBehavior.DECODER:
-            return StatefulSeq2SeqDecoderPatcher(self, model, model_kwargs)
-        return super().patch_model_for_export(model, model_kwargs)
+        return OVSeq2SeqModelPatcher(self, model, model_kwargs)
 
     @property
     def inputs(self):
@@ -4056,8 +4046,6 @@ class BlenderbotOpenVINOConfig(BlenderbotOnnxConfig):
     def patch_model_for_export(
         self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
     ) -> "ModelPatcher":
-        if getattr(self, "stateful", False) and self._behavior == ConfigBehavior.DECODER:
-            return BlenderbotStatefulSeq2SeqDecoderPatcher(self, model, model_kwargs)
         return BlenderbotModelPatcher(self, model, model_kwargs=model_kwargs)
 
     @property
@@ -4084,8 +4072,6 @@ class BlenderbotSmallOpenVINOConfig(BlenderbotSmallOnnxConfig):
     def patch_model_for_export(
         self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
     ) -> "ModelPatcher":
-        if getattr(self, "stateful", False) and self._behavior == ConfigBehavior.DECODER:
-            return BlenderbotSmallStatefulSeq2SeqDecoderPatcher(self, model, model_kwargs)
         return BlenderbotSmallModelPatcher(self, model, model_kwargs=model_kwargs)
 
     @property
@@ -4112,8 +4098,6 @@ class PegasusOpenVINOConfig(PegasusOnnxConfig):
     def patch_model_for_export(
         self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
     ) -> "ModelPatcher":
-        if getattr(self, "stateful", False) and self._behavior == ConfigBehavior.DECODER:
-            return PegasusStatefulSeq2SeqDecoderPatcher(self, model, model_kwargs)
         return PegasusModelPatcher(self, model, model_kwargs=model_kwargs)
 
     @property
@@ -4140,8 +4124,6 @@ class MarianOpenVINOConfig(MarianOnnxConfig):
     def patch_model_for_export(
         self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
     ) -> "ModelPatcher":
-        if getattr(self, "stateful", False) and self._behavior == ConfigBehavior.DECODER:
-            return MarianStatefulSeq2SeqDecoderPatcher(self, model, model_kwargs)
         return MarianModelPatcher(self, model, model_kwargs=model_kwargs)
 
     @property

optimum/exporters/openvino/model_patcher.py

@@ -27,7 +27,7 @@
 from transformers.models.speecht5.modeling_speecht5 import SpeechT5EncoderWithSpeechPrenet
 from transformers.utils import is_tf_available
 
-from optimum.exporters.onnx.base import OnnxConfig
+from optimum.exporters.onnx.base import ConfigBehavior, OnnxConfig
 from optimum.exporters.onnx.model_patcher import (
     UNSUPPORTED_OPS_PATCHING_SPEC,
     DecoderModelPatcher,
@@ -327,7 +327,11 @@ def eager_mask_without_vmap(*args, **kwargs) -> Optional[torch.Tensor]:
     mask = sdpa_mask_without_vmap(*args, allow_is_causal_skip=False, **kwargs)
     # we use torch.finfo(torch.float16).min instead torch.finfo(dtype).min to avoid an overflow but not
     # sure this is the right way to handle this, we are basically pretending that -65,504 is -inf
-    mask = torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), torch.finfo(torch.float16).min)
+    mask = torch.where(
+        mask,
+        torch.tensor(0.0, device=mask.device, dtype=dtype),
+        torch.tensor(torch.finfo(torch.float16).min, device=mask.device, dtype=dtype),
+    )
     return mask
 
 
@@ -4711,52 +4715,77 @@ def __exit__(self, exc_type, exc_value, traceback):
                 layer.attn._attn = layer.attn._orig_attn
 
 
-class StatefulSeq2SeqDecoderPatcher(Seq2SeqModelPatcher):
+class OVSeq2SeqModelPatcher(Seq2SeqModelPatcher):
     def __init__(
         self,
         config: "OnnxConfig",
         model: Union["PreTrainedModel", "TFPreTrainedModel"],
         model_kwargs: Optional[Dict[str, Any]] = None,
     ):
-        model.__orig_forward = model.forward
+        if getattr(config, "stateful", False) and config._behavior == ConfigBehavior.DECODER:
+            model.__orig_forward = model.forward
 
-        @functools.wraps(model.__orig_forward)
-        def patched_forward(*args, **kwargs):
-            from transformers.cache_utils import EncoderDecoderCache
+            @functools.wraps(model.__orig_forward)
+            def patched_forward(*args, **kwargs):
+                from transformers.cache_utils import EncoderDecoderCache
+
+                signature = inspect.signature(self.orig_forward)
+                args, kwargs = override_arguments(args, kwargs, signature, model_kwargs=self.model_kwargs)
+
+                return_legacy_cache = False
+                pkv_in_args = False
+                legacy_pkv = None
+                if "past_key_values" in kwargs:
+                    legacy_pkv = kwargs.pop("past_key_values", None)
+                sign_names = list(signature.parameters.keys())
+                pkv_argument_index = sign_names.index("past_key_values")
+                if legacy_pkv is None and len(args) > pkv_argument_index:
+                    legacy_pkv = args[pkv_argument_index]
+                    pkv_in_args = True
+                if legacy_pkv is not None:
+                    if isinstance(legacy_pkv, EncoderDecoderCache):
+                        legacy_pkv = legacy_pkv.to_legacy_cache()
+                    only_self_cache = [cache_item[:2] for cache_item in legacy_pkv]
+                    pkv = EncoderDecoderCache.from_legacy_cache(only_self_cache)
+                    return_legacy_cache = True
+                    if not pkv_in_args:
+                        kwargs["past_key_values"] = pkv
+                    else:
+                        args[pkv_argument_index] = pkv
 
-            signature = inspect.signature(self.orig_forward)
-            args, kwargs = override_arguments(args, kwargs, signature, model_kwargs=self.model_kwargs)
+                outputs = model.__orig_forward(*args, **kwargs)
+                if return_legacy_cache:
+                    outputs.past_key_values = outputs.past_key_values.to_legacy_cache()
 
-            return_legacy_cache = False
-            pkv_in_args = False
-            legacy_pkv = None
-            if "past_key_values" in kwargs:
-                legacy_pkv = kwargs.pop("past_key_values", None)
-            sign_names = list(signature.parameters.keys())
-            pkv_argument_index = sign_names.index("past_key_values")
-            if legacy_pkv is None and len(args) > pkv_argument_index:
-                legacy_pkv = args[pkv_argument_index]
-                pkv_in_args = True
-            if legacy_pkv is not None:
-                if isinstance(legacy_pkv, EncoderDecoderCache):
-                    legacy_pkv = legacy_pkv.to_legacy_cache()
-                only_self_cache = [cache_item[:2] for cache_item in legacy_pkv]
-                pkv = EncoderDecoderCache.from_legacy_cache(only_self_cache)
-                return_legacy_cache = True
-                if not pkv_in_args:
-                    kwargs["past_key_values"] = pkv
-                else:
-                    args[pkv_argument_index] = pkv
+                return outputs
 
-            outputs = model.__orig_forward(*args, **kwargs)
-            if return_legacy_cache:
-                outputs.past_key_values = outputs.past_key_values.to_legacy_cache()
+            model.forward = patched_forward
 
-            return outputs
+        super().__init__(config, model, model_kwargs)
+
+    def __enter__(self):
+        super().__enter__()
 
-        model.forward = patched_forward
+        if is_transformers_version(">=", "4.53.0"):
+            # for OpenVINO, we use torch.finfo(torch.float16).min instead of torch.finfo(dtype).min
+            # Although I'm not sure this is the right way to handle this, we are basically pretending that -65,504 is -inf
+            ALL_MASK_ATTENTION_FUNCTIONS.register("eager", eager_mask_without_vmap)
 
-        super().__init__(config, model, model_kwargs)
+            # for non-stateful decoder models, we use eager mask without vmap for sdpa as well
+            # to avoid a nan output issue in OpenVINO that only happens in case of non-stateful models
+            if not getattr(self.real_config, "stateful", False):
+                logger.warning(
+                    "Exporting a non-stateful decoder model currently results in a nan output in OpenVINO. "
+                    "There might be a performance impact due to the use of eager mask (floats) instead of sdpa mask (bools). "
+                )
+                ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa", eager_mask_without_vmap)
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        super().__exit__(exc_type, exc_value, traceback)
+
+        if is_transformers_version(">=", "4.53.0"):
+            ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa", sdpa_mask)
+            ALL_MASK_ATTENTION_FUNCTIONS.register("eager", eager_mask)
 
 
 class SanaTextEncoderModelPatcher(ModelPatcher):
@@ -5376,7 +5405,7 @@ def modulewise_unpatch(model, module_cls):
                 modulewise_unpatch(module, module_cls)
 
 
-class BlenderbotModelPatcher(Seq2SeqModelPatcher):
+class BlenderbotModelPatcher(OVSeq2SeqModelPatcher):
     def __enter__(self):
         super().__enter__()
         if is_transformers_version(">=", "4.49.0"):
@@ -5392,7 +5421,7 @@ def __exit__(self, exc_type, exc_value, traceback):
             modulewise_unpatch(self._model, BlenderbotAttention)
 
 
-class BlenderbotSmallModelPatcher(Seq2SeqModelPatcher):
+class BlenderbotSmallModelPatcher(OVSeq2SeqModelPatcher):
     def __enter__(self):
         super().__enter__()
         if is_transformers_version(">=", "4.49.0"):
@@ -5408,15 +5437,7 @@ def __exit__(self, exc_type, exc_value, traceback):
             modulewise_unpatch(self._model, BlenderbotSmallAttention)
 
 
-class BlenderbotStatefulSeq2SeqDecoderPatcher(StatefulSeq2SeqDecoderPatcher, BlenderbotModelPatcher):
-    pass
-
-
-class BlenderbotSmallStatefulSeq2SeqDecoderPatcher(StatefulSeq2SeqDecoderPatcher, BlenderbotSmallModelPatcher):
-    pass
-
-
-class PegasusModelPatcher(Seq2SeqModelPatcher):
+class PegasusModelPatcher(OVSeq2SeqModelPatcher):
     def __enter__(self):
         super().__enter__()
         if is_transformers_version(">=", "4.49.0"):
@@ -5495,11 +5516,7 @@ def __exit__(self, exc_type, exc_value, traceback):
             modulewise_unpatch(self._model, Qwen2MoeSparseMoeBlock)
 
 
-class PegasusStatefulSeq2SeqDecoderPatcher(StatefulSeq2SeqDecoderPatcher, PegasusModelPatcher):
-    pass
-
-
-class MarianModelPatcher(Seq2SeqModelPatcher):
+class MarianModelPatcher(OVSeq2SeqModelPatcher):
     def __enter__(self):
         super().__enter__()
         if is_transformers_version(">=", "4.49.0"):
@@ -5515,10 +5532,6 @@ def __exit__(self, exc_type, exc_value, traceback):
             modulewise_unpatch(self._model, MarianAttention)
 
 
-class MarianStatefulSeq2SeqDecoderPatcher(StatefulSeq2SeqDecoderPatcher, MarianModelPatcher):
-    pass
-
-
 # Adopted from https://github.com/huggingface/transformers/blob/v4.51.3/src/transformers/models/speecht5/modeling_speecht5.py#L698
 # this is a patch to avoid PyTorch FE issue
 # with the same tensor names on input and intermediate tensor for speaker_embeddings

tests/openvino/test_modeling.py

@@ -1960,6 +1960,9 @@ def test_compare_to_transformers(self, model_arch):
         model_id = MODEL_NAMES[model_arch]
         set_seed(SEED)
         ov_model = OVModelForSeq2SeqLM.from_pretrained(model_id, export=True, ov_config=F32_CONFIG)
+        ov_stateless_model = OVModelForSeq2SeqLM.from_pretrained(
+            model_id, export=True, use_cache=False, stateful=False, ov_config=F32_CONFIG
+        )
         expected_stateful = is_transformers_version(">", "4.43") and model_arch in self.SUPPORT_STATEFUL
         self.assertEqual(ov_model.decoder.stateful, expected_stateful)
         self.assertEqual(model_has_state(ov_model.decoder.model), expected_stateful)
@@ -1977,6 +1980,7 @@ def test_compare_to_transformers(self, model_arch):
         decoder_start_token_id = transformers_model.config.decoder_start_token_id if model_arch != "mbart" else 2
         decoder_inputs = {"decoder_input_ids": torch.ones((1, 1), dtype=torch.long) * decoder_start_token_id}
         ov_outputs = ov_model(**tokens, **decoder_inputs)
+        ov_stateless_outputs = ov_stateless_model(**tokens, **decoder_inputs)
 
         self.assertTrue("logits" in ov_outputs)
         self.assertIsInstance(ov_outputs.logits, torch.Tensor)
@@ -1985,6 +1989,7 @@ def test_compare_to_transformers(self, model_arch):
             transformers_outputs = transformers_model(**tokens, **decoder_inputs)
         # Compare tensor outputs
         self.assertTrue(torch.allclose(ov_outputs.logits, transformers_outputs.logits, atol=5e-3))
+        self.assertTrue(torch.allclose(ov_stateless_outputs.logits, transformers_outputs.logits, atol=5e-3))
         gen_config = GenerationConfig(
             max_new_tokens=10,
             min_new_tokens=10,
@@ -1997,8 +2002,11 @@ def test_compare_to_transformers(self, model_arch):
         generated_tokens = transformers_model.generate(**tokens, generation_config=gen_config)
         set_seed(SEED)
         ov_generated_tokens = ov_model.generate(**tokens, generation_config=gen_config)
+        set_seed(SEED)
+        ov_stateless_generated_tokens = ov_stateless_model.generate(**tokens, generation_config=gen_config)
 
         self.assertTrue(torch.equal(generated_tokens, ov_generated_tokens))
+        self.assertTrue(torch.equal(generated_tokens, ov_stateless_generated_tokens))
 
         del transformers_model
         del ov_model
@@ -2850,6 +2858,9 @@ def test_compare_to_transformers(self, model_arch):
         model_id = MODEL_NAMES[model_arch]
         transformers_model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id)
         ov_model = OVModelForSpeechSeq2Seq.from_pretrained(model_id, export=True, ov_config=F32_CONFIG)
+        ov_model_stateless = OVModelForSpeechSeq2Seq.from_pretrained(
+            model_id, export=True, ov_config=F32_CONFIG, stateful=False
+        )
         self.assertIsInstance(ov_model.config, PretrainedConfig)
         # whisper cache class support implemented in 4.43
         expected_stateful = is_transformers_version(">", "4.43")
@@ -2874,9 +2885,13 @@ def test_compare_to_transformers(self, model_arch):
                 decoder_inputs = {"decoder_input_ids": np.ones((1, 1), dtype=np.int64) * decoder_start_token_id}
 
             ov_outputs = ov_model(**features, **decoder_inputs)
+            ov_stateless_outputs = ov_model_stateless(**features, **decoder_inputs)
             self.assertIn("logits", ov_outputs)
             # Compare tensor outputs
             self.assertTrue(torch.allclose(torch.Tensor(ov_outputs.logits), transformers_outputs.logits, atol=1e-3))
+            self.assertTrue(
+                torch.allclose(torch.Tensor(ov_stateless_outputs.logits), transformers_outputs.logits, atol=1e-3)
+            )
 
         generate_kwrgs = {}
         if is_transformers_version(">=", "4.50"):
@@ -2894,8 +2909,13 @@ def test_compare_to_transformers(self, model_arch):
         generated_tokens = transformers_model.generate(**pt_features, generation_config=gen_config, **generate_kwrgs)
         set_seed(SEED)
         ov_generated_tokens = ov_model.generate(**pt_features, generation_config=gen_config, **generate_kwrgs)
+        set_seed(SEED)
+        ov_stateless_generated_tokens = ov_model_stateless.generate(
+            **pt_features, generation_config=gen_config, **generate_kwrgs
+        )
 
         self.assertTrue(torch.equal(generated_tokens, ov_generated_tokens))
+        self.assertTrue(torch.equal(generated_tokens, ov_stateless_generated_tokens))
 
         del transformers_model
         del ov_model