code de-deuplication

Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>

Author: galagam

modelopt/onnx/autocast/referencerunner.py

@@ -128,10 +128,7 @@ def _get_ort_runner(self, model):
         # 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)
-        has_external_data = any(
-            init.HasField("data_location") and init.data_location == onnx.TensorProto.EXTERNAL
-            for init in self.model.graph.initializer
-        )
+        has_external_data = onnx_utils.check_model_uses_external_data(self.model)
 
         # Also check if model would be too large (>2GB) for SerializeToString
         # This handles cases where model was loaded with external data already loaded

modelopt/onnx/utils.py

@@ -696,6 +696,21 @@ def get_opset_version(model: onnx.ModelProto) -> int:
     return ai_onnx_domain[0].version
 
 
+def check_model_uses_external_data(model: onnx.ModelProto) -> bool:
+    """Checks if the model uses external data.
+
+    Args:
+        model: Loaded in-memory onnx ModelProto.
+
+    Returns:
+        True if any initializer tensor has data_location set to EXTERNAL.
+    """
+    return any(
+        init.HasField("data_location") and init.data_location == onnx.TensorProto.EXTERNAL
+        for init in model.graph.initializer
+    )
+
+
 def bfloat16_to_float32(bf16_array):
     """Converts a bfloat16 array (as raw data) to a float32 array."""
     uint32_array = bf16_array.astype(np.uint32) << 16

modelopt/torch/_deploy/utils/onnx_utils.py

@@ -45,14 +45,3 @@ def _get_onnx_external_data_tensors(model: onnx.ModelProto) -> list[str]:
         if tensor.HasField("data_location") and tensor.data_location == onnx.TensorProto.EXTERNAL
     ]
     return model_tensors_ext
-
-
-def check_model_uses_external_data(model: onnx.ModelProto) -> bool:
-    """
-    Checks if the model uses external data.
-    """
-    model_tensors = _get_initializer_tensors(model)
-    return any(
-        tensor.HasField("data_location") and tensor.data_location == onnx.TensorProto.EXTERNAL
-        for tensor in model_tensors
-    )

modelopt/torch/_deploy/utils/torch_onnx.py

@@ -42,6 +42,7 @@
 )
 from modelopt.onnx.quantization.qdq_utils import qdq_to_dq, replace_zero_scale_with_smallest_nonzero
 from modelopt.onnx.utils import (
+    check_model_uses_external_data,
     get_input_names,
     get_input_shapes,
     get_node_names,
@@ -55,7 +56,6 @@
 from modelopt.torch.utils._pytree import TreeSpec
 
 from ..utils.onnx_optimizer import Optimizer
-from .onnx_utils import check_model_uses_external_data
 
 ModelMetadata = dict[str, Any]
 ModelType = Any