diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5a20007eef7..7d1be7313b3 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -66,6 +66,7 @@ Changelog **Bug Fixes** +- Fix ONNX AutoCast failing on models with external initializers larger than 2 GiB. - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled. - Fix NVFP4 ONNX export to quantize FP4 weights with the published FP8 block scales, matching eager ModelOpt packed weights. Block scales below ``2**-9`` are now clamped to that minimum, and non-finite or negative scales raise an error. - Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint. diff --git a/modelopt/onnx/autocast/convert.py b/modelopt/onnx/autocast/convert.py index 65cabe86974..aef5f03cce8 100644 --- a/modelopt/onnx/autocast/convert.py +++ b/modelopt/onnx/autocast/convert.py @@ -23,6 +23,7 @@ nodes. """ +import os from copy import deepcopy import numpy as np @@ -112,7 +113,7 @@ def convert_to_mixed_precision( onnx.ModelProto: The converted mixed precision model. """ # Load and process model - model = onnx.load(onnx_path, load_external_data=True) + model = onnx.load(onnx_path, load_external_data=False) assert low_precision_type in ["fp16", "bf16"], "low_precision_type must be either fp16 or bf16" original_network_io_metadata = _capture_network_io_metadata(model, keep_io_types) @@ -150,6 +151,7 @@ def convert_to_mixed_precision( trt_plugins=trt_plugins, trt_plugins_precision=trt_plugins_precision, max_ir_version=LATEST_IR_VERSION_SUPPORTED_BY_ORT, + onnx_path=onnx_path, ) graph_sanitizer.sanitize() model = graph_sanitizer.model @@ -158,6 +160,9 @@ def convert_to_mixed_precision( # as an exception (triggering infer_types' standalone type-inference fallback) instead of # silently leaving tensors untyped, which would break later type lookups. model = onnx_utils.infer_types(model, use_standalone_type_inference, strict_mode=True) + onnx.external_data_helper.load_external_data_for_model( + model, os.path.dirname(os.path.abspath(onnx_path)) + ) value_info_map, initializer_map, node_to_init_map = utils.setup_mappings(model) # Automatically add 'trt' to list of providers if custom ops are detected @@ -191,6 +196,7 @@ def convert_to_mixed_precision( custom_ops=graph_sanitizer.custom_ops, use_standalone_type_inference=use_standalone_type_inference, original_network_io_metadata=original_network_io_metadata, + sanitize_model=False, ) # Obtain reference data diff --git a/modelopt/onnx/autocast/graphsanitizer.py b/modelopt/onnx/autocast/graphsanitizer.py index 2154a42568e..e9be72c0c27 100644 --- a/modelopt/onnx/autocast/graphsanitizer.py +++ b/modelopt/onnx/autocast/graphsanitizer.py @@ -15,6 +15,8 @@ """Graph sanitization and optimization for ONNX models.""" +import os + import numpy as np import onnx import onnx_graphsurgeon as gs @@ -25,7 +27,12 @@ import modelopt.onnx.utils as onnx_utils from modelopt.onnx.autocast.logging_config import logger from modelopt.onnx.quantization.graph_utils import cast_custom_ops -from modelopt.onnx.trt_utils import interpret_trt_plugins_precision_flag +from modelopt.onnx.trt_utils import ( + get_custom_layers, + infer_types_shapes_tensorrt, + interpret_trt_plugins_precision_flag, + set_trt_plugin_domain, +) class GraphSanitizer: @@ -38,6 +45,7 @@ def __init__( max_ir_version: int | None = None, trt_plugins: list[str] | None = [], trt_plugins_precision: list[str] | None = [], + onnx_path: str | None = None, ) -> None: """Initialize GraphSanitizer. @@ -46,6 +54,7 @@ def __init__( min_opset: minimum opset version to use max_ir_version: maximum IR version supported by ORT trt_plugins: list of TensorRT plugin library paths in .so format (compiled shared library). + onnx_path: path to the source ONNX model, used to resolve external data. """ self.model = model self.min_opset = min_opset @@ -55,6 +64,8 @@ def __init__( self.custom_ops_low_precision_nodes = [] self.trt_plugins = trt_plugins self.trt_plugins_precision = trt_plugins_precision or [] + self.onnx_path = os.path.abspath(onnx_path) if onnx_path is not None else None + self.external_data_dir = os.path.dirname(self.onnx_path) if self.onnx_path else "" def sanitize(self) -> None: """Sanitize the model graph. @@ -118,13 +129,14 @@ def find_custom_nodes(self) -> None: node.op_type for node in self.model.graph.node if node.op_type not in self.standard_ops } if self.custom_ops: - from modelopt.onnx.trt_utils import infer_types_shapes_tensorrt, set_trt_plugin_domain - # Set TensorRT plugin domain info in the graph for ORT compatibility self.model = set_trt_plugin_domain(self.model, self.custom_ops) # Infer types and shapes in the graph for ORT compatibility - self.model = infer_types_shapes_tensorrt(self.model, self.trt_plugins) + _, all_tensor_info = get_custom_layers(self.onnx_path or self.model, self.trt_plugins) + self.model = infer_types_shapes_tensorrt( + self.model, self.trt_plugins, all_tensor_info=all_tensor_info + ) def remove_disconnected_outputs(self) -> None: """Remove disconnected outputs from the model.""" @@ -501,7 +513,7 @@ def _get_initializer_value(self, name: str, return_array: bool = False) -> np.nd """Get value from an initializer by name.""" for init in self.model.graph.initializer: if init.name == name: - value = numpy_helper.to_array(init) + value = numpy_helper.to_array(init, base_dir=self.external_data_dir) return value if return_array else value.item() return None @@ -516,7 +528,7 @@ def _convert_fp64_initializers(self) -> bool: for initializer in self.model.graph.initializer: if initializer.data_type == onnx.TensorProto.DOUBLE: # Convert the data to FP32 - fp64_data = numpy_helper.to_array(initializer) + fp64_data = numpy_helper.to_array(initializer, base_dir=self.external_data_dir) fp32_data = fp64_data.astype(np.float32) # Create new initializer with FP32 data @@ -575,7 +587,7 @@ def _convert_fp64_nodes(self) -> bool: for attr in node.attribute: if attr.name == "value" and attr.t.data_type == onnx.TensorProto.DOUBLE: # Convert the tensor value to FP32 - fp64_data = numpy_helper.to_array(attr.t) + fp64_data = numpy_helper.to_array(attr.t, base_dir=self.external_data_dir) fp32_data = fp64_data.astype(np.float32) new_tensor = numpy_helper.from_array(fp32_data) attr.t.CopyFrom(new_tensor) diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py index bd5c0fc13f3..7db660b29b4 100644 --- a/modelopt/onnx/autocast/precisionconverter.py +++ b/modelopt/onnx/autocast/precisionconverter.py @@ -99,6 +99,7 @@ def __init__( tensor_block_dict: dict[str, dict[str, list[int]]] = {}, use_standalone_type_inference: bool = False, original_network_io_metadata: dict[str, list[onnx.ValueInfoProto]] | None = None, + sanitize_model: bool = True, ) -> None: """Initialize PrecisionConverter. @@ -118,11 +119,18 @@ def __init__( tensor_block_dict: Dictionary of tensors (operation type and I/O indices) that should remain in FP32. use_standalone_type_inference: Use standalone type inference instead of ONNX's infer_shapes. original_network_io_metadata: Original public input/output metadata captured at the API boundary. + sanitize_model: Whether to sanitize the model before precision conversion. """ self.model = deepcopy(model) - self.value_info_map = value_info_map - self.initializer_map = initializer_map - self.node_to_init_map = node_to_init_map + self.sanitize_model = sanitize_model + if sanitize_model: + self.value_info_map = value_info_map + self.initializer_map = initializer_map + self.node_to_init_map = node_to_init_map + else: + self.value_info_map, self.initializer_map, self.node_to_init_map = utils.setup_mappings( + self.model + ) self.keep_io_types = keep_io_types self.init_conversion_max_bytes = ( np.inf if init_conversion_max_bytes is None else init_conversion_max_bytes @@ -195,7 +203,8 @@ def convert( "AutoCast can only operate on valid ONNX models, but the input model is invalid. See log for details." ) - self._sanitize_model() + if self.sanitize_model: + self._sanitize_model() # Filter out nodes that are not allowed to be in low precision # This is done here and not in NodeClassifier because it is required for the model to be valid diff --git a/modelopt/onnx/autocast/referencerunner.py b/modelopt/onnx/autocast/referencerunner.py index 6dcf7137689..4db54b3b366 100644 --- a/modelopt/onnx/autocast/referencerunner.py +++ b/modelopt/onnx/autocast/referencerunner.py @@ -26,6 +26,7 @@ import copy import io +import os import sys import tempfile from collections import OrderedDict @@ -112,8 +113,6 @@ def _load_inputs_from_npz(self, input_data_path): Returns: List of input dictionaries, one per batch. """ - import os - if os.path.isdir(input_data_path): # Load all NPZ files in the directory as multiple batches npz_files = sorted([f for f in os.listdir(input_data_path) if f.endswith(".npz")]) @@ -135,7 +134,7 @@ def _load_inputs_from_npz(self, input_data_path): def _validate_inputs(self, data_loader): """Validate that input names and shapes match the model.""" if isinstance(data_loader, list) and ( - isinstance(data_loader[0], (dict, np.lib.npyio.NpzFile)) + isinstance(data_loader[0], dict | np.lib.npyio.NpzFile) ): if sorted(self.input_names) != sorted(data_loader[0].keys()): raise ValueError("Input names from ONNX model do not match provided input names.") @@ -165,8 +164,6 @@ def _load_inputs(self, inputs): # If no inputs are provided, use random inputs data_loader = DataLoader(val_range={"": (-1, 1)}) - import os - if inputs is not None: if isinstance(inputs, str): if inputs.endswith(".json"): @@ -178,7 +175,7 @@ def _load_inputs(self, inputs): f"Invalid input file: {inputs}. Supported input types: .json (Polygraphy JSON format), " ".npz (Numpy), or a directory containing .npz files" ) - elif isinstance(inputs, (dict, OrderedDict)): + elif isinstance(inputs, dict | OrderedDict): data_loader = [inputs] else: raise ValueError( @@ -193,32 +190,32 @@ def _get_ort_runner(self, model): from polygraphy.backend.onnx import BytesFromOnnx from polygraphy.backend.onnxrt import OnnxrtRunner, SessionFromOnnx - # Check if model has external data by checking: - # 1. If any initializer has data_location set to EXTERNAL (even if data is loaded) - # 2. If model size would exceed 2GB (indicating need for external data) - needs_external_data = onnx_utils.check_model_uses_external_data( - self.model - ) or self.model.ByteSize() > 2 * (1024**3) - if needs_external_data: - logger.debug("Model has external data, using file-based approach") - # Get the actual ONNX ModelProto from ModifyOutputs wrapper - modified_model = model() - - # Use a persistent temp file, because we need the file to be present in an broader context - tmp_file = tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) - tmp_file.close() - tmp_file_path = tmp_file.name - onnx_utils.save_onnx(modified_model, tmp_file_path, save_as_external_data=True) - logger.debug(f"Model with all outputs saved to {tmp_file_path}") - build_onnxrt_session = SessionFromOnnx(tmp_file_path, providers=self.providers) + # Get the actual ONNX ModelProto from ModifyOutputs wrapper + modified_model = model() - else: - # For models without external data, use the original BytesFromOnnx approach (no tmp files) - logger.debug("Model has no external data, using BytesFromOnnx approach") - serialize_onnx = BytesFromOnnx(model) - build_onnxrt_session = SessionFromOnnx(serialize_onnx, providers=self.providers) - runners = [OnnxrtRunner(build_onnxrt_session)] - return runners + needs_file_backed_model = onnx_utils.check_model_uses_external_data( + modified_model + ) or onnx_utils.is_model_too_large_for_protobuf(modified_model) + model_temp_dir = None + try: + if needs_file_backed_model: + logger.debug("Model has external data, using file-based approach") + model_temp_dir = tempfile.TemporaryDirectory() + tmp_file_path = os.path.join(model_temp_dir.name, "model.onnx") + onnx_utils.save_onnx(modified_model, tmp_file_path, save_as_external_data=True) + logger.debug(f"Model with all outputs saved to {tmp_file_path}") + build_onnxrt_session = SessionFromOnnx(tmp_file_path, providers=self.providers) + else: + # For models without external data, use the original BytesFromOnnx approach (no tmp files) + logger.debug("Model has no external data, using BytesFromOnnx approach") + serialize_onnx = BytesFromOnnx(modified_model) + build_onnxrt_session = SessionFromOnnx(serialize_onnx, providers=self.providers) + runners = [OnnxrtRunner(build_onnxrt_session)] + except Exception: + if model_temp_dir is not None: + model_temp_dir.cleanup() + raise + return runners, model_temp_dir def _aggregate_tensor_stats(self, all_batch_data: list[OrderedDict]) -> OrderedDict: """Aggregate tensor statistics across multiple batches. @@ -300,22 +297,25 @@ def run(self, inputs=None): modify_outputs = ModifyOnnxOutputs(model_copy, outputs=constants.MARK_ALL) # Load the modified model and create an inference session - runners = self._get_ort_runner(modify_outputs) - - # Comparator is used despite the fact that we are using ONNXRuntime - # because it provides the ability to generate random inputs using DataLoader - data_loader = self._load_inputs(inputs) - - # Temporarily redirect stdout to suppress Comparator.run() output - stdout = sys.stdout - string_buffer = io.StringIO() - sys.stdout = string_buffer + runners, model_temp_dir = self._get_ort_runner(modify_outputs) try: - results = Comparator.run(runners, data_loader=data_loader) + # Comparator is used despite the fact that we are using ONNXRuntime + # because it provides the ability to generate random inputs using DataLoader + data_loader = self._load_inputs(inputs) + + # Temporarily redirect stdout to suppress Comparator.run() output + stdout = sys.stdout + string_buffer = io.StringIO() + sys.stdout = string_buffer + try: + results = Comparator.run(runners, data_loader=data_loader) + finally: + # Capture the output before restoring stdout + captured_output = string_buffer.getvalue() + sys.stdout = stdout finally: - # Capture the output before restoring stdout - captured_output = string_buffer.getvalue() - sys.stdout = stdout + if model_temp_dir is not None: + model_temp_dir.cleanup() if not results: logger.error(f"ONNXRuntime execution failed with output:\n{captured_output}") diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index f8b5a41a41a..aa7d27bc8fd 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -18,9 +18,9 @@ import copy import io import os +import sys import tempfile import uuid -from collections import defaultdict from typing import Any import numpy as np @@ -527,50 +527,114 @@ def name_onnx_nodes(graph: onnx.GraphProto) -> bool: def duplicate_shared_constants(onnx_model: onnx.ModelProto) -> tuple[onnx.ModelProto, bool]: """Duplicate constant tensors if they are shared.""" - graph = gs.import_onnx(onnx_model) - name_dict = defaultdict(lambda: 0) - - def _get_unique_name(old_name): - name_dict[old_name] += 1 - return old_name + "_" + str(name_dict[old_name]) - - # Get tensors with shared constant inputs - tensors = [] - for node in graph.nodes: - for inp_idx, tensor in enumerate(node.inputs): - # constant is shared across multiple nodes - if isinstance(tensor, Constant) and len(tensor.outputs) > 1: - tensors.append({"tensor": tensor, "inp_node": node, "inp_idx": inp_idx}) - - # Duplicate shared tensors - for tensor_dict in tensors: - tensor = tensor_dict["tensor"] - new_tensor = Constant( - name=_get_unique_name(tensor.name), - values=tensor.values, + graph = onnx_model.graph + initializers = {initializer.name: initializer for initializer in graph.initializer} + use_counts = dict.fromkeys(initializers, 0) + for node in graph.node: + for input_name in node.input: + if input_name in use_counts: + use_counts[input_name] += 1 + + shared_names = {name for name, count in use_counts.items() if count > 1} + if not shared_names: + return onnx_model, False + + used_names = set(initializers) + used_names.update( + sparse_initializer.values.name for sparse_initializer in graph.sparse_initializer + ) + for value_info in list(graph.input) + list(graph.output) + list(graph.value_info): + used_names.add(value_info.name) + for node in graph.node: + used_names.update(name for name in node.input if name) + used_names.update(name for name in node.output if name) + + def _get_nested_graphs(node: onnx.NodeProto): + for attribute in node.attribute: + if attribute.type == onnx.AttributeProto.GRAPH: + yield attribute.g + elif attribute.type == onnx.AttributeProto.GRAPHS: + yield from attribute.graphs + + def _find_captured_names(nested_graph: onnx.GraphProto, outer_names: set[str]) -> set[str]: + local_names = {value_info.name for value_info in nested_graph.input} + local_names.update(initializer.name for initializer in nested_graph.initializer) + local_names.update( + sparse_initializer.values.name for sparse_initializer in nested_graph.sparse_initializer ) - tensor_dict["inp_node"].inputs[tensor_dict["inp_idx"]] = new_tensor + for nested_node in nested_graph.node: + local_names.update(name for name in nested_node.output if name) - onnx_model = gs.export_onnx(graph) - is_modified = bool(tensors) - return onnx_model, is_modified + visible_outer_names = outer_names - local_names + captured_names = { + output.name for output in nested_graph.output if output.name in visible_outer_names + } + for nested_node in nested_graph.node: + captured_names.update( + input_name for input_name in nested_node.input if input_name in visible_outer_names + ) + for child_graph in _get_nested_graphs(nested_node): + captured_names.update(_find_captured_names(child_graph, visible_outer_names)) + return captured_names + captured_names = set() + for node in graph.node: + for nested_graph in _get_nested_graphs(node): + captured_names.update(_find_captured_names(nested_graph, shared_names)) -def check_model(model: onnx.ModelProto) -> None: - """Checks if the given model is valid.""" - save_as_external_data = False + next_suffix: dict[str, int] = {} + + def _get_unique_name(old_name: str) -> str: + suffix = next_suffix.get(old_name, 1) + new_name = f"{old_name}_{suffix}" + while new_name in used_names: + suffix += 1 + new_name = f"{old_name}_{suffix}" + next_suffix[old_name] = suffix + 1 + used_names.add(new_name) + return new_name + + for node in graph.node: + for input_index, input_name in enumerate(node.input): + if input_name not in shared_names: + continue + + duplicated_initializer = graph.initializer.add() + duplicated_initializer.CopyFrom(initializers[input_name]) + duplicated_initializer.name = _get_unique_name(input_name) + node.input[input_index] = duplicated_initializer.name + + retained_names = captured_names | { + output.name for output in graph.output if output.name in shared_names + } + removed_initializer_names = shared_names - retained_names + for initializer in initializers.values(): + if initializer.name in removed_initializer_names: + graph.initializer.remove(initializer) + for graph_input in list(graph.input): + if graph_input.name in removed_initializer_names: + graph.input.remove(graph_input) + + return onnx_model, True + + +def is_model_too_large_for_protobuf(model: onnx.ModelProto) -> bool: + """Return whether a model cannot safely use an in-memory protobuf API.""" try: model_size = model.ByteSize() except Exception as e: logger.warning( "Failed to compute model size with ByteSize (%s). Using external data path.", e ) - save_as_external_data = True - else: - if model_size <= 0 or model_size > (2 * (1024**3)): - save_as_external_data = True + return True - if save_as_external_data: + max_model_size = onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b"") + return model_size <= 0 or model_size > max_model_size + + +def check_model(model: onnx.ModelProto) -> None: + """Checks if the given model is valid.""" + if is_model_too_large_for_protobuf(model): with tempfile.TemporaryDirectory() as temp_dir: # ONNX also looks in CWD, so we need to use a unique id unique_id = str(uuid.uuid4())[:8] @@ -1171,19 +1235,7 @@ def infer_types_verification(model: onnx.ModelProto) -> onnx.ModelProto: def infer_shapes(model: onnx.ModelProto, **kwargs): """Infers shapes of the onnx graph, handles large models.""" - save_as_external_data = False - try: - model_size = model.ByteSize() - except Exception as e: - logger.warning( - "Failed to compute model size with ByteSize (%s). Using external data path.", e - ) - save_as_external_data = True - else: - if model_size <= 0 or model_size > (2 * (1024**3)): - save_as_external_data = True - - if save_as_external_data: + if is_model_too_large_for_protobuf(model): with tempfile.TemporaryDirectory() as temp_dir: # ONNX also looks in CWD, so we need to use a unique id unique_id = str(uuid.uuid4())[:8] diff --git a/tests/unit/onnx/autocast/test_autocast.py b/tests/unit/onnx/autocast/test_autocast.py index ccad0c7201b..1377b20af9f 100644 --- a/tests/unit/onnx/autocast/test_autocast.py +++ b/tests/unit/onnx/autocast/test_autocast.py @@ -27,6 +27,7 @@ from modelopt.onnx.autocast import convert_to_mixed_precision from modelopt.onnx.autocast.__main__ import get_parser, main from modelopt.onnx.autocast.convert import convert_to_f16 +from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer from modelopt.onnx.autocast.logging_config import configure_logging configure_logging("DEBUG") @@ -153,6 +154,42 @@ def test_convert_simple_model(temp_model_path, temp_output_path, keep_io_types): onnx.checker.check_model(loaded_model) +def test_convert_external_data_sanitizes_once_and_materializes_initializers( + tmp_path, simple_model, monkeypatch +): + model_path = tmp_path / "external_model.onnx" + onnx.save_model( + simple_model, + model_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="external_model.data", + size_threshold=0, + ) + + sanitize_calls = [] + original_sanitize = GraphSanitizer.sanitize + + def record_sanitize(sanitizer): + initializer = sanitizer.model.graph.initializer[0] + sanitize_calls.append( + (sanitizer.onnx_path, initializer.data_location, bool(initializer.raw_data)) + ) + return original_sanitize(sanitizer) + + monkeypatch.setattr(GraphSanitizer, "sanitize", record_sanitize) + + converted_model = convert_to_mixed_precision(onnx_path=str(model_path), data_max=np.inf) + + assert sanitize_calls == [(str(model_path.resolve()), onnx.TensorProto.EXTERNAL, False)] + for initializer in converted_model.graph.initializer: + assert initializer.data_location != onnx.TensorProto.EXTERNAL + assert not initializer.external_data + assert initializer.raw_data + assert initializer.data_type == onnx.TensorProto.FLOAT16 + onnx.checker.check_model(converted_model) + + def assert_input_precision(nodes, dtype="float16"): for node in nodes: for inp in node.inputs: @@ -199,11 +236,6 @@ def test_conv_resize_conversion(tmp_path): # Convert the model converted_model = convert_to_mixed_precision(onnx_path=onnx_path) - # Output model should be produced in the same tmp_path - output_onnx_path = onnx_path.replace(".onnx", ".fp16.onnx") - onnx.save(converted_model, output_onnx_path) - - # Load the output model graph = gs.import_onnx(converted_model) # Check that Resize is correctly converted: @@ -213,6 +245,7 @@ def test_conv_resize_conversion(tmp_path): assert all(inp.dtype == np.float16 for inp in resize_node.inputs[0:2]), ( "Resize data and ROI inputs should be FP16" ) + assert resize_node.inputs[1].name != resize_node.inputs[2].name @pytest.mark.parametrize("target_opset", [13, 17, 19, 21]) diff --git a/tests/unit/onnx/autocast/test_graphsanitizer.py b/tests/unit/onnx/autocast/test_graphsanitizer.py index cb487b56bf0..f42324a6338 100644 --- a/tests/unit/onnx/autocast/test_graphsanitizer.py +++ b/tests/unit/onnx/autocast/test_graphsanitizer.py @@ -13,10 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import Mock + import numpy as np +import onnx import pytest from onnx import TensorProto, helper, numpy_helper +import modelopt.onnx.autocast.graphsanitizer as graphsanitizer from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer @@ -407,3 +411,63 @@ def test_convert_fp64_no_changes_needed(): assert sanitizer._convert_fp64_initializers() is False assert sanitizer._convert_fp64_io_types() is False assert sanitizer._convert_fp64_nodes() is False + + +def test_sanitize_large_external_initializer_metadata(): + external_data_bytes = onnx.checker.MAXIMUM_PROTOBUF + 2048 + initializer = TensorProto( + name="weight", + data_type=TensorProto.FLOAT, + dims=[external_data_bytes // 4], + data_location=TensorProto.EXTERNAL, + ) + for key, value in ( + ("location", "weight.bin"), + ("offset", "0"), + ("length", str(external_data_bytes)), + ): + initializer.external_data.add(key=key, value=value) + + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, initializer.dims) + identity = helper.make_node("Identity", [initializer.name], [output.name]) + graph = helper.make_graph( + [identity], "large_external_initializer", [], [output], initializer=[initializer] + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 22)]) + model.ir_version = 10 + assert model.ByteSize() < 1024 + + sanitizer = GraphSanitizer(model, min_opset=22) + sanitizer.sanitize() + + initializer = sanitizer.model.graph.initializer[0] + assert initializer.data_location == TensorProto.EXTERNAL + assert not initializer.raw_data + assert {entry.key: entry.value for entry in initializer.external_data} == { + "location": "weight.bin", + "offset": "0", + "length": str(external_data_bytes), + } + assert [node.op_type for node in sanitizer.model.graph.node] == ["Identity"] + + +def test_find_custom_nodes_uses_source_model_path(tmp_path, monkeypatch): + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1]) + custom_node = helper.make_node("CustomOp", [x.name], [y.name], name="custom") + graph = helper.make_graph([custom_node], "custom_graph", [x], [y]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 22)]) + model_path = tmp_path / "custom.onnx" + tensor_info = {"Y": {"dtype": TensorProto.FLOAT, "shape": [1]}} + get_custom_layers = Mock(return_value=([custom_node.name], tensor_info)) + infer_types_shapes = Mock(return_value=model) + + monkeypatch.setattr(graphsanitizer, "set_trt_plugin_domain", Mock(return_value=model)) + monkeypatch.setattr(graphsanitizer, "get_custom_layers", get_custom_layers) + monkeypatch.setattr(graphsanitizer, "infer_types_shapes_tensorrt", infer_types_shapes) + + sanitizer = GraphSanitizer(model, min_opset=22, onnx_path=str(model_path)) + sanitizer.find_custom_nodes() + + get_custom_layers.assert_called_once_with(str(model_path.resolve()), []) + infer_types_shapes.assert_called_once_with(model, [], all_tensor_info=tensor_info) diff --git a/tests/unit/onnx/autocast/test_referencerunner.py b/tests/unit/onnx/autocast/test_referencerunner.py index 00075714c69..5c5c3c00ab2 100644 --- a/tests/unit/onnx/autocast/test_referencerunner.py +++ b/tests/unit/onnx/autocast/test_referencerunner.py @@ -17,9 +17,12 @@ import os import tempfile from collections import OrderedDict +from unittest.mock import Mock import numpy as np import onnx +import polygraphy.backend.onnx as polygraphy_onnx +import polygraphy.backend.onnxrt as polygraphy_onnxrt import pytest from onnx import TensorProto, helper @@ -67,6 +70,54 @@ def test_init(simple_model): assert isinstance(runner.model, onnx.ModelProto) +def test_get_ort_runner_uses_file_when_modified_byte_size_fails(monkeypatch, reference_runner): + modified_model = Mock() + modified_model.ByteSize.side_effect = ValueError("model size unavailable") + check_external_data = Mock(return_value=False) + save_onnx = Mock() + bytes_from_onnx = Mock(side_effect=AssertionError("Byte serialization should not be used")) + session_from_onnx = Mock(return_value="session") + onnxrt_runner = Mock(return_value="runner") + + monkeypatch.setattr(onnx_utils, "check_model_uses_external_data", check_external_data) + monkeypatch.setattr(onnx_utils, "save_onnx", save_onnx) + monkeypatch.setattr(polygraphy_onnx, "BytesFromOnnx", bytes_from_onnx) + monkeypatch.setattr(polygraphy_onnxrt, "SessionFromOnnx", session_from_onnx) + monkeypatch.setattr(polygraphy_onnxrt, "OnnxrtRunner", onnxrt_runner) + + runners, model_temp_dir = reference_runner._get_ort_runner(lambda: modified_model) + + assert model_temp_dir is not None + with model_temp_dir as model_temp_path: + save_onnx.assert_called_once() + model_path = save_onnx.call_args.args[1] + assert runners == ["runner"] + assert os.path.dirname(model_path) == model_temp_path + check_external_data.assert_called_once_with(modified_model) + modified_model.ByteSize.assert_called_once_with() + save_onnx.assert_called_once_with(modified_model, model_path, save_as_external_data=True) + bytes_from_onnx.assert_not_called() + session_from_onnx.assert_called_once_with(model_path, providers=reference_runner.providers) + onnxrt_runner.assert_called_once_with("session") + + +def test_run_cleans_model_tempdir_when_input_loading_fails(monkeypatch, reference_runner): + model_temp_dir = Mock() + monkeypatch.setattr( + reference_runner, "_get_ort_runner", Mock(return_value=([], model_temp_dir)) + ) + monkeypatch.setattr( + reference_runner, + "_load_inputs", + Mock(side_effect=ValueError("input loading failed")), + ) + + with pytest.raises(ValueError, match="input loading failed"): + reference_runner.run() + + model_temp_dir.cleanup.assert_called_once_with() + + def test_run_with_random_inputs(reference_runner): """Test running inference with random inputs.""" results = reference_runner.run() diff --git a/tests/unit/onnx/test_onnx_utils.py b/tests/unit/onnx/test_onnx_utils.py index 36face35b90..590a9f24186 100644 --- a/tests/unit/onnx/test_onnx_utils.py +++ b/tests/unit/onnx/test_onnx_utils.py @@ -14,6 +14,8 @@ # limitations under the License. import os +import sys +from unittest.mock import Mock import numpy as np import onnx @@ -28,6 +30,7 @@ make_tensor_value_info, ) +import modelopt.onnx.utils as onnx_utils from modelopt.onnx.trt_utils import load_onnx_model from modelopt.onnx.utils import ( clear_stale_value_info, @@ -56,6 +59,190 @@ def test_save_onnx(tmp_path): assert os.path.exists(os.path.join(tmp_path, "test.onnx")) +def _make_external_initializer( + name: str, location: str = "missing-shared-data.bin" +) -> onnx.TensorProto: + initializer = onnx.TensorProto(name=name, data_type=onnx.TensorProto.FLOAT, dims=[1]) + initializer.data_location = onnx.TensorProto.EXTERNAL + for key, value in [ + ("location", location), + ("offset", "0"), + ("length", "4"), + ]: + entry = initializer.external_data.add() + entry.key = key + entry.value = value + return initializer + + +def test_duplicate_shared_constants_preserves_and_materializes_external_data(tmp_path): + external_values = np.array([3.25], dtype=np.float32) + external_data_path = tmp_path / "shared.bin" + external_data_path.write_bytes(external_values.tobytes()) + shared = _make_external_initializer("shared", external_data_path.name) + collision = make_tensor("shared_1", onnx.TensorProto.FLOAT, [1], [0.0]) + sparse_collision = onnx.helper.make_sparse_tensor( + make_tensor("shared_2", onnx.TensorProto.FLOAT, [1], [0.0]), + make_tensor("", onnx.TensorProto.INT64, [1], [0]), + [1], + ) + nodes = [ + make_node("Add", ["input", "shared"], ["intermediate"]), + make_node("Add", ["intermediate", "shared"], ["output"]), + ] + graph = make_graph( + nodes, + "shared_external_initializer", + [ + make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1]), + make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]), + ], + [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])], + initializer=[shared, collision], + ) + graph.sparse_initializer.append(sparse_collision) + model = make_model(graph) + + result, modified = onnx_utils.duplicate_shared_constants(model) + + assert result is model + assert modified + duplicated_names = ("shared_3", "shared_4") + assert [node.input[1] for node in result.graph.node] == list(duplicated_names) + initializers = {initializer.name: initializer for initializer in result.graph.initializer} + assert set(initializers) == {"shared_1", *duplicated_names} + assert result.graph.sparse_initializer[0].values.name == "shared_2" + assert {graph_input.name for graph_input in result.graph.input} == {"input"} + for name in duplicated_names: + initializer = initializers[name] + assert initializer.data_location == onnx.TensorProto.EXTERNAL + assert [(entry.key, entry.value) for entry in initializer.external_data] == [ + ("location", "shared.bin"), + ("offset", "0"), + ("length", "4"), + ] + assert not initializer.HasField("raw_data") + + onnx.external_data_helper.load_external_data_for_model(result, str(tmp_path)) + for name in duplicated_names: + np.testing.assert_array_equal( + onnx.numpy_helper.to_array(initializers[name]), external_values + ) + + +def test_duplicate_shared_constants_fast_path_returns_original_model(): + shared = _make_external_initializer("single_use") + graph = make_graph( + [make_node("Add", ["input", "single_use"], ["output"])], + "single_external_initializer", + [make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1])], + [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])], + initializer=[shared], + ) + model = make_model(graph) + serialized_model = model.SerializeToString() + + result, modified = onnx_utils.duplicate_shared_constants(model) + + assert result is model + assert not modified + assert result.SerializeToString() == serialized_model + + +def test_duplicate_shared_constants_retains_initializer_captured_by_subgraphs(): + def make_branch(name): + output = make_tensor_value_info("branch_output", onnx.TensorProto.FLOAT, [1]) + return make_graph([make_node("Identity", ["shared"], [output.name])], name, [], [output]) + + shared = make_tensor("shared", onnx.TensorProto.FLOAT, [1], [1.0]) + nodes = [ + make_node("Add", ["input", "shared"], ["left"]), + make_node("Add", ["input", "shared"], ["right"]), + make_node( + "If", + ["condition"], + ["branch_value"], + then_branch=make_branch("then_branch"), + else_branch=make_branch("else_branch"), + ), + make_node("Sum", ["left", "right", "branch_value"], ["output"]), + ] + graph = make_graph( + nodes, + "captured_initializer", + [ + make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1]), + make_tensor_value_info("condition", onnx.TensorProto.BOOL, []), + make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]), + ], + [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])], + initializer=[shared], + ) + model = make_model(graph, opset_imports=[make_opsetid("", 17)]) + + result, modified = onnx_utils.duplicate_shared_constants(model) + + assert modified + assert [node.input[1] for node in result.graph.node[:2]] == ["shared_1", "shared_2"] + assert {initializer.name for initializer in result.graph.initializer} == { + "shared", + "shared_1", + "shared_2", + } + assert "shared" in {graph_input.name for graph_input in result.graph.input} + onnx.checker.check_model(result) + + +def test_duplicate_shared_constants_retains_initializer_exposed_as_graph_output(): + shared = make_tensor("shared", onnx.TensorProto.FLOAT, [1], [1.0]) + graph = make_graph( + [ + make_node("Add", ["input", "shared"], ["left"]), + make_node("Add", ["input", "shared"], ["right"]), + ], + "initializer_graph_output", + [make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1])], + [ + make_tensor_value_info("left", onnx.TensorProto.FLOAT, [1]), + make_tensor_value_info("right", onnx.TensorProto.FLOAT, [1]), + make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]), + ], + initializer=[shared], + ) + model = make_model(graph) + + result, modified = onnx_utils.duplicate_shared_constants(model) + + assert modified + assert [node.input[1] for node in result.graph.node] == ["shared_1", "shared_2"] + assert {initializer.name for initializer in result.graph.initializer} == { + "shared", + "shared_1", + "shared_2", + } + assert [output.name for output in result.graph.output] == ["left", "right", "shared"] + onnx.checker.check_model(result) + + +@pytest.mark.parametrize( + ("model_size", "expected"), + [ + (1, False), + (onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b""), False), + (onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b"") + 1, True), + (0, True), + (ValueError("model size unavailable"), True), + ], +) +def test_is_model_too_large_for_protobuf(model_size, expected): + model = Mock() + if isinstance(model_size, Exception): + model.ByteSize.side_effect = model_size + else: + model.ByteSize.return_value = model_size + assert onnx_utils.is_model_too_large_for_protobuf(model) is expected + + def make_onnx_model_for_matmul_op(): input_left = np.array([1, 2]) input_right = np.array([1, 3])