ORT Seq2Seq modeling refactorization (#34)

This PR:
- Creates and separates the testing of seq2seq / encoder-decoder models
(text, speech and vision).
- Enables full testing, including batched inference with padded inputs.
logits, last hidden state and pkv matching assertions. and beam search,
random sampling...
- Fixes issues emerging/detected when running the full testing suite,
like missing dynamic axis, missing attention masks in the
encoder/decoder, untested model types for some tasks, unpatched models,
etc..
- Rewrites and simplifies the seq2seq modeling code, removing a lot of
hackiness around pkv (self attn vs cross attn) and making the forward
pass (especially the seq2seqdecoder forward) more linear, finally
enabling cases like use_merged=True and use_cache=False, which weren't
supported before.
- Rewrite and simplifies the filename inference when loading a seq2seq
model.
- Documents behaviors like m2m100 consuming the RNG state, bigbird
pegasus switching between attention types.
- Enables Moonshine speech recognition and patch it for dynamic audio
input length (only supported static).
- Removes some deprecated stuff, like onnx config with loss which was
only relevant for ort training.
- Correct the output dynamic axis for different ctc models.
- [maybe more 🥲]

Author: IlyasMoutawwakil

optimum/exporters/onnx/__init__.py

@@ -20,7 +20,7 @@
 
 
 _import_structure = {
-    "base": ["OnnxConfig", "OnnxConfigWithLoss", "OnnxConfigWithPast", "OnnxSeq2SeqConfigWithPast"],
+    "base": ["OnnxConfig", "OnnxConfigWithPast", "OnnxSeq2SeqConfigWithPast"],
     "config": ["TextDecoderOnnxConfig", "TextEncoderOnnxConfig", "TextSeq2SeqOnnxConfig"],
     "convert": [
         "export",
@@ -40,12 +40,7 @@
 
 if TYPE_CHECKING:
     from optimum.exporters.onnx.__main__ import main_export
-    from optimum.exporters.onnx.base import (
-        OnnxConfig,
-        OnnxConfigWithLoss,
-        OnnxConfigWithPast,
-        OnnxSeq2SeqConfigWithPast,
-    )
+    from optimum.exporters.onnx.base import OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast
     from optimum.exporters.onnx.config import TextDecoderOnnxConfig, TextEncoderOnnxConfig, TextSeq2SeqOnnxConfig
     from optimum.exporters.onnx.convert import (
         export,

optimum/exporters/onnx/base.py

@@ -22,6 +22,8 @@
 
 from optimum.exporters.onnx.base import ConfigBehavior, OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast
 from optimum.exporters.onnx.constants import ONNX_DECODER_MERGED_NAME, ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME
+from optimum.exporters.tasks import TasksManager
+from optimum.onnx import merge_decoders
 from optimum.utils import (
     DummyAudioInputGenerator,
     DummyBboxInputGenerator,
@@ -31,19 +33,13 @@
     DummySeq2SeqPastKeyValuesGenerator,
     DummyTextInputGenerator,
     DummyVisionInputGenerator,
-    is_diffusers_available,
     logging,
 )
 
 
-# TODO : moved back onnx imports applied in https://github.com/huggingface/optimum/pull/2114/files after refactorization
-
-
 if TYPE_CHECKING:
     from transformers import PretrainedConfig, PreTrainedModel
 
-    if is_diffusers_available():
-        from diffusers import ModelMixin
 
 logger = logging.get_logger(__name__)
 
@@ -110,7 +106,7 @@ def outputs(self) -> dict[str, dict[int, str]]:
     def post_process_exported_models(
         self,
         path: Path,
-        models_and_onnx_configs: dict[str, tuple[PreTrainedModel | ModelMixin, OnnxConfig]],
+        models_and_onnx_configs: dict[str, tuple[PreTrainedModel, OnnxConfig]],
         onnx_files_subpaths: list[str],
     ):
         models_and_onnx_configs, onnx_files_subpaths = super().post_process_exported_models(
@@ -119,8 +115,6 @@ def post_process_exported_models(
 
         # Attempt to merge only if the decoder-only was exported separately without/with past
         if self.use_past is True and len(models_and_onnx_configs) == 2:
-            from optimum.onnx import merge_decoders
-
             decoder_path = Path(path, onnx_files_subpaths[0])
             decoder_with_past_path = Path(path, onnx_files_subpaths[1])
             decoder_merged_path = Path(path, ONNX_DECODER_MERGED_NAME + ".onnx")
@@ -171,35 +165,19 @@ class TextSeq2SeqOnnxConfig(OnnxSeq2SeqConfigWithPast):
         DummySeq2SeqPastKeyValuesGenerator,
     )
 
-    @property
-    def torch_to_onnx_input_map(self) -> dict[str, str]:
-        if self._behavior is ConfigBehavior.DECODER:
-            return {
-                "decoder_input_ids": "input_ids",
-                "encoder_outputs": "encoder_hidden_states",
-                "attention_mask": "encoder_attention_mask",
-            }
-        return {}
-
     @property
     def inputs(self) -> dict[str, dict[int, str]]:
         common_inputs = {}
-        if self._behavior is not ConfigBehavior.DECODER:
+        if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}:
             common_inputs["input_ids"] = {0: "batch_size", 1: "encoder_sequence_length"}
-
+        else:
+            common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"}
         common_inputs["attention_mask"] = {0: "batch_size", 1: "encoder_sequence_length"}
 
-        if self._behavior is not ConfigBehavior.ENCODER:
+        if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}:
+            common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"}
             if self.use_past_in_inputs:
-                # TODO: validate the axis name for attention_mask
-                # common_inputs["attention_mask"][1] = "past_encoder_sequence_length + sequence_length"
-                common_inputs["decoder_input_ids"] = {0: "batch_size"}
                 self.add_past_key_values(common_inputs, direction="inputs")
-            else:
-                common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"}
-
-        if self._behavior is ConfigBehavior.DECODER:
-            common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"}
 
         return common_inputs
 
@@ -260,31 +238,18 @@ class AudioToTextOnnxConfig(OnnxSeq2SeqConfigWithPast):
     def inputs(self) -> dict[str, dict[int, str]]:
         common_inputs = {}
 
-        if self._behavior is not ConfigBehavior.DECODER:
+        if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}:
             common_inputs["input_features"] = {0: "batch_size", 1: "feature_size", 2: "encoder_sequence_length"}
+        else:
+            common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"}
 
-        if self._behavior is not ConfigBehavior.ENCODER:
+        if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}:
+            common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"}
             if self.use_past_in_inputs:
-                common_inputs["decoder_input_ids"] = {0: "batch_size"}
                 self.add_past_key_values(common_inputs, direction="inputs")
-            else:
-                common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"}
-
-        if self._behavior is ConfigBehavior.DECODER:
-            common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"}
 
         return common_inputs
 
-    @property
-    def torch_to_onnx_input_map(self) -> dict[str, str]:
-        if self._behavior is ConfigBehavior.DECODER:
-            return {
-                "decoder_input_ids": "input_ids",
-                "encoder_outputs": "encoder_hidden_states",
-                "attention_mask": "encoder_attention_mask",
-            }
-        return {}
-
 
 class EncoderDecoderBaseOnnxConfig(OnnxSeq2SeqConfigWithPast):
     DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator,)
@@ -313,8 +278,6 @@ def __init__(
             legacy=legacy,
         )
 
-        from optimum.exporters.tasks import TasksManager
-
         self.is_decoder_with_past = False
 
         # Set up the encoder ONNX config.
@@ -382,41 +345,19 @@ def __init__(
     @property
     def inputs(self) -> dict[str, dict[int, str]]:
         common_inputs = {}
-        if self._behavior is not ConfigBehavior.DECODER:
+        if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}:
             common_inputs["input_ids"] = {0: "batch_size", 1: "encoder_sequence_length"}
-
+        else:
+            common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"}
         common_inputs["attention_mask"] = {0: "batch_size", 1: "encoder_sequence_length"}
 
-        if self._behavior is not ConfigBehavior.ENCODER:
-            # TODO: it is likely this pop() is unwanted as we then always hit
-            # https://github.com/huggingface/transformers/blob/v4.26.0/src/transformers/models/t5/modeling_t5.py#L965-L969
-            common_inputs.pop("attention_mask")
-
-            if self.use_past_in_inputs:
-                # TODO: validate the axis name for attention_mask
-                # common_inputs["attention_mask"][1] = "past_encoder_sequence_length + sequence_length"
-                common_inputs["decoder_input_ids"] = {0: "batch_size"}
-            else:
-                common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"}
-
+        if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}:
+            common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"}
             if self.use_past_in_inputs:
                 self.add_past_key_values(common_inputs, direction="inputs")
 
-        if self._behavior is ConfigBehavior.DECODER:
-            common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"}
-
         return common_inputs
 
-    @property
-    def torch_to_onnx_input_map(self) -> dict[str, str]:
-        if self._behavior is ConfigBehavior.DECODER:
-            return {
-                "decoder_input_ids": "input_ids",
-                "encoder_outputs": "encoder_hidden_states",
-                "attention_mask": "encoder_attention_mask",
-            }
-        return {}
-
     def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], direction: str):
         if self.is_decoder_with_past:
             return self._decoder_onnx_config.add_past_key_values(inputs_or_outputs, direction)
@@ -429,26 +370,34 @@ def flatten_output_collection_property(self, name: str, field: Iterable[Any]) ->
         return self._decoder_onnx_config.flatten_output_collection_property(name, field)
 
     def generate_dummy_inputs_for_validation(
-        self, reference_model_inputs: dict[str, Any], onnx_input_names: list[str] | None = None
+        self, reference_model_inputs: dict[str, Any], onnx_input_names: list[str]
     ) -> dict[str, Any]:
         if self._behavior is ConfigBehavior.ENCODER:
-            return self._encoder_onnx_config.generate_dummy_inputs_for_validation(reference_model_inputs)
+            return self._encoder_onnx_config.generate_dummy_inputs_for_validation(
+                reference_model_inputs, onnx_input_names
+            )
         else:
             if self._behavior is ConfigBehavior.DECODER:
-                reference_model_inputs["input_ids"] = reference_model_inputs.pop("decoder_input_ids")
-
-            if "encoder_outputs" in reference_model_inputs:
-                if "encoder_hidden_states" in onnx_input_names:
-                    reference_model_inputs["encoder_hidden_states"] = reference_model_inputs.pop("encoder_outputs")[0]
-                else:
-                    reference_model_inputs.pop("encoder_outputs")
-
-            return self._decoder_onnx_config.generate_dummy_inputs_for_validation(reference_model_inputs)
+                if "decoder_input_ids" in reference_model_inputs:
+                    reference_model_inputs["input_ids"] = reference_model_inputs.pop("decoder_input_ids")
+                if "attention_mask" in reference_model_inputs:
+                    reference_model_inputs["encoder_attention_mask"] = reference_model_inputs.pop("attention_mask")
+                if "encoder_outputs" in reference_model_inputs:
+                    if "encoder_hidden_states" in onnx_input_names:
+                        reference_model_inputs["encoder_hidden_states"] = reference_model_inputs.pop(
+                            "encoder_outputs"
+                        )[0]
+                    else:
+                        reference_model_inputs.pop("encoder_outputs")
+
+            return self._decoder_onnx_config.generate_dummy_inputs_for_validation(
+                reference_model_inputs, onnx_input_names
+            )
 
     def post_process_exported_models(
         self,
         path: Path,
-        models_and_onnx_configs: dict[str, tuple[PreTrainedModel | ModelMixin, OnnxConfig]],
+        models_and_onnx_configs: dict[str, tuple[PreTrainedModel, OnnxConfig]],
         onnx_files_subpaths: list[str],
     ):
         models_and_onnx_configs, onnx_files_subpaths = super().post_process_exported_models(

optimum/exporters/onnx/config.py

@@ -333,9 +333,7 @@ def _run_validation(
 
     # Possibly edit the input for the onnxruntime.InferenceSession, this is for example the case for merged
     # models where the input `use_cache_branch` is added
-    reference_ort_inputs = config.generate_dummy_inputs_for_validation(
-        reference_model_inputs, onnx_input_names=onnx_input_names
-    )
+    reference_ort_inputs = config.generate_dummy_inputs_for_validation(reference_model_inputs, onnx_input_names)
 
     # We flatten potential collection of inputs (i.e. past_keys)
     onnx_inputs = {}

optimum/exporters/onnx/convert.py

@@ -13,7 +13,12 @@
 # limitations under the License.
 from __future__ import annotations
 
-from optimum.utils import DummyPastKeyValuesGenerator, NormalizedTextConfig, is_transformers_version
+from optimum.utils import (
+    DummyAudioInputGenerator,
+    DummyPastKeyValuesGenerator,
+    NormalizedTextConfig,
+    is_transformers_version,
+)
 
 
 class GPTBigCodeDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator):
@@ -64,3 +69,25 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int
             ]
 
         return pkv
+
+
+class DummyMoonshineAudioInputGenerator(DummyAudioInputGenerator):
+    SUPPORTED_INPUT_NAMES = ("input_values", "attention_mask")
+
+    def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
+        if input_name == "input_values":  # raw waveform
+            return self.random_float_tensor(
+                shape=[self.batch_size, self.sequence_length],
+                min_value=-1,
+                max_value=1,
+                framework=framework,
+                dtype=float_dtype,
+            )
+        elif input_name == "attention_mask":  # attention mask
+            return self.random_mask_tensor(
+                shape=[self.batch_size, self.sequence_length],
+                framework=framework,
+                dtype=int_dtype,
+            )
+        else:
+            raise ValueError(f"Unsupported input name: {input_name}")