call sanitizer in precision converter

Signed-off-by: Ali Boubezari <[email protected]>

Author: aboubezari

modelopt/onnx/autocast/precisionconverter.py

@@ -32,6 +32,7 @@
 
 import modelopt.onnx.autocast.utils as utils
 import modelopt.onnx.utils as onnx_utils
+from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer
 from modelopt.onnx.autocast.logging_config import configure_logging, logger
 
 configure_logging()
@@ -73,6 +74,9 @@ def __init__(
         low_precision_type: str = "fp16",
         init_conversion_max_bytes: int | None = None,
         custom_ops: set[str] | None = None,
+        min_opset: int = 13,
+        max_ir_version: int | None = None,
+        trt_plugins: list[str] | None = [],
     ) -> None:
         """Initialize PrecisionConverter.
 
@@ -109,6 +113,9 @@ def __init__(
         self.original_network_io.update(
             {io.name: io.type.tensor_type.elem_type for io in self.model.graph.output}
         )
+        self.min_opset = min_opset
+        self.max_ir_version = max_ir_version
+        self.trt_plugins = trt_plugins
 
     def convert(
         self,
@@ -132,6 +139,8 @@ def convert(
                 "AutoCast can only operate on valid ONNX models, but the input model is invalid. See log for details."
             )
 
+        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
         high_precision_nodes, low_precision_nodes = self._filter_unsupported_op_types(
@@ -1030,3 +1039,13 @@ def _is_foldable_constant_cast_pattern(self, node: onnx.NodeProto) -> bool:
             get_consumer_nodes = utils.get_consumer_nodes(self.model, const_producer.output[0])
             return len(get_consumer_nodes) == 1 and get_consumer_nodes[0] == node
         return False
+
+    def _sanitize_model(self):
+        graph_sanitizer = GraphSanitizer(
+            self.model,
+            self.min_opset,
+            trt_plugins=self.trt_plugins,
+            max_ir_version=self.max_ir_version,
+        )
+        graph_sanitizer.sanitize()
+        self.model = graph_sanitizer.model

tests/unit/onnx/autocast/test_precisionconverter.py

@@ -20,7 +20,6 @@
 
 import modelopt.onnx.autocast.utils as utils
 import modelopt.onnx.utils as onnx_utils
-from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer
 from modelopt.onnx.autocast.logging_config import configure_logging
 from modelopt.onnx.autocast.precisionconverter import PrecisionConverter
 
@@ -31,6 +30,9 @@ def low_precision_onnx_type(low_precision_type_str):
     return TensorProto.FLOAT16 if low_precision_type_str == "fp16" else TensorProto.BFLOAT16
 
 
+LATEST_IR_VERSION_SUPPORTED_BY_ORT = 10
+
+
 ####################################################################################################
 # Testing with a basic GEMM->Add->Relu graph
 ####################################################################################################
@@ -1079,17 +1081,16 @@ def test_casted_input_to_output_model(
 ):
     model, value_info_map, initializer_map, node_to_init_map = model_with_casted_input_to_output
 
-    min_opset = 22 if low_precision_type == "bf16" else 13
-    graph_sanitizer = GraphSanitizer(model, min_opset)
-    graph_sanitizer.sanitize()
-    model = graph_sanitizer.model
     converter = PrecisionConverter(
         model,
         value_info_map,
         initializer_map,
         node_to_init_map,
         keep_io_types=keep_io_types,
         low_precision_type=low_precision_type,
+        min_opset=22 if low_precision_type == "bf16" else 13,
+        max_ir_version=LATEST_IR_VERSION_SUPPORTED_BY_ORT,
+        trt_plugins=[],
     )
     converted_model = converter.convert(
         high_precision_nodes=["cast_input"], low_precision_nodes=["add1", "add2"]