Add onnxsim as an alternative ONNX simplification backend - #2018
Add onnxsim as an alternative ONNX simplification backend#2018take-cheeze wants to merge 1 commit into
Conversation
📝 WalkthroughWalkthroughThe quantization workflow adds selectable ONNX simplification backends. The API and CLI support ChangesONNX simplification backend selection
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
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
345acfa to
d98e6c9
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unit/onnx/quantization/test_quantize_api.py (1)
92-95: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExercise 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
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (5)
CHANGELOG.rstmodelopt/onnx/quantization/__main__.pymodelopt/onnx/quantization/quantize.pypyproject.tomltests/unit/onnx/quantization/test_quantize_api.py
d98e6c9 to
8c7f03c
Compare
|
Hi @take-cheeze, can you please elaborate on the motivation for this? Are you observing issues with |
|
I know the situation in #478 , though onnx-simplifier is now maintained. |
@take-cheeze If onnxsim is now well-maintained, I would certainly pick onnxsim up, it's the only one. |
Do you mean we should take |
Keep onnxslim as the default simplifier and offer onnxsim as an equivalent, user-selectable option for modelopt.onnx.quantization.quantize(..., simplify=True) and the --simplify CLI flag. - quantize.py / __main__.py: add a simplify_backend argument (--simplify_backend) that selects between "onnxslim" (default, unchanged behavior) and "onnxsim". An unknown backend name or a missing onnxsim install fails loudly; genuine simplification failures still fall back gracefully to the original model. - pyproject.toml / uv.lock: add onnxsim>=0.7.0 to the onnx extra alongside onnxslim. onnxsim 0.7.0 ships wheels for Python 3.12+ and aarch64. - test_quantize_api.py: update the _preprocess_onnx mock signature for the new argument. - CHANGELOG.rst: document the new backend option under 0.47. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PYET7LC6SwmWpkCSN18T1Z Signed-off-by: Claude <noreply@anthropic.com>
8c7f03c to
f07c462
Compare
There was a problem hiding this comment.
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.
Actionable comments posted: 1
🤖 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 `@CHANGELOG.rst`:
- Line 63: Update the changelog entry for the simplify_backend option to qualify
the backend equivalence claim: state that both onnxslim and onnxsim aim to
preserve model semantics, while allowing their generated graphs to differ. Keep
the existing default backend and CLI/API option details unchanged.
🪄 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: 6dc045b0-eec4-495a-90c5-4fe93c8440cf
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (5)
CHANGELOG.rstmodelopt/onnx/quantization/__main__.pymodelopt/onnx/quantization/quantize.pypyproject.tomltests/unit/onnx/quantization/test_quantize_api.py
🚧 Files skipped from review as they are similar to previous changes (4)
- pyproject.toml
- tests/unit/onnx/quantization/test_quantize_api.py
- modelopt/onnx/quantization/main.py
- modelopt/onnx/quantization/quantize.py
| - Add a ``constant_amax`` ``QuantizerAttributeConfig`` field that pins a quantizer's ``amax`` to a fixed value and skips activation calibration (no forward statistics collected). Unlike ``use_constant_amax`` (which hardcodes 448.0 for KV-cache cast math and registers no buffer), ``constant_amax`` stores the constant on the ``_amax`` buffer so it is used by both the fake-quant forward and the exported scaling factor. For NVFP4 activations the exported ``input_scale`` equals ``amax / (E2M1_MAX * E4M3_MAX)``, so ``constant_amax: 2688.0`` yields ``input_scale == 1.0``. Ships a new recipe ``modelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml`` that applies this to the MoE expert activation quantizers. | ||
| - Add ``MaxCalibConfig.skip_forward_without_activation_calib`` (opt-in, default ``False``): when enabled, max calibration skips the ``forward_loop`` if no enabled quantizer needs data-driven activation statistics — e.g. an experts-only recipe whose activation quantizers all use ``constant_amax`` / ``use_constant_amax``, or dynamic / MX (MXFP4/MXFP8) quantization (and none has a static ``bias_calibrator``). Weight calibration still runs on the weight tensors directly, so the quantized weights are unchanged; only the wasted forward is avoided. It is opt-in because the ``forward_loop`` can carry caller-side effects (notably materializing sharded parameters under DeepSpeed ZeRO-3). The advanced algorithms that always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) call ``max_calibrate`` directly and are unaffected. The ``nvfp4_experts_only_input_scale1-kv_fp8_cast`` recipe enables it. | ||
| - Add ``examples/minimax_m3/hf_ptq_mixed_mxfp8_nvfp4.py`` for streaming MiniMax-M3 export and a model-specific ``hf_ptq.py`` recipe that produces an MXFP8 language-model base with MSE-calibrated NVFP4 routed experts directly from BF16. The NVFP4 expert ``input_scale`` is fixed to 1.0. | ||
| - 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. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- relevant files ---'
git ls-files | rg '(^|/)(CHANGELOG\.rst|.*onnx.*|.*quantiz.*|.*simplif.*)' | head -200
printf '%s\n' '--- changelog context ---'
cat -n CHANGELOG.rst | sed -n '45,75p'
printf '%s\n' '--- backend references ---'
rg -n --hidden -g '!CHANGELOG.rst' 'simplify_backend|onnxslim|onnxsim|simplify=True|--simplify' . | head -250Repository: NVIDIA/Model-Optimizer
Length of output: 30895
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- quantize implementation ---'
cat -n modelopt/onnx/quantization/quantize.py | sed -n '105,145p;205,260p;400,435p;570,600p;640,665p'
printf '%s\n' '--- CLI definition ---'
cat -n modelopt/onnx/quantization/__main__.py | sed -n '315,345p'
printf '%s\n' '--- simplification tests ---'
cat -n tests/gpu/onnx/test_simplify.py | sed -n '1,130p'
printf '%s\n' '--- API tests ---'
cat -n tests/unit/onnx/quantization/test_quantize_api.py | sed -n '1,125p'
printf '%s\n' '--- test and source inventory ---'
rg -n 'simplify_backend|onnxsim|onnxslim|equivalent simplified|semantics|checker' tests modelopt/onnx/quantization CHANGELOG.rstRepository: NVIDIA/Model-Optimizer
Length of output: 33029
🏁 Script executed:
#!/bin/bash
set -e
python3 - <<'PY'
import ast
from pathlib import Path
source_path = Path("modelopt/onnx/quantization/quantize.py")
tree = ast.parse(source_path.read_text())
calls = []
for node in ast.walk(tree):
if isinstance(node, ast.Call):
fn = node.func
if isinstance(fn, ast.Attribute) and isinstance(fn.value, ast.Name):
if (fn.value.id, fn.attr) in {("onnxslim", "slim"), ("onnxsim", "simplify")}:
calls.append((fn.value.id, fn.attr, ast.unparse(node)))
print("simplifier_calls:")
for item in calls:
print(item)
test_text = Path("tests/gpu/onnx/test_simplify.py").read_text()
api_text = Path("tests/unit/onnx/quantization/test_quantize_api.py").read_text()
for name, text in [("tests/gpu/onnx/test_simplify.py", test_text),
("tests/unit/onnx/quantization/test_quantize_api.py", api_text)]:
print(f"{name}:")
print(" mentions simplify_backend =", "simplify_backend" in text)
print(" mentions onnxsim =", "onnxsim" in text)
print(" calls quantize with simplify=True =", "simplify=True" in text)
PY
printf '%s\n' '--- dependency metadata context ---'
cat -n pyproject.toml | sed -n '60,80p'
printf '%s\n' '--- historical changelog context ---'
cat -n CHANGELOG.rst | sed -n '330,347p'Repository: NVIDIA/Model-Optimizer
Length of output: 3085
Qualify the equivalence claim. onnxslim and onnxsim use different simplification calls, and no test compares their outputs. State that both aim to preserve model semantics, but the generated graph may differ by backend.
🤖 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 `@CHANGELOG.rst` at line 63, Update the changelog entry for the
simplify_backend option to qualify the backend equivalence claim: state that
both onnxslim and onnxsim aim to preserve model semantics, while allowing their
generated graphs to differ. Keep the existing default backend and CLI/API option
details unchanged.
4 participants
What does this PR do?
Type of change: New feature
Adds
onnxsimas an alternative ONNX simplification backend formodelopt.onnx.quantization.quantize(..., simplify=True)and the--simplifyCLI flag, while keeping
onnxslimas the default. Previously the simplifierwas hard-wired to
onnxslim; this makes the two interchangeable via a newsimplify_backendargument (--simplify_backendon the CLI) that accepts"onnxslim"(default) or"onnxsim".Details:
quantize.py:quantize()and_preprocess_onnx()gain asimplify_backend: str = "onnxslim"argument. The simplify block dispatcheson it —
onnxslim.slim(..., skip_fusion_patterns=["FusionGemm"])foronnxslim,onnxsim.simplify(...)foronnxsim. An unknown backend name ora missing
onnxsiminstall fails loudly (raised before the guarded block),whereas genuine simplification failures still fall back gracefully to the
original model, exactly as before.
onnxsimis imported lazily so it is onlyrequired when actually selected.
__main__.py: adds--simplify_backendwithchoices=["onnxslim", "onnxsim"],default
onnxslim, wired into thequantize(...)call.pyproject.toml/uv.lock: addsonnxsim>=0.7.0to theonnxextraalongside the existing
onnxslim>=0.1.76.onnxsim0.7.0 ships wheels forPython 3.12+ and aarch64, so no platform markers are needed.
Usage
Testing
_preprocess_onnxmock intests/unit/onnx/quantization/test_quantize_api.pyfor the new argument;the existing unit suite for the quantize API continues to pass.
tests/gpu/onnx/test_simplify.pyexercises thedefault (
onnxslim) simplify path end-to-end.pyproject.tomlanduv.lockparse, the lock resolvesonnxsim0.7.0 (depsonnx,rich) withonnxslimuntouched, and thechanged Python files pass
ruff format --check.onnxsimpath(needs a GPU and
onnxsiminstalled).Before your PR is "Ready for review"
simplifydefaults to the existingonnxslimbackend; no existing call site or CLI invocation changes behavior.you follow guidance in
CONTRIBUTING.md: ✅ — addedonnxsimto the optionalonnxextra and relockeduv.lock; no code copied._preprocess_onnxmock only; a dedicatedsimplify_backend-selection unittest is not yet added (happy to add one covering the dispatch and the
invalid-backend
ValueErrorif reviewers want it).section.
/claude reviewonceopened.
Additional Information
No related issue. Default simplification remains
onnxslim;onnxsimis purelyadditive and opt-in.
Summary by CodeRabbit
--simplify_backendCLI option and corresponding parameter to choose betweenonnxslim(default) andonnxsimwhen simplification is enabled.onnxsim>=0.7.0, including compatible Python and aarch64 environments.