Add onnxsim as an alternative ONNX simplification backend

CHANGELOG.rst

@@ -6,6 +6,10 @@ Changelog
 
 **New Features**
 
+*Misc*
+
+- Add ``onnxsim`` as an alternative ONNX simplification backend for ``modelopt.onnx.quantization.quantize(..., simplify=True)`` (and the ``--simplify`` CLI flag). The new ``simplify_backend`` argument / ``--simplify_backend`` flag selects between ``"onnxslim"`` (default, unchanged) and ``"onnxsim"``; both produce an equivalent simplified model. ``onnxsim>=0.7.0`` now ships wheels for Python 3.12+ and aarch64.
+
 **Backward Breaking Changes**
 
 **Deprecations**

modelopt/onnx/quantization/__main__.py

@@ -327,6 +327,16 @@ def get_parser() -> argparse.ArgumentParser:
         action="store_true",
         help="If True, the given ONNX model will be simplified before quantization is performed.",
     )
+    argparser.add_argument(
+        "--simplify_backend",
+        type=str,
+        default="onnxslim",
+        choices=["onnxslim", "onnxsim"],
+        help=(
+            "ONNX simplification package to use when --simplify is set. "
+            "Both produce an equivalent simplified model (default: onnxslim)."
+        ),
+    )
     argparser.add_argument(
         "--calibrate_per_node",
         action="store_true",
@@ -552,6 +562,7 @@ def main():
         use_zero_point=args.use_zero_point,
         passes=args.passes,
         simplify=args.simplify,
+        simplify_backend=args.simplify_backend,
         calibrate_per_node=args.calibrate_per_node,
         direct_io_types=args.direct_io_types,
         opset=args.opset,

modelopt/onnx/quantization/quantize.py

@@ -122,6 +122,7 @@ def _preprocess_onnx(
     trt_plugins_precision: list[str] | None,
     override_shapes: str,
     simplify: bool = False,
+    simplify_backend: str = "onnxslim",
     quantize_mode: str = "int8",
     opset: int | None = None,
 ) -> tuple[str, onnx.ModelProto, list[str], bool, bool, bool, dict, dict]:
@@ -218,10 +219,32 @@ def _preprocess_onnx(
 
     # Simplify model if requested
     if simplify:
-        logger.info("Attempting to simplify model")
+        logger.info(f"Attempting to simplify model with '{simplify_backend}'")
+
+        # Resolve the backend before attempting simplification so that a missing
+        # optional dependency or an unknown backend name fails loudly instead of
+        # being silently swallowed by the graceful fallback below.
+        if simplify_backend == "onnxsim":
+            try:
+                import onnxsim
+            except ModuleNotFoundError as e:
+                logger.warning(
+                    "onnxsim is not installed. Please install it with 'pip install onnxsim'."
+                )
+                raise e
+        elif simplify_backend != "onnxslim":
+            raise ValueError(
+                f"Unsupported simplify_backend '{simplify_backend}'. "
+                "Choose one of: 'onnxslim', 'onnxsim'."
+            )
+
         try:
-            model_simp = onnxslim.slim(onnx_model, skip_fusion_patterns=["FusionGemm"])
-            if model_simp:
+            if simplify_backend == "onnxslim":
+                model_simp = onnxslim.slim(onnx_model, skip_fusion_patterns=["FusionGemm"])
+                check = model_simp is not None
+            else:
+                model_simp, check = onnxsim.simplify(onnx_model)
+            if check:
                 onnx_model = model_simp
                 onnx_path = os.path.join(output_dir, f"{model_name}_simp.onnx")
                 save_onnx(onnx_model, onnx_path, use_external_data_format)
@@ -395,6 +418,7 @@ def quantize(
     autotune_warmup_runs: int = 50,
     autotune_timing_runs: int = 100,
     autotune_trtexec_args: str | None = None,
+    simplify_backend: str = "onnxslim",
     **kwargs: Any,
 ) -> None:
     """Quantizes the provided ONNX model.
@@ -558,6 +582,9 @@ def quantize(
         autotune_trtexec_args:
             Additional trtexec arguments as a single quoted string.
             Example: --autotune_trtexec_args '--fp16 --workspace=4096'
+        simplify_backend:
+            ONNX simplification package to use when ``simplify`` is set. One of ``"onnxslim"``
+            (default) or ``"onnxsim"``; both produce an equivalent simplified model.
         kwargs:
             Additional keyword arguments for int4 quantization, including:
             - awqlite_alpha_step (float): Alpha step for lite, range [0, 1].
@@ -627,6 +654,7 @@ def quantize(
         trt_plugins_precision,
         override_shapes,  # type: ignore[arg-type]
         simplify,
+        simplify_backend,
         quantize_mode,
         opset,
     )

pyproject.toml

@@ -69,6 +69,7 @@ onnx = [
     "onnxruntime~=1.24.2; python_version > '3.10' and (platform_machine == 'aarch64' or platform_system == 'Darwin')",
     "onnxruntime-gpu~=1.24.2; python_version > '3.10' and platform_machine != 'aarch64' and platform_system != 'Darwin' and platform_system != 'Windows'",
     "onnxscript",
+    "onnxsim>=0.7.0",
     "onnxslim>=0.1.76",
     "polygraphy>=0.49.22",
 ]

tests/unit/onnx/quantization/test_quantize_api.py

@@ -89,6 +89,7 @@ def fake_preprocess(
         trt_plugins_precision,
         override_shapes,
         simplify,
+        simplify_backend,
         quantize_mode,
         opset,
     ):