[6410139] Fix ONNX AutoCast for large external initializers

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.

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

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)

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

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}")