Saving external data for large ONNX models

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

@@ -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,13 +901,16 @@ 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]
+        src_paths = [self.encoder_model_path, self.decoder_model_path]
         dst_file_names = [encoder_file_name, decoder_file_name]
         if self.use_cache:
-            src_file_names.append(self.decoder_with_past_model_path)
+            src_paths.append(self.decoder_with_past_model_path)
             dst_file_names.append(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
             shutil.copyfile(src_path, dst_path)
 

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(model_tensors_ext)
+    return src_paths, dst_file_names