Add onnxsim as an alternative ONNX simplification backend

[take-cheeze]

What does this PR do?

Type of change: New feature

Adds onnxsim as an alternative ONNX simplification backend for
modelopt.onnx.quantization.quantize(..., simplify=True) and the --simplify
CLI flag, while keeping onnxslim as the default. Previously the simplifier
was hard-wired to onnxslim; this makes the two interchangeable via a new
simplify_backend argument (--simplify_backend on the CLI) that accepts
"onnxslim" (default) or "onnxsim".

Details:

• quantize.py: quantize() and _preprocess_onnx() gain a
  simplify_backend: str = "onnxslim" argument. The simplify block dispatches
  on it — onnxslim.slim(..., skip_fusion_patterns=["FusionGemm"]) for
  onnxslim, onnxsim.simplify(...) for onnxsim. An unknown backend name or
  a missing onnxsim install fails loudly (raised before the guarded block),
  whereas genuine simplification failures still fall back gracefully to the
  original model, exactly as before. onnxsim is imported lazily so it is only
  required when actually selected.
• __main__.py: adds --simplify_backend with choices=["onnxslim", "onnxsim"],
  default onnxslim, wired into the quantize(...) call.
• pyproject.toml / uv.lock: adds onnxsim>=0.7.0 to the onnx extra
  alongside the existing onnxslim>=0.1.76. onnxsim 0.7.0 ships wheels for
  Python 3.12+ and aarch64, so no platform markers are needed.

Usage

from modelopt.onnx.quantization.quantize import quantize

# Default behavior is unchanged (uses onnxslim):
quantize("model.onnx", simplify=True)

# Opt into onnxsim instead:
quantize("model.onnx", simplify=True, simplify_backend="onnxsim")

# CLI — default (onnxslim):
python -m modelopt.onnx.quantization --onnx_path model.onnx --simplify

# CLI — onnxsim backend:
python -m modelopt.onnx.quantization --onnx_path model.onnx --simplify --simplify_backend onnxsim

Testing

• Updated the _preprocess_onnx mock in
  tests/unit/onnx/quantization/test_quantize_api.py for the new argument;
  the existing unit suite for the quantize API continues to pass.
• The existing GPU test tests/gpu/onnx/test_simplify.py exercises the
  default (onnxslim) simplify path end-to-end.
• Locally verified pyproject.toml and uv.lock parse, the lock resolves
  onnxsim 0.7.0 (deps onnx, rich) with onnxslim untouched, and the
  changed Python files pass ruff format --check.
• Not run in this environment: the GPU simplify test on the onnxsim path
  (needs a GPU and onnxsim installed).

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ — simplify defaults to the existing
  onnxslim backend; no existing call site or CLI invocation changes behavior.
• If you copied code from any other sources or added a new PIP dependency, did
  you follow guidance in CONTRIBUTING.md: ✅ — added onnxsim to the optional
  onnx extra and relocked uv.lock; no code copied.
• Did you write any new necessary tests?: ❌ — updated the existing
  _preprocess_onnx mock only; a dedicated simplify_backend-selection unit
  test is not yet added (happy to add one covering the dispatch and the
  invalid-backend ValueError if reviewers want it).
• Did you update Changelog?: ✅ — added an entry under the 0.46 New Features
  section.
• Did you get Claude approval on this PR?: N/A — will run /claude review once
  opened.

Additional Information

No related issue. Default simplification remains onnxslim; onnxsim is purely
additive and opt-in.

Summary by CodeRabbit

• New Features
  • Added a --simplify_backend CLI option (and simplify_backend parameter) to choose between onnxslim and onnxsim when model simplification is enabled.
  • Documented and enabled onnxsim>=0.7.0 support for relevant Python versions and architectures.
• Bug Fixes
  • Improved validation and error handling for unsupported or missing simplification backends.
  • Simplified models are now persisted only after successful validation.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ad6c93bc-b2c3-4dc5-997d-7bfe1a2eeba2

📥 Commits

Reviewing files that changed from the base of the PR and between d98e6c9 and 8c7f03c.

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (5)

• CHANGELOG.rst
• modelopt/onnx/quantization/__main__.py
• modelopt/onnx/quantization/quantize.py
• pyproject.toml
• tests/unit/onnx/quantization/test_quantize_api.py

🚧 Files skipped from review as they are similar to previous changes (5)

• tests/unit/onnx/quantization/test_quantize_api.py
• pyproject.toml
• modelopt/onnx/quantization/main.py
• CHANGELOG.rst
• modelopt/onnx/quantization/quantize.py

---

📝 Walkthrough

Walkthrough

The quantization workflow adds selectable ONNX simplification backends, exposing onnxslim and onnxsim through the API and CLI. Backend validation, optional dependency packaging, tests, and changelog documentation are updated.

Changes

ONNX simplification backend selection

Layer / File(s)	Summary
Backend API and simplification execution modelopt/onnx/quantization/quantize.py, tests/unit/onnx/quantization/test_quantize_api.py	quantize accepts and forwards simplify_backend; preprocessing validates and invokes onnxslim or onnxsim, while the test mock matches the updated call signature.
CLI backend selection modelopt/onnx/quantization/__main__.py	Adds --simplify_backend with onnxslim as the default and forwards the selected value to quantize.
onnxsim packaging and changelog pyproject.toml, CHANGELOG.rst	Adds onnxsim>=0.7.0 to the ONNX optional dependencies and documents the backend selector and wheel support.

Estimated code review effort: 2 (Simple) | ~10 minutes

Sequence Diagram(s)

sequenceDiagram
  participant User
  participant CLI
  participant quantize
  participant _preprocess_onnx
  participant onnxslim_or_onnxsim
  User->>CLI: select --simplify_backend
  CLI->>quantize: pass simplify_backend
  quantize->>_preprocess_onnx: request simplification
  _preprocess_onnx->>onnxslim_or_onnxsim: invoke selected backend
  onnxslim_or_onnxsim-->>_preprocess_onnx: return simplified model

Loading

Suggested reviewers: kevalmorabia97, gcunhase

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: adding onnxsim as an alternative ONNX simplification backend.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	No changed modelopt/example Python lines add banned patterns; the new onnxsim dependency is permissively licensed.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Warning

CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.

Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.

👉 Steps to fix this

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/unit/onnx/quantization/test_quantize_api.py (1)

92-95: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Exercise backend selection rather than only updating mock arity.

Capture and assert the forwarded default, then add focused tests for onnxsim, unsupported backends, and missing optional dependencies.

Minimal assertion

         simplify_backend,
         quantize_mode,
         opset,
     ):
+        captured["simplify_backend"] = simplify_backend
         return onnx_path, object(), [], True, False, False, {}, {}
...
+    assert captured["simplify_backend"] == "onnxslim"

As per path instructions, tests must add coverage for new features and exercise the behavior they validate.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/onnx/quantization/test_quantize_api.py` around lines 92 - 95,
Expand the quantization API tests around the backend-selection parameter in the
test function receiving simplify_backend, quantize_mode, and opset: capture and
assert its forwarded default, then add focused coverage for the onnxsim backend,
unsupported backend handling, and missing optional dependencies. Ensure the
tests exercise the actual selection behavior rather than only adjusting mock
argument arity.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@modelopt/onnx/quantization/quantize.py`:
- Line 401: Reorder the new simplify_backend parameter in the quantization
function signature so it appears after autotune_trtexec_args, preserving the
existing positional binding of later parameters such as calibrate_per_node. Keep
its default value unchanged and update any related signature handling
accordingly.

---

Nitpick comments:
In `@tests/unit/onnx/quantization/test_quantize_api.py`:
- Around line 92-95: Expand the quantization API tests around the
backend-selection parameter in the test function receiving simplify_backend,
quantize_mode, and opset: capture and assert its forwarded default, then add
focused coverage for the onnxsim backend, unsupported backend handling, and
missing optional dependencies. Ensure the tests exercise the actual selection
behavior rather than only adjusting mock argument arity.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 22ab02b7-001d-4812-ae75-19e815c5d0fc

📥 Commits

Reviewing files that changed from the base of the PR and between 33d05b0 and d98e6c9.

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (5)

• CHANGELOG.rst
• modelopt/onnx/quantization/__main__.py
• modelopt/onnx/quantization/quantize.py
• pyproject.toml
• tests/unit/onnx/quantization/test_quantize_api.py