Protobuf 7.34 fix for model size > 2gib (#1058)

### What does this PR do?

Type of change:  Bug fix

For the newer version of protobuf when model size > 2gb the model throws
an error when serializing to string specifically in windows.

```
Traceback (most recent call last):
  File "quantize.py", line 643, in <module>
    main(args)
  File "quantize.py", line 428, in main
    quantized_onnx_model = quantize_int4(
        args.onnx_path,
        ...
        use_column_major=args.use_column_major,
    )
  File ".venv\Lib\site-packages\modelopt\onnx\quantization\int4.py", line 1529, in quantize
    onnx_model = _quantize_awq_lite(
        onnx_model,
        ...
    )
  File ".venv\Lib\site-packages\modelopt\onnx\quantization\int4.py", line 1088, in _quantize_awq_lite
    save_onnx(augmented_model, augmented_onnx_path, use_external_data_format)
  File ".venv\Lib\site-packages\modelopt\onnx\utils.py", line 646, in save_onnx
    model_proto = model.SerializeToString()
google.protobuf.message.EncodeError: Failed to serialize proto
```

See Nvbugid 5989474 for more context

Suggested fix is to force it to save as external data when this happens.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Improved handling of model size checks so failures no longer crash
workflows—serialization attempts now log warnings and fall back to
external-data storage.
* Loading now tolerates size-query errors and enables external-data mode
when size is invalid or exceeds limits.

* **Improvements**
* Resilient opset upgrade path: version-conversion failures log a
warning and apply a safe opset import fallback before saving.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

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

Author: hthadicherla

modelopt/onnx/quantization/quantize.py

@@ -176,7 +176,21 @@ def _preprocess_onnx(
         )
 
     if original_opset_version < target_opset and original_opset_version != 1:
-        onnx_model = onnx.version_converter.convert_version(onnx_model, target_opset)
+        try:
+            onnx_model = onnx.version_converter.convert_version(onnx_model, target_opset)
+        except Exception as e:
+            logger.warning(
+                "onnx.version_converter failed (%s). Performing lightweight opset update.", e
+            )
+            current_opset = {opset.domain: opset.version for opset in onnx_model.opset_import}
+            new_opset_imports = [onnx.helper.make_opsetid("", target_opset)]
+            if "com.microsoft" not in current_opset:
+                new_opset_imports.append(onnx.helper.make_opsetid("com.microsoft", 1))
+            for domain, version in current_opset.items():
+                if domain not in ["", "ai.onnx"]:
+                    new_opset_imports.append(onnx.helper.make_opsetid(domain, version))
+            onnx_model.ClearField("opset_import")
+            onnx_model.opset_import.extend(new_opset_imports)
         onnx_path = os.path.join(output_dir, f"{model_name}_opset{target_opset}.onnx")
         save_onnx(onnx_model, onnx_path, use_external_data_format)
         logger.info(f"Model is cloned to {onnx_path} with opset_version {target_opset}")

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

@@ -558,7 +558,19 @@ def _get_unique_name(old_name):
 
 def check_model(model: onnx.ModelProto) -> None:
     """Checks if the given model is valid."""
-    if model.ByteSize() > (2 * (1024**3)):  # 2GB limit
+    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:
         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]
@@ -642,21 +654,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 as e:
-        if "Message onnx.ModelProto exceeds maximum protobuf size of 2GB" in str(e):
-            logger.warning("Model exceeds 2GB limit, switching to external data storage")
+    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:
-            logger.error(f"Failed to serialize model: {e!s}")
-            raise
+            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
@@ -1162,7 +1171,19 @@ 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."""
-    if model.ByteSize() > (2 * (1024**3)):  # 2GB limit
+    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:
         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]