huggingface · fxmarty · Mar 31, 2023 · Mar 28, 2023 · Mar 28, 2023 · Mar 28, 2023
diff --git a/optimum/exporters/onnx/base.py b/optimum/exporters/onnx/base.py
@@ -753,8 +753,11 @@ def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], dire
                 or (self._behavior is ConfigBehavior.DECODER and self.use_past is False)
                 or direction == "inputs"
             ):
-                inputs_or_outputs[f"{name}.{i}.encoder.key"] = {0: "batch_size", 2: "encoder_sequence_length"}
-                inputs_or_outputs[f"{name}.{i}.encoder.value"] = {0: "batch_size", 2: "encoder_sequence_length"}
+                # TODO: we only need to call it encoder_sequence_length_out in the merge case - but at torch.onnx.export()
+                # time we have currently no case to check whether we will merge at a later step or not (self.is_merged is
+                # not yet set at this time)
+                inputs_or_outputs[f"{name}.{i}.encoder.key"] = {0: "batch_size", 2: "encoder_sequence_length_out"}
+                inputs_or_outputs[f"{name}.{i}.encoder.value"] = {0: "batch_size", 2: "encoder_sequence_length_out"}
 
     def flatten_past_key_values(self, flattened_output, name, idx, t):
         flattened_output[f"{name}.{idx}.decoder.key"] = t[0]

diff --git a/optimum/onnxruntime/base.py b/optimum/onnxruntime/base.py
diff --git a/optimum/onnxruntime/constants.py b/optimum/onnxruntime/constants.py
@@ -13,6 +13,6 @@
 #  limitations under the License.
 
 ENCODER_ONNX_FILE_PATTERN = r"(.*)?encoder(.*)?\.onnx"
-DECODER_ONNX_FILE_PATTERN = r"(.*)?decoder((?!with_past).)*?\.onnx"
+DECODER_ONNX_FILE_PATTERN = r"(.*)?decoder((?!(with_past|merged)).)*?\.onnx"
 DECODER_WITH_PAST_ONNX_FILE_PATTERN = r"(.*)?decoder(.*)?with_past(.*)?\.onnx"
 DECODER_MERGED_ONNX_FILE_PATTERN = r"(.*)?decoder(.*)?merged(.*)?\.onnx"
diff --git a/optimum/onnxruntime/modeling_decoder.py b/optimum/onnxruntime/modeling_decoder.py
@@ -294,6 +294,8 @@ def _save_pretrained(self, save_directory: Union[str, Path]):
         for src_path, dst_path in zip(src_paths, dst_paths):
             shutil.copyfile(src_path, dst_path)
 
+        self.generation_config.save_pretrained(save_directory)
+
     @classmethod
     def _from_pretrained(
         cls,
@@ -330,6 +332,7 @@ def _from_pretrained(
                     revision=revision,
                 )
                 use_merged = True
+                decoder_path = decoder_merged_path
             except FileNotFoundError as e:
                 if use_merged is True:
                     raise FileNotFoundError(
@@ -339,74 +342,74 @@ def _from_pretrained(
                     )
                 use_merged = False
 
-        if not validate_file_exists(model_id, decoder_file_name, subfolder=subfolder, revision=revision):
-            decoder_without_past_path = ORTModelDecoder.infer_onnx_filename(
-                model_id,
-                [DECODER_ONNX_FILE_PATTERN],
-                "decoder_file_name",
-                subfolder=subfolder,
-                use_auth_token=use_auth_token,
-                revision=revision,
-            )
-        else:
-            decoder_without_past_path = model_path / subfolder / decoder_file_name
-
-        if use_merged is True:
-            decoder_path = decoder_merged_path
-        else:
-            decoder_path = decoder_without_past_path
-
-        decoder_regular_onnx_filenames = ORTModelDecoder._generate_regular_names_for_filename(ONNX_DECODER_NAME)
-        if decoder_path.name not in decoder_regular_onnx_filenames:
-            logger.warning(
-                f"The ONNX file {decoder_path.name} is not a regular name used in optimum.onnxruntime that are {decoder_regular_onnx_filenames}, the "
-                f"{cls.__name__} might not behave as expected."
-            )
-
-        # If the decoder without / with past has been merged, we do not need to look for any additional file
+        decoder_without_past_path = None
         decoder_with_past_path = None
-        if not validate_file_exists(model_id, decoder_with_past_file_name, subfolder=subfolder, revision=revision):
-            try:
-                decoder_with_past_path = ORTModelDecoder.infer_onnx_filename(
+        if use_merged is False:
+            if not validate_file_exists(model_id, decoder_file_name, subfolder=subfolder, revision=revision):
+                decoder_without_past_path = ORTModelDecoder.infer_onnx_filename(
                     model_id,
-                    [DECODER_WITH_PAST_ONNX_FILE_PATTERN],
-                    "decoder_with_past_file_name",
+                    [DECODER_ONNX_FILE_PATTERN],
+                    "decoder_file_name",
                     subfolder=subfolder,
                     use_auth_token=use_auth_token,
                     revision=revision,
                 )
-            except FileNotFoundError as e:
-                if use_cache is True and use_merged is False:
-                    raise FileNotFoundError(
-                        "The parameter `use_cache=True` was passed to ORTModelForCausalLM.from_pretrained()"
-                        " but no ONNX file using past key values could be found in"
-                        f" {str(Path(model_id, subfolder))}, with the error: {e}"
-                    )
-        else:
-            decoder_with_past_path = model_path / subfolder / decoder_with_past_file_name
+            else:
+                decoder_without_past_path = model_path / subfolder / decoder_file_name
 
-        if use_cache is True and use_merged is False:
-            decoder_with_past_regular_onnx_filenames = ORTModelDecoder._generate_regular_names_for_filename(
-                ONNX_DECODER_WITH_PAST_NAME
-            )
+            decoder_path = decoder_without_past_path
 
-            if decoder_with_past_path.name not in decoder_with_past_regular_onnx_filenames:
+            decoder_regular_onnx_filenames = ORTModelDecoder._generate_regular_names_for_filename(ONNX_DECODER_NAME)
+            if decoder_path.name not in decoder_regular_onnx_filenames:
                 logger.warning(
-                    f"The ONNX file {decoder_with_past_path.name} is not a regular name used in optimum.onnxruntime that are {decoder_with_past_regular_onnx_filenames}, "
-                    f"the {cls.__name__} might not behave as expected."
+                    f"The ONNX file {decoder_path.name} is not a regular name used in optimum.onnxruntime that are {decoder_regular_onnx_filenames}, the "
+                    f"{cls.__name__} might not behave as expected."
+                )
+
+            # If the decoder without / with past has been merged, we do not need to look for any additional file
+            if use_cache is True:
+                if not validate_file_exists(
+                    model_id, decoder_with_past_file_name, subfolder=subfolder, revision=revision
+                ):
+                    try:
+                        decoder_with_past_path = ORTModelDecoder.infer_onnx_filename(
+                            model_id,
+                            [DECODER_WITH_PAST_ONNX_FILE_PATTERN],
+                            "decoder_with_past_file_name",
+                            subfolder=subfolder,
+                            use_auth_token=use_auth_token,
+                            revision=revision,
+                        )
+                    except FileNotFoundError as e:
+                        raise FileNotFoundError(
+                            "The parameter `use_cache=True` was passed to ORTModelForCausalLM.from_pretrained()"
+                            " but no ONNX file using past key values could be found in"
+                            f" {str(Path(model_id, subfolder))}, with the error: {e}"
+                        )
+                else:
+                    decoder_with_past_path = model_path / subfolder / decoder_with_past_file_name
+
+                decoder_with_past_regular_onnx_filenames = ORTModelDecoder._generate_regular_names_for_filename(
+                    ONNX_DECODER_WITH_PAST_NAME
                 )
 
+                if decoder_with_past_path.name not in decoder_with_past_regular_onnx_filenames:
+                    logger.warning(
+                        f"The ONNX file {decoder_with_past_path.name} is not a regular name used in optimum.onnxruntime that are {decoder_with_past_regular_onnx_filenames}, "
+                        f"the {cls.__name__} might not behave as expected."
+                    )
+
         preprocessors = None
         if model_path.is_dir():
             new_model_save_dir = model_path
             preprocessors = maybe_load_preprocessors(model_id)
         else:
             attribute_name_to_filename = {
-                "last_decoder_model_name": decoder_path.name,
+                "last_decoder_model_name": decoder_path.name if use_merged is False else None,
                 "last_decoder_with_past_model_name": decoder_with_past_path.name
-                if decoder_with_past_path is not None
+                if (use_cache is True and use_merged is False)
                 else None,
-                "last_decoder_merged_name": decoder_merged_path.name if decoder_merged_path is not None else None,
+                "last_decoder_merged_name": decoder_merged_path.name if use_merged is True else None,
             }
             paths = {}
             for attr_name, filename in attribute_name_to_filename.items():
@@ -443,18 +446,15 @@ def _from_pretrained(
             new_model_save_dir = Path(model_cache_path).parent
             preprocessors = maybe_load_preprocessors(model_id, subfolder=subfolder)
 
-            last_decoder_with_past_name = paths.get("last_decoder_with_past_model_name", None)
-            if last_decoder_with_past_name is not None:
-                decoder_with_past_path = new_model_save_dir / last_decoder_with_past_name
-
             if use_merged is True:
                 decoder_path = new_model_save_dir / paths["last_decoder_merged_name"]
+                decoder_merged_path = new_model_save_dir / paths["last_decoder_merged_name"]
             else:
                 decoder_path = new_model_save_dir / paths["last_decoder_model_name"]
+                decoder_without_past_path = new_model_save_dir / paths["last_decoder_model_name"]
 
-            decoder_without_past_path = new_model_save_dir / paths["last_decoder_model_name"]
-            if decoder_merged_path is not None:
-                decoder_merged_path = new_model_save_dir / paths["last_decoder_merged_name"]
+                if use_cache is True:
+                    decoder_with_past_path = new_model_save_dir / paths["last_decoder_with_past_model_name"]
 
         ort_inference_sessions = cls.load_model(
             decoder_path=decoder_path,
@@ -481,11 +481,13 @@ def _from_pretrained(
         except OSError:
             logger.info("Generation config file not found, using a generation config created from the model config.")
 
-        onnx_paths = [decoder_without_past_path]
-        if decoder_merged_path is not None:
+        onnx_paths = []
+        if use_merged is False:
+            onnx_paths.append(decoder_without_past_path)
+            if use_cache is True:
+                onnx_paths.append(decoder_with_past_path)
+        else:
             onnx_paths.append(decoder_merged_path)
-        if decoder_with_past_path is not None:
-            onnx_paths.append(decoder_with_past_path)
 
         return cls(
             ort_inference_sessions[0],

diff --git a/optimum/onnxruntime/modeling_ort.py b/optimum/onnxruntime/modeling_ort.py
@@ -783,6 +783,7 @@ def _prepare_io_binding(
                 for axis_name in output_node.shape:
                     output_shape.append(self._output_shape_inference(axis_name, dimensions))
             output_buffer = self._prepare_output_buffer(model, output_shape, output_name)
+
             io_binding.bind_output(
                 output_name,
                 output_buffer.device.type,