Added code changes to have a pathway for when serialize succedes and when serialize fails

Signed-off-by: Hrishith Thadicherla <hthadicherla@nvidia.com>

Author: hthadicherla

modelopt/onnx/trt_utils.py

@@ -322,8 +322,22 @@ def load_onnx_model(
 
     # Load the model and weights
     onnx_model = onnx.load(onnx_path, load_external_data=True)
-    size_threshold = 2 * (1024**3)  # 2GB
-    use_external_data_format = onnx_model.ByteSize() > size_threshold or use_external_data_format
+    if not use_external_data_format:
+        try:
+            model_size = onnx_model.ByteSize()
+        except Exception as e:
+            logger.warning(
+                "Failed to compute model size with ByteSize (%s). Saving tensors as external data.",
+                e,
+            )
+            use_external_data_format = True
+        else:
+            if model_size <= 0 or model_size >= onnx.checker.MAXIMUM_PROTOBUF:
+                use_external_data_format = True
+                logger.debug(
+                    "Model is too large to save as a single file but 'use_external_data_format'"
+                    " is False. Saving tensors as external data, regardless."
+                )
 
     # If inputs are dynamic and override shapes are given, set them as static
     dynamic_inputs = get_dynamic_graph_inputs(onnx_model)

modelopt/onnx/utils.py

@@ -23,7 +23,6 @@
 from collections import defaultdict
 from typing import Any
 
-import google.protobuf.message
 import numpy as np
 import onnx
 import onnx_graphsurgeon as gs
@@ -643,19 +642,18 @@ def get_variable_inputs(node: Node) -> list[Variable]:
 def save_onnx(model: onnx.ModelProto, onnx_path: str, save_as_external_data: bool = False):
     """Save an ONNX model to given path. If a model is larger than 2GB, will save with external data."""
     size_threshold = 2 * (1024**3)  # 2GB
-    try:
-        model_proto = model.SerializeToString()
-        model_size = len(model_proto)
-        save_as_external_data = save_as_external_data or model_size > size_threshold
-        logger.debug(
-            f"Model size: {model_size} bytes, using external data: {save_as_external_data}"
-        )
-
-    except (ValueError, google.protobuf.message.EncodeError) as e:
-        logger.warning(
-            "Model exceeds 2GB limit, switching to external data storage. Error message: [%s]", e
-        )
-        save_as_external_data = True
+    if not save_as_external_data:
+        try:
+            model_proto = model.SerializeToString()
+        except Exception as e:
+            logger.warning("Failed to serialize model. Saving tensors as external data. (%s)", e)
+            save_as_external_data = True
+        else:
+            model_size = len(model_proto)
+            save_as_external_data = model_size > size_threshold
+            logger.debug(
+                f"Model size: {model_size} bytes, using external data: {save_as_external_data}"
+            )
 
     # Set ir_version to 10, remove it once ORT supports ir_version 11
     model.ir_version = 10