[iree-import-onnx] improve handling of large models (#19217)
This pr adds a few options:
1. `--large-model` allows disabling the onnx model checker if a user
knows ahead of time that the model is too large. It will also not load
the external weights in memory unless saving the parameters.
2. `--num-initializers-threshold` allows storing initializers to the
irpa file in batches with a specified number of entries. This can reduce
the memory overhead of first gathering all of the initializers, then
saving them to the irpa at once.
3. `--externalize-inputs-threshold` allows converting inputs to
externalized weights. This is useful for the following workflow:
exporting a HF pytorch model with safetensors, saving a `.irpa` from the
safetensor weights directly, and exporting to onnx with
`export_params=False` and `do_constant_folding=False` (which converts
weights to inputs and avoids folding weights with things like
transposes). When importing to mlir, you can set
`externalize-inputs-threshold=<num_original_inputs>` and it will convert
the inputs from and beyond that threshold to `util.global` ops.
4. `--save-params`/`--no-save-params` factors saving parameters out of
`import_initializer`, and one can avoid saving parameters with
`--no-save-params`. Useful for debugging compilation failures.
## TODO:
Figure out what to do about loading the onnx model and updating opset
version. It's possible to do opset version updating without weights in a
somewhat hacky way, since models > 2GB fail on opset version updating.
Add documentation
---------
Signed-off-by: zjgarvey <zjgarvey@gmail.com>
diff --git a/compiler/bindings/python/iree/compiler/tools/import_onnx/__main__.py b/compiler/bindings/python/iree/compiler/tools/import_onnx/__main__.py
index fc108d8..81ec745 100644
--- a/compiler/bindings/python/iree/compiler/tools/import_onnx/__main__.py
+++ b/compiler/bindings/python/iree/compiler/tools/import_onnx/__main__.py
@@ -20,7 +20,6 @@
from pathlib import Path
import sys
import tempfile
-
from .importer_externalization_overrides import *
@@ -32,25 +31,45 @@
imp: Any = None
if args.externalize_params:
- imp = IREENodeImporter.define_function(
- model_info.main_graph, m, args.num_elements_threshold, args.params_scope
+ if not args.save_params:
+ param_path = None
+ elif args.save_params_to:
+ param_path = args.save_params_to
+ elif (args.output_file is not None) and (args.output_file != "-"):
+ output_dir = Path(args.output_file).parent
+ output_stem = Path(args.output_file).stem
+ param_path = output_dir / (output_stem + "_params.irpa")
+ else:
+ raise ValueError(
+ "If `--externalize-params` is set and `--output-file` is stdout, either `--save-params-to` or `--no-save-params` must be set."
+ )
+ data_dir = (
+ args.data_dir
+ if args.data_dir is not None
+ else str(Path(args.input_file).parent)
)
+ param_bit_threshold = (
+ None
+ if args.param_gb_threshold is None
+ else int(args.param_gb_threshold * 8 * (10**9))
+ )
+ param_data = ParamData(
+ param_bit_threshold=param_bit_threshold,
+ num_elements_threshold=args.num_elements_threshold,
+ params_scope=args.params_scope,
+ data_dir=data_dir,
+ param_path=str(param_path),
+ input_index_threshold=args.externalize_inputs_threshold,
+ )
+ imp = IREENodeImporter.define_function(model_info.main_graph, m, param_data)
else:
imp = onnx_importer.NodeImporter.define_function(model_info.main_graph, m)
+
imp.import_all()
if not args.no_verify:
m.verify()
- if args.externalize_params:
- default_param_path = Path(args.output_file).parent / Path(args.output_file).stem
- param_path = (
- (str(default_param_path) + "_params.irpa")
- if args.save_params_to is None
- else str(args.save_params_to)
- )
- imp.param_archive.create_archive_file(param_path)
-
# TODO: This isn't very efficient output. If these files ever
# get large, enable bytecode and direct binary emission to save
# some copies.
@@ -60,41 +79,54 @@
else:
print(m.get_asm(assume_verified=not args.no_verify))
+ if args.externalize_params and args.save_params:
+ imp.save_params()
+
def load_onnx_model(args: argparse.Namespace) -> onnx.ModelProto:
input_dir = os.path.dirname(os.path.abspath(args.input_file))
- # Load the model, with possible external data coming from the default
- # location, or the location specified on the command line.
- if args.data_dir is None:
- raw_model = onnx.load(args.input_file)
- else:
- raw_model = onnx.load(args.input_file, load_external_data=False)
- onnx.load_external_data_for_model(raw_model, str(args.data_dir))
-
- # Only change the opset version if it is greater than the current one.
- if args.opset_version and args.opset_version > raw_model.opset_import[0].version:
- raw_model = onnx.version_converter.convert_version(
- raw_model, args.opset_version
+ # TODO: setup updating opset version without loading external weights.
+ if args.opset_version and args.large_model:
+ raise NotImplementedError(
+ "Updating the opset version for large models is currently unsupported."
)
- # Do shape inference two ways. First, attempt in-memory to avoid redundant
- # loading and the need for writing a temporary file somewhere. If that
- # fails, typically because of the 2 GB protobuf size limit, try again via
- # files. See
- # https://onnx.ai/onnx/repo-docs/PythonAPIOverview.html#shape-inference-a-large-onnx-model-2gb
- # for details about the file-based technique.
+ if not args.large_model:
+ # Load the model, with possible external data coming from the default
+ # location, or the location specified on the command line.
+ if args.data_dir is None:
+ raw_model = onnx.load(args.input_file)
+ else:
+ raw_model = onnx.load(args.input_file, load_external_data=False)
+ onnx.load_external_data_for_model(raw_model, str(args.data_dir))
- # Run the checker to test whether the file is above the threshold for
- # in-memory shape inference. If not, go ahead and do the shape inference.
- try:
- onnx.checker.check_model(raw_model)
- inferred_model = onnx.shape_inference.infer_shapes(
- raw_model, data_prop=args.data_prop
- )
- return inferred_model
- except ValueError:
- pass
+ # Only change the opset version if it is greater than the current one.
+ if (
+ args.opset_version
+ and args.opset_version > raw_model.opset_import[0].version
+ ):
+ raw_model = onnx.version_converter.convert_version(
+ raw_model, args.opset_version
+ )
+
+ # Do shape inference two ways. First, attempt in-memory to avoid redundant
+ # loading and the need for writing a temporary file somewhere. If that
+ # fails, typically because of the 2 GB protobuf size limit, try again via
+ # files. See
+ # https://onnx.ai/onnx/repo-docs/PythonAPIOverview.html#shape-inference-a-large-onnx-model-2gb
+ # for details about the file-based technique.
+
+ # Run the checker to test whether the file is above the threshold for
+ # in-memory shape inference. If not, go ahead and do the shape inference.
+ try:
+ onnx.checker.check_model(raw_model)
+ inferred_model = onnx.shape_inference.infer_shapes(
+ raw_model, data_prop=args.data_prop
+ )
+ return inferred_model
+ except ValueError:
+ pass
# Model is too big for in-memory inference: do file-based shape inference
# to a temp file.
@@ -111,7 +143,9 @@
# Load the temp file and the external data.
inferred_model = onnx.load(temp_inferred_file, load_external_data=False)
data_dir = Path(input_dir if args.data_dir is None else args.data_dir)
- onnx.load_external_data_for_model(inferred_model, str(data_dir))
+ # we don't need to load the model weights in-memory when externalizing params
+ if not args.externalize_params:
+ onnx.load_external_data_for_model(inferred_model, str(data_dir))
return inferred_model
@@ -146,30 +180,63 @@
type=int,
)
parser.add_argument(
- "--num-elements-threshold",
- help="Minimum number of elements for an initializer to be externalized. Only has an effect if 'externalize-params' is true.",
- type=int,
- default=100,
- )
- parser.add_argument(
- "--externalize-params",
- help="Externalize large parameters and store them on the disk, to load at runtime.",
+ "--large-model",
+ help="Setting this to true is recommended for large models that do not require --opset-version."
+ " It will bypass loading external weights and running the onnx checker to determine the model size.",
action=argparse.BooleanOptionalAction,
default=False,
)
- parser.add_argument(
+ # args for saving a file with externalized params
+ externalization_args = parser.add_argument_group(
+ "externalization", "args used to customize the externalization of model weights"
+ )
+ externalization_args.add_argument(
+ "--externalize-params",
+ help="Import the mlir file with large weights replaced by external reference calls.",
+ action=argparse.BooleanOptionalAction,
+ default=False,
+ )
+ externalization_args.add_argument(
+ "--externalize-inputs-threshold",
+ help="Treats inputs at or after the provided index as external parameters of the model."
+ " Only has an effect if 'externalize-params' is true.",
+ type=int,
+ )
+ externalization_args.add_argument(
+ "--num-elements-threshold",
+ help="Minimum number of elements for an initializer to be externalized."
+ " Only has an effect if 'externalize-params' is true.",
+ type=int,
+ default=100,
+ )
+ externalization_args.add_argument(
+ "--params-scope",
+ help="The namespace or the scope in which the externalized parameters are placed."
+ " Default is 'model'.",
+ type=str,
+ default="model",
+ )
+ # args for creating an external weight file
+ externalization_args.add_argument(
+ "--save-params",
+ help="Whether to save the params to a file. Setting this to false will generate mlir with externalized weights"
+ " without creating an associated .irpa file.",
+ action=argparse.BooleanOptionalAction,
+ default=True,
+ )
+ externalization_args.add_argument(
+ "--param-gb-threshold",
+ help="Setting this will flush params to a temp file when total in-memory param size exceeds the Gigabyte threshold."
+ " This is less efficient (about x2 slower) and only recommended for machines with limited RAM.",
+ type=float,
+ )
+ externalization_args.add_argument(
"--save-params-to",
help="Location to save the externalized parameters. When not set, the parameters will be written to '<output_file_name>_params.irpa'"
" under the namespace 'model', which can be configured by passing the namespace string to 'params-scope'.",
default=None,
type=Path,
)
- parser.add_argument(
- "--params-scope",
- help="The namespace or the scope in which the externalized parameters are placed. Default is 'model'.",
- type=str,
- default="model",
- )
args = parser.parse_args(argv)
return args
diff --git a/compiler/bindings/python/iree/compiler/tools/import_onnx/importer_externalization_overrides.py b/compiler/bindings/python/iree/compiler/tools/import_onnx/importer_externalization_overrides.py
index d99b269..8146333 100644
--- a/compiler/bindings/python/iree/compiler/tools/import_onnx/importer_externalization_overrides.py
+++ b/compiler/bindings/python/iree/compiler/tools/import_onnx/importer_externalization_overrides.py
@@ -5,15 +5,20 @@
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
import copy
+import logging
import random
import string
-import iree.runtime as rt
+import sys
+import time
+import tempfile
+from pathlib import Path
+from typing import Optional, Tuple, Any, NamedTuple, Union
-from ...dialects import util
-from typing import Optional, Tuple, Any
+import numpy
try:
import onnx
+ from onnx import numpy_helper
except ModuleNotFoundError as e:
raise ModuleNotFoundError(
f"iree-import-onnx requires that the `onnx` Python package is installed "
@@ -27,7 +32,7 @@
"iree-import-onnx is only available if IREE was built with Torch support"
) from e
-from onnx import numpy_helper
+from ...dialects import util
from ...ir import (
Context,
@@ -44,6 +49,17 @@
IntegerType,
)
+logger = logging.getLogger(__name__)
+
+
+class ParamData(NamedTuple):
+ param_bit_threshold: Optional[int]
+ num_elements_threshold: int
+ params_scope: str
+ data_dir: str
+ param_path: str
+ input_index_threshold: Optional[int]
+
class IREENodeImporter(onnx_importer.NodeImporter):
def __init__(
@@ -55,8 +71,7 @@
context_cache: "onnx_importer.ContextCache",
module_op: Operation,
module_cache: "onnx_importer.ModuleCache",
- num_elements_threshold: int,
- params_scope: str,
+ param_data: ParamData,
):
super().__init__(
graph_info,
@@ -69,9 +84,8 @@
self.last_global_op = None
self.symbol_table = SymbolTable(module_op)
self.symbol_table.insert(parent_op)
- self.num_elements_threshold = num_elements_threshold
- self.param_archive = rt.ParameterIndex()
- self.params_scope = params_scope
+ self.param_data = param_data
+ self.globals = []
def sanitize_name(self, name: str) -> str:
# There are often some initializers in the models that have no name
@@ -92,9 +106,17 @@
new_name = str(random.randrange(1, 1000)) + "__" + ch
return new_name
+ def get_type_info_from_type(self, tp: onnx.TypeProto):
+ tt = tp.tensor_type
+ if tt.elem_type:
+ dims = tuple(
+ (d.dim_value if not d.dim_param else None) for d in tt.shape.dim
+ )
+ return dims, tt.elem_type
+
def create_tensor_global(
self,
- t: onnx.TensorProto,
+ t: Union[onnx.TensorProto, onnx.ValueInfoProto],
) -> Tuple[str, IrType]:
# Always create globals at the top. Then after created, if there was
# a prior one, move the new one to after it to maintain declaration
@@ -106,8 +128,18 @@
# After lowering to linalg-on-tensors, the data type needs to be signless.
# So, we construct the globals to have signless types, and use
# torch_c.from_builtin_tensor to convert to the correct frontend type.
+ if isinstance(t, onnx.TensorProto):
+ dims = tuple(t.dims)
+ data_type = t.data_type
+ elif isinstance(t, onnx.ValueInfoProto):
+ dims, data_type = self.get_type_info_from_type(t.type)
+ else:
+ raise TypeError(
+ f"Expected an onnx.TensorProto or an onnx.ValueInfoProto, recieved {type(t)} from {name}"
+ )
+
vtensor_type = RankedTensorType.get(
- tuple(t.dims), ELEM_TYPE_TO_SIGNLESS_IR_TYPE[t.data_type]()
+ dims, ELEM_TYPE_TO_SIGNLESS_IR_TYPE[data_type]()
)
ir_attrs = {
"sym_name": StringAttr.get(name),
@@ -115,7 +147,7 @@
"type": TypeAttr.get(vtensor_type),
}
- external_scope_attr = StringAttr.get(self.params_scope)
+ external_scope_attr = StringAttr.get(self.param_data.params_scope)
external_name_attr = StringAttr.get(name)
ir_attrs["initial_value"] = Attribute.parse(
f"#stream.parameter.named<{external_scope_attr}::{external_name_attr}> : {vtensor_type}"
@@ -138,8 +170,7 @@
cls,
graph_info: onnx_importer.GraphInfo,
module_op: Operation,
- num_elements_threshold: int,
- params_scope: str,
+ param_data: ParamData,
context_cache: Optional["onnx_importer.ContextCache"] = None,
module_cache: Optional["onnx_importer.ModuleCache"] = None,
private: bool = False,
@@ -155,6 +186,16 @@
# Recover per-module caches of various attributes.
# Allows modification in the same module_op without
# loss of current state.
+ all_input_map = copy.deepcopy(graph_info.input_map)
+ num_inputs = len(all_input_map.items())
+ if param_data.input_index_threshold is not None:
+ if param_data.input_index_threshold not in range(0, num_inputs):
+ raise ValueError(
+ f"input_index_threshold must be in the range [0,num_inputs={num_inputs})"
+ )
+ for _index in range(param_data.input_index_threshold, num_inputs):
+ _discarded = graph_info.input_map.popitem()
+
mc = (
module_cache
if module_cache is not None
@@ -187,32 +228,49 @@
context_cache=cc,
module_op=module_op,
module_cache=mc,
- num_elements_threshold=num_elements_threshold,
- params_scope=params_scope,
+ param_data=param_data,
)
for node_name, input_value in zip(graph_info.input_map.keys(), block.arguments):
imp._nv_map[node_name] = input_value
imp._populate_graph_attrs(func_op)
+ imp._gi.input_map = all_input_map
return imp
+ def import_all(self, func=True):
+ num_inputs = len(self._gi.input_map.items())
+ if self.param_data.input_index_threshold is not None:
+ for i in range(self.param_data.input_index_threshold, num_inputs):
+ self.import_initializer(list(self._gi.input_map.values())[i])
+ super().import_all(func)
+
def import_initializer(
- self, initializer: onnx.TensorProto, extern_name: Optional[str] = None
+ self,
+ initializer: Union[onnx.TensorProto, onnx.ValueInfoProto],
+ extern_name: Optional[str] = None,
) -> Value:
# If an explicitly specified name is given, use that; otherwise, pick
# up the name from the tensor proto itself
initializer_name = extern_name if extern_name else initializer.name
- dims = list(initializer.dims)
- num_elements = 1
- for d in dims:
- num_elements = num_elements * d
- if num_elements < self.num_elements_threshold:
- imported_tensor = super().import_initializer(initializer)
- self._nv_map[initializer_name] = imported_tensor
- return imported_tensor
+ if isinstance(initializer, onnx.TensorProto):
+ dims = tuple(initializer.dims)
+ num_elements = 1
+ for d in dims:
+ num_elements = num_elements * d
+ if num_elements < self.param_data.num_elements_threshold:
+ imported_tensor = super().import_initializer(initializer)
+ self._nv_map[initializer_name] = imported_tensor
+ return imported_tensor
+ data_type = initializer.data_type
+ elif isinstance(initializer, onnx.ValueInfoProto):
+ dims, data_type = self.get_type_info_from_type(initializer.type)
+ else:
+ raise TypeError(
+ f"Expected an onnx.TensorProto or an onnx.ValueInfoProto, recieved {type(initializer)} from {initializer_name}"
+ )
actual_symbol_name, tensor_type = self.create_tensor_global(initializer)
vtensor_type = self._cc.get_vtensor_type(
- tuple(initializer.dims), self._cc.tensor_element_type(initializer.data_type)
+ dims, self._cc.tensor_element_type(data_type)
)
with InsertionPoint(self._b), Location.name(initializer_name):
@@ -224,10 +282,101 @@
).result
self._nv_map[initializer_name] = converted_value
- tensor_as_array = numpy_helper.to_array(initializer)
- self.param_archive.add_buffer(actual_symbol_name, tensor_as_array)
+ if isinstance(initializer, onnx.TensorProto):
+ self.globals.append((initializer_name, actual_symbol_name))
return converted_value
+ def save_params(self):
+ """
+ Only gets called if the arg `--save-params` is set to `True`.
+ Saving params requires iree-runtime, so putting the import here will avoid requiring it uniformly for the importer.
+ """
+ try:
+ import iree.runtime as rt
+ except ModuleNotFoundError as e:
+ raise ModuleNotFoundError(
+ "iree-import-onnx requires iree runtime api for externalizing parameters. "
+ "For example: `pip install iree-base-runtime`"
+ ) from e
+
+ param_archive = rt.ParameterIndex()
+ # if we don't need to save params in batches, gather all tensors in param_archive
+ if self.param_data.param_bit_threshold is None:
+ for name, actual_symbol_name in self.globals:
+ initializer = self._gi.initializer_map[name]
+ tensor_as_array = numpy_helper.to_array(
+ initializer, base_dir=self.param_data.data_dir
+ )
+ param_archive.add_buffer(actual_symbol_name, tensor_as_array)
+ param_archive.create_archive_file(self.param_data.param_path)
+ return
+
+ # else we need to save in batches:
+ # 1. setup a temporary directory to save smaller param files
+ # 2. keep a target_index for storing references to saved tensors
+ # 3. create an archive file for target_index at param_path to gather all temp data
+ target_index = rt.ParameterIndex()
+ in_memory_param_bits = 0
+ iter = 0
+ t00 = time.time()
+ with tempfile.TemporaryDirectory(
+ dir=Path(self.param_data.param_path).parent
+ ) as temp_dir_name:
+ get_curr_path = lambda: str(Path(temp_dir_name) / f"params_{iter}.irpa")
+ for name, actual_symbol_name in self.globals:
+ initializer = self._gi.initializer_map[name]
+ tensor_as_array = numpy_helper.to_array(
+ initializer, base_dir=self.param_data.data_dir
+ )
+ param_archive.add_buffer(actual_symbol_name, tensor_as_array)
+ # get the new param size
+ elem_dtype = tensor_as_array.dtype
+ elem_kind = elem_dtype.kind
+ if elem_kind not in ["i", "f"]:
+ raise TypeError(f"Unhandled numpy dtype: {elem_dtype}")
+ elem_info = (
+ numpy.iinfo(elem_dtype)
+ if elem_kind == "i"
+ else numpy.finfo(elem_dtype)
+ )
+ elem_bits = elem_info.bits
+ param_bits = tensor_as_array.size * elem_bits
+ # update the running total memory use
+ in_memory_param_bits += param_bits
+ if param_bits >= self.param_data.param_bit_threshold:
+ logger.warning(
+ f"Single parameter {name} is {param_bits} bits, "
+ + f"which exceeds threshold of {self.param_data.param_bit_threshold} bits."
+ )
+ # flush the param archive to a temp file if the threshold is exceeded
+ if in_memory_param_bits >= self.param_data.param_bit_threshold:
+ t0 = time.time()
+ param_archive.create_archive_file(
+ get_curr_path(),
+ target_index=target_index,
+ )
+ logger.info(
+ f"iter {iter} with {in_memory_param_bits} bits took {time.time() - t0}s to flush"
+ )
+ iter += 1
+ del param_archive
+ param_archive = rt.ParameterIndex()
+ in_memory_param_bits = 0
+
+ # write any remaining params to a temp file
+ t0 = time.time()
+ param_archive.create_archive_file(
+ get_curr_path(), target_index=target_index
+ )
+ logger.info(
+ f"iter {iter} with {in_memory_param_bits} bits took {time.time() - t0}s to flush"
+ )
+ # combine all temporary param files into the final result
+ t0 = time.time()
+ target_index.create_archive_file(self.param_data.param_path)
+ logger.info(f"combining {iter + 1} irpa files took {time.time() - t0}s")
+ logger.info(f"total time to save params: {time.time()-t00}")
+
ELEM_TYPE_TO_SIGNLESS_IR_TYPE = copy.deepcopy(onnx_importer.ELEM_TYPE_TO_IR_TYPE_CB)
diff --git a/compiler/bindings/python/test/tools/import_onnx_test.py b/compiler/bindings/python/test/tools/import_onnx_test.py
index 6089b65..0692e4c 100644
--- a/compiler/bindings/python/test/tools/import_onnx_test.py
+++ b/compiler/bindings/python/test/tools/import_onnx_test.py
@@ -8,6 +8,7 @@
import sys
import tempfile
import unittest
+from pathlib import Path
def run_tool(*argv: str):
@@ -125,6 +126,70 @@
# there should be no inlined constants.
self.assertNotIn("onnx.Constant", contents)
+ def testExternalizeInitializersThreshold(self):
+ run_tool(
+ LARGE_WEIGHTS_ONNX_FILE_PATH,
+ "--externalize-params",
+ "--param-gb-threshold",
+ "0.0",
+ "-o",
+ self.outputPath,
+ )
+ with open(self.outputPath, "rt") as f:
+ contents = f.read()
+ self.assertIn("util.global", contents)
+ self.assertIn("util.global.load", contents)
+ self.assertIn("onnx.Constant", contents)
+
+ def testExternalizeInputsThreshold(self):
+ run_tool(
+ LARGE_WEIGHTS_ONNX_FILE_PATH,
+ "--externalize-params",
+ "--externalize-inputs-threshold",
+ "0",
+ "-o",
+ self.outputPath,
+ )
+ with open(self.outputPath, "rt") as f:
+ contents = f.read()
+ self.assertIn("util.global", contents)
+ self.assertIn("util.global.load", contents)
+ self.assertNotIn("%arg0", contents)
+ self.assertIn("onnx.Constant", contents)
+
+ def testNoSaveParams(self):
+ param_path = str(Path(self.outputPath).parent / "test.irpa")
+ run_tool(
+ LARGE_WEIGHTS_ONNX_FILE_PATH,
+ "--externalize-params",
+ "--no-save-params",
+ "--save-params-to",
+ param_path,
+ "-o",
+ self.outputPath,
+ )
+ if os.path.exists(param_path):
+ self.fail("expected param_path to not exist with --no-save-params")
+ with open(self.outputPath, "rt") as f:
+ contents = f.read()
+ self.assertIn("util.global", contents)
+ self.assertIn("util.global.load", contents)
+ self.assertIn("onnx.Constant", contents)
+
+ def testLargeModelFlag(self):
+ run_tool(
+ LARGE_WEIGHTS_ONNX_FILE_PATH,
+ "--externalize-params",
+ "--large-model",
+ "-o",
+ self.outputPath,
+ )
+ with open(self.outputPath, "rt") as f:
+ contents = f.read()
+ self.assertIn("util.global", contents)
+ self.assertIn("util.global.load", contents)
+ self.assertIn("onnx.Constant", contents)
+
if __name__ == "__main__":
try:
