Skip to content

Add onnxsim as an alternative ONNX simplification backend #2018

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
take-cheeze wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add onnxsim as an alternative ONNX simplification backend #2018

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ Experimental
- 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.

@coderabbitai coderabbitai Bot Aug 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 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 -250

Repository: 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.rst

Repository: 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.


**Bug Fixes**

Expand Down
11 changes: 11 additions & 0 deletions modelopt/onnx/quantization/__main__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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,
Expand Down
34 changes: 31 additions & 3 deletions modelopt/onnx/quantization/quantize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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].
Expand Down Expand Up @@ -627,6 +654,7 @@ def quantize(
trt_plugins_precision,
override_shapes, # type: ignore[arg-type]
simplify,
simplify_backend,
quantize_mode,
opset,
)
Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
]
Expand Down
1 change: 1 addition & 0 deletions tests/unit/onnx/quantization/test_quantize_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@ def fake_preprocess(
trt_plugins_precision,
override_shapes,
simplify,
simplify_backend,
quantize_mode,
opset,
):
Expand Down
29 changes: 29 additions & 0 deletions uv.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.