Handling ONNX models with external data

optimum/exporters/onnx/convert.py

@@ -444,11 +444,12 @@ def export_models(
 
     for i, model_name in enumerate(models_for_export.keys()):
         submodel, sub_onnx_config = models_for_export[model_name]
-        output_path = (
-            output_dir.joinpath(output_names[i])
-            if output_names is not None
-            else output_dir.joinpath(model_name + ".onnx")
-        )
+        output_name = output_names[i] if output_names is not None else Path(model_name + ".onnx")
+
+        # when the model uses several ONNX files, save each in subfolders to avoid conflicting external files
+        output_path = output_dir / model_name / output_name
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
         outputs.append(
             export(
                 submodel,

optimum/onnxruntime/modeling_decoder.py

@@ -35,7 +35,13 @@
 from ..utils.save_utils import maybe_load_preprocessors, maybe_save_preprocessors
 from .io_binding import TypeHelper
 from .modeling_ort import ORTModel
-from .utils import ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME, get_provider_for_device, parse_device
+from .utils import (
+    ONNX_DECODER_NAME,
+    ONNX_DECODER_WITH_PAST_NAME,
+    _get_external_data_paths,
+    get_provider_for_device,
+    parse_device,
+)
 
 
 if TYPE_CHECKING:
@@ -475,12 +481,16 @@ def _save_pretrained(
         """
         src_paths = [self.decoder_model_path]
         dst_file_names = [decoder_file_name]
+
         if self.use_cache:
             src_paths.append(self.decoder_with_past_model_path)
             dst_file_names.append(decoder_with_past_file_name)
 
+        # add external data paths in case of large models
+        src_paths, dst_file_names = _get_external_data_paths(src_paths, dst_file_names)
+
         for src_path, dst_file_name in zip(src_paths, dst_file_names):
-            dst_path = Path(save_directory).joinpath(dst_file_name)
+            dst_path = Path(save_directory) / dst_file_name
             shutil.copyfile(src_path, dst_path)
 
     @classmethod

optimum/onnxruntime/modeling_ort.py

@@ -52,6 +52,7 @@
 from .io_binding import IOBindingHelper, TypeHelper
 from .utils import (
     ONNX_WEIGHTS_NAME,
+    _get_external_data_paths,
     get_device_for_provider,
     get_provider_for_device,
     parse_device,
@@ -301,9 +302,15 @@ def _save_pretrained(self, save_directory: Union[str, Path], file_name: str = ON
             file_name (`str`, *optional*, defaults to the value of `optimum.onnxruntime.utils.ONNX_WEIGHTS_NAME`):
                 The filename to use when saving the model.
         """
-        # TODO: support models with external data
-        dst_path = Path(save_directory).joinpath(file_name)
-        shutil.copyfile(self.model_path, dst_path)
+        src_paths = [self.model_path]
+        dst_file_names = [file_name]
+
+        # add external data paths in case of large models
+        src_paths, dst_file_names = _get_external_data_paths(src_paths, dst_file_names)
+
+        for src_path, dst_file_name in zip(src_paths, dst_file_names):
+            dst_path = Path(save_directory) / dst_file_name
+            shutil.copyfile(src_path, dst_path)
 
     @staticmethod
     def _generate_regular_names_for_filename(filename: str):

optimum/onnxruntime/modeling_seq2seq.py

@@ -36,7 +36,7 @@
 from ..exporters.onnx import export_models, get_encoder_decoder_models_for_export
 from ..exporters.tasks import TasksManager
 from ..utils import NormalizedConfigManager, check_if_transformers_greater
-from ..utils.file_utils import validate_file_exists
+from ..utils.file_utils import validate_file_exists, find_files_matching_pattern
 from ..utils.save_utils import maybe_load_preprocessors, maybe_save_preprocessors
 from .io_binding import TypeHelper
 from .modeling_decoder import ORTDecoder
@@ -45,6 +45,7 @@
     ONNX_DECODER_NAME,
     ONNX_DECODER_WITH_PAST_NAME,
     ONNX_ENCODER_NAME,
+    _get_external_data_paths,
     get_provider_for_device,
     parse_device,
     validate_provider_availability,
@@ -900,14 +901,18 @@ def _save_pretrained(
                 The decoder with past key values model file name overwriting the default file name, allowing to save
                 the decoder model with a different name.
         """
-        src_file_names = [self.encoder_model_path, self.decoder_model_path]
-        dst_file_names = [encoder_file_name, decoder_file_name]
+        src_paths = [self.encoder_model_path, self.decoder_model_path]
+        dst_file_names = ["encoder_model/" + encoder_file_name, "decoder_model/" + decoder_file_name]
         if self.use_cache:
-            src_file_names.append(self.decoder_with_past_model_path)
-            dst_file_names.append(decoder_with_past_file_name)
+            src_paths.append(self.decoder_with_past_model_path)
+            dst_file_names.append("decoder_with_past_model/" + decoder_with_past_file_name)
 
-        for src_path, dst_file_name in zip(src_file_names, dst_file_names):
+        # add external data paths in case of large models
+        src_paths, dst_file_names = _get_external_data_paths(src_paths, dst_file_names)
+
+        for src_path, dst_file_name in zip(src_paths, dst_file_names):
             dst_path = Path(save_directory) / dst_file_name
+            dst_path.parent.mkdir(parents=True, exist_ok=True)
             shutil.copyfile(src_path, dst_path)
 
     @classmethod
@@ -1152,6 +1157,42 @@ def to(self, device: Union[torch.device, str, int]):
         return self
 
 
+    @staticmethod
+    def infer_onnx_filename(
+        model_name_or_path: Union[str, Path],
+        pattern: str,
+        argument_name: str,
+        subfolder: str = "",
+        use_auth_token: Optional[Union[bool, str]] = None,
+        revision: Optional[str] = None,
+        fail_if_not_found: bool = True,
+    ) -> str:
+        onnx_files = find_files_matching_pattern(
+            model_name_or_path,
+            pattern,
+            glob_pattern="**/*.onnx",
+            subfolder=subfolder,
+            use_auth_token=use_auth_token,
+            revision=revision,
+        )
+
+        path = model_name_or_path
+        if subfolder != "":
+            path = f"{path}/{subfolder}"
+
+        if len(onnx_files) == 0:
+            if fail_if_not_found:
+                raise FileNotFoundError(f"Could not find any ONNX model file in {path}")
+            return None
+        elif len(onnx_files) > 1:
+            if argument_name is not None:
+                raise RuntimeError(
+                    f"Too many ONNX model files were found in {path}, specify which one to load by using the "
+                    f"{argument_name} argument."
+                )
+        return onnx_files[0].parent.name + "/" + onnx_files[0].name
+
+
 class ORTModelForSeq2SeqLM(ORTModelForConditionalGeneration, GenerationMixin):
     """
     Sequence-to-sequence model with a language modeling head for ONNX Runtime inference.

optimum/onnxruntime/utils.py

@@ -16,7 +16,8 @@
 import importlib.util
 import os
 from enum import Enum
-from typing import Dict, Tuple, Union
+from pathlib import Path
+from typing import Dict, List, Tuple, Union
 
 import torch
 from transformers.onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast
@@ -25,6 +26,7 @@
 import onnx
 import onnxruntime as ort
 import pkg_resources
+from onnx.external_data_helper import ExternalDataInfo, _get_initializer_tensors
 
 from ..onnx import OnnxConfigWithLoss, OnnxConfigWithPastAndLoss, OnnxSeq2SeqConfigWithPastAndLoss
 
@@ -270,3 +272,22 @@ class ORTQuantizableOperator(Enum):
     Resize = "Resize"
     AveragePool = "AveragePool"
     Concat = "Concat"
+
+
+def _get_external_data_paths(src_paths: List[Path], dst_file_names: List[str]) -> Tuple[List[Path], List[str]]:
+    """
+    Get external data paths from the model and add them to the list of files to copy.
+    """
+    model_paths = src_paths.copy()
+    for model_path in model_paths:
+        model = onnx.load(str(model_path), load_external_data=False)
+        # filter out tensors that are not external data
+        model_tensors = _get_initializer_tensors(model)
+        model_tensors_ext = [
+            ExternalDataInfo(tensor).location
+            for tensor in model_tensors
+            if tensor.HasField("data_location") and tensor.data_location == onnx.TensorProto.EXTERNAL
+        ]
+        src_paths.extend([model_path.parent / tensor_name for tensor_name in model_tensors_ext])
+        dst_file_names.extend(str(model_path.parent.name / tensor_name) for tensor_name in model_tensors_ext)
+    return src_paths, dst_file_names