[OMNIML-5613] Quantize ResNet residual adds in torch ONNX example - #2024
[OMNIML-5613] Quantize ResNet residual adds in torch ONNX example#2024ajrasane wants to merge 11 commits into
Conversation
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR adds timm ResNet quantization plugins and PTQ recipes, integrates recipe selection into the Torch ONNX example, extends AutoQuantize matching and cost handling, and updates ONNX FP8/INT8 FP16 transformations with structural validation. Changestimm ResNet quantization and export
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2024 +/- ##
===========================================
+ Coverage 60.23% 77.20% +16.96%
===========================================
Files 519 520 +1
Lines 59399 59556 +157
===========================================
+ Hits 35778 45979 +10201
+ Misses 23621 13577 -10044
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Small, focused change (+90/-2) that adds per-block residual-add quantizers to the timm ResNet path of the torch→ONNX example, with CHANGELOG + README updates and an ONNX-level assertion in the existing example test. No licensing concerns; no prompt-injection content in the PR metadata. A few things worth addressing before merge:
-
Duplicated calibration logic / extra data pass.
_add_resnet_residual_quantizersre-implements the enable_calib → forward-loop →load_calib_amaxdance that_calibrate_uncalibrated_quantizers(same file) already performs, and adds a second full pass over the calibration set for ResNet (the first happens insidequantize_modelfor the FP8 Conv overrides). If the residual quantizer were attached asinput_quantizeron the activation module and created beforequantize_model, both the calibration and the dead-quantizer guard would come for free from the existing helpers. Even more idiomatic: modelopt already supports this viaQuantModuleRegistry.register({nn.ReLU: "nn.ReLU"})(QuantInputBase)(exactly whatmodelopt/torch/quantization/nn/modules/quant_activations.pydoes fornn.LeakyReLU) plus a{"parent_class": "nn.ReLU", "quantizer_name": "*input_quantizer"}config entry, which gets calibrated bymtq.quantize's forward loop and is visible tomtq.print_quant_summary/ modelopt state. Please either reuse one of these paths or note in the PR/comment why the manual hook is needed. -
Residual quantizers bypass the file's own
amax<=0/NaN guard, andload_calib_amax()is strict._disable_dead_quantizersonly inspectsinput_quantizer/output_quantizer/weight_quantizer, and it runs insidequantize_model— i.e. before these quantizers exist. A residual quantizer that calibrates toamax == 0(or NaN) will therefore reach the FP8 exporter, which is precisely thescale = 448 / amaxdivision the guard exists to prevent. Also, unlike_calibrate_uncalibrated_quantizers, this code callsload_calib_amax()withoutstrict=False, so any block that didn't see data raises. -
automode picks the residual format from the search space, not the search result. Aftermtq.auto_quantize, each block's actual format is known (e.g.block.conv3.input_quantizer._num_bits); derivingnum_bitsfrom the union of requested formats can give an FP8 residual Q/DQ next to an INT8-quantized block (or the reverse), which is what the rest of this file goes to some length to avoid for TRT.
Minor: the two new functions are the only helpers in this file without docstrings; and assert len(residual_adds) == 16 asserts that every Add in the exported graph is a residual add, which will break confusingly if the exporter ever emits an unrelated Add — consider filtering to the 16 residual adds (e.g. by producer/consumer pattern) before the count assertion.
![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
🤖 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 `@examples/torch_onnx/torch_quant_to_onnx.py`:
- Around line 223-239: Update _add_resnet_residual_quantizers and the
surrounding auto-quantization flow so residual quantizers are installed and
configured before mtq.auto_quantize() runs. For auto mode, include these
residual quantizer modules in every candidate format configuration used by the
search, ensuring their forced INT8/FP8 precision is scored and counted toward
the effective-bits constraint while preserving the existing non-auto behavior.
🪄 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: 9cd948d5-f05e-4d18-9714-517039e16f2d
📒 Files selected for processing (4)
CHANGELOG.rstexamples/torch_onnx/README.mdexamples/torch_onnx/torch_quant_to_onnx.pytests/examples/torch_onnx/test_torch_quant_to_onnx.py
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/torch_onnx/torch_quant_to_onnx.py (1)
636-642:⚠️ Potential issue | 🟠 MajorInstall residual quantizers before the quantization/search pass.
At Line 636, residual quantizers are added only after
quantized_modelhas been created. Inautomode, they are therefore absent from candidate scoring and effective-bits constraints; the later heuristic can also choose a format different from the per-block format selected by AutoQuantize. The standard path additionally requires a second calibration pass and runs dead-quantizer cleanup before these modules exist.Move installation/configuration before quantization, or explicitly integrate these quantizers into AutoQuantize and rerun cleanup after calibration.
🤖 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 `@examples/torch_onnx/torch_quant_to_onnx.py` around lines 636 - 642, Move the `_add_resnet_residual_quantizers` installation and configuration before the quantization/search pass creates `quantized_model`, so residual quantizers participate in AutoQuantize candidate scoring and effective-bits constraints. Ensure calibration and dead-quantizer cleanup operate on these modules, and remove the current post-quantization-only installation path.
🤖 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.
Outside diff comments:
In `@examples/torch_onnx/torch_quant_to_onnx.py`:
- Around line 636-642: Move the `_add_resnet_residual_quantizers` installation
and configuration before the quantization/search pass creates `quantized_model`,
so residual quantizers participate in AutoQuantize candidate scoring and
effective-bits constraints. Ensure calibration and dead-quantizer cleanup
operate on these modules, and remove the current post-quantization-only
installation path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0ba0d581-0d68-4e4d-993c-47c5819b72dd
📒 Files selected for processing (4)
CHANGELOG.rstexamples/torch_onnx/README.mdexamples/torch_onnx/torch_quant_to_onnx.pytests/examples/torch_onnx/test_torch_quant_to_onnx.py
🚧 Files skipped from review as they are similar to previous changes (3)
- examples/torch_onnx/README.md
- CHANGELOG.rst
- tests/examples/torch_onnx/test_torch_quant_to_onnx.py
Co-Authored-By: Codex <noreply@openai.com> Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
|
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: 4
🧹 Nitpick comments (5)
tests/unit/onnx/test_fold_casts.py (1)
86-97: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRedundant dtype plumbing in
_initializer_q_model.
input_dtypeandoutput_dtypeare computed identically, andoutputs[0]is built withFLOAT16at Line 86 only to be overwritten at Line 97. Collapse into a single variable used at construction time.♻️ Proposed simplification
- outputs = [helper.make_tensor_value_info("y", TensorProto.FLOAT16, [None, 4])] + float_dtype = TensorProto.FLOAT16 if opset >= 19 else TensorProto.FLOAT + outputs = [helper.make_tensor_value_info("y", float_dtype, [None, 4])] if shared: nodes.append(helper.make_node("Identity", ["w"], ["w_out"], "identity")) outputs.append(helper.make_tensor_value_info("w_out", TensorProto.FLOAT, [4, 4])) @@ - input_dtype = TensorProto.FLOAT16 if opset >= 19 else TensorProto.FLOAT - output_dtype = TensorProto.FLOAT16 if opset >= 19 else TensorProto.FLOAT - outputs[0].type.tensor_type.elem_type = output_dtype return helper.make_model( helper.make_graph( nodes, "g", - [helper.make_tensor_value_info("x", input_dtype, [None, 4])], + [helper.make_tensor_value_info("x", float_dtype, [None, 4])],🤖 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/test_fold_casts.py` around lines 86 - 97, In `_initializer_q_model`, replace the redundant `input_dtype` and `output_dtype` calculations with one shared dtype variable derived from `opset`, and use it when constructing the primary `y` output instead of creating it as `FLOAT16` and mutating it afterward. Remove the subsequent `outputs[0].type.tensor_type.elem_type` assignment while preserving the existing opset-dependent dtype.examples/torch_onnx/torch_quant_to_onnx.py (1)
232-259: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low value
_calibrate_new_quantizersmutates_if_quantstate without honoring pre-existing disabled-quant quantizers.
enabled_quantizersis filtered byis_enabled(i.e._disabled), not by_if_quant. Any quantizer that was enabled but intentionally haddisable_quant()applied earlier getsenable_quant()in thefinallyblock, silently turning quantization back on. Filtering onmodule._if_quantinstead would make the save/restore symmetric.♻️ Suggested tweak
- enabled_quantizers = [ - module - for module in model.modules() - if isinstance(module, TensorQuantizer) and module.is_enabled - ] + enabled_quantizers = [ + module + for module in model.modules() + if isinstance(module, TensorQuantizer) and module.is_enabled and module._if_quant + ]🤖 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 `@examples/torch_onnx/torch_quant_to_onnx.py` around lines 232 - 259, Update _calibrate_new_quantizers to track quantizers based on their pre-existing _if_quant state rather than is_enabled when building enabled_quantizers. Restore quantization only for quantizers whose _if_quant state was originally active, preserving intentionally disabled quantizers through the finally block.tests/examples/torch_onnx/test_torch_quant_to_onnx.py (2)
35-62: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUnguarded
producers[...]lookups will raiseKeyErrorinstead of a useful failure.If an
Addinput is a graph input or an initializer (no producer node), Line 54/56 raisesKeyErrorrather than an assertion explaining what the graph looks like. Same for thenext(...)lookups at Lines 88 and 95, which raiseStopIterationif the exporter ever emitsMatMul/ReduceMeaninstead ofGemm/GlobalAveragePool. Preferproducers.get(...)plus explicit asserts so failures are diagnosable.🤖 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/examples/torch_onnx/test_torch_quant_to_onnx.py` around lines 35 - 62, Harden _assert_residual_adds_are_quantized against missing graph producers by replacing direct producers[...] accesses with producers.get(...) and explicit assertions that identify the missing input or node. Apply the same pattern to the MatMul/ReduceMean lookup paths around the next(...) calls, asserting the expected producer exists before dereferencing it while preserving the current validation behavior.
64-116: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHard-coded node counts (54/52, 53, 12/4/4) need a comment explaining their derivation.
These numbers encode the exact ResNet-50 Q/DQ topology, but nothing in the test says where they come from, so a future exporter change produces an unexplainable
assert 53 == 52. A one-line comment per assertion (e.g. "53 = 53 Conv weights, fc weight quantizer disabled") would make the failures actionable. As per path instructions, checked-in tests should "document expected behavior".🤖 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/examples/torch_onnx/test_torch_quant_to_onnx.py` around lines 64 - 116, Add concise comments immediately before each hard-coded topology-count assertion in the quantization test, explaining how the expected values derive from the ResNet-50 Q/DQ structure and mode-specific behavior. Cover the activation quantizer counts, DQ fanout counts, and int8 weight quantizer count, including details such as the number of convolution weights and disabled fully connected weight quantization; leave the assertions unchanged.Source: Path instructions
modelopt/torch/_deploy/utils/torch_onnx.py (1)
48-48: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueImporting a private helper across modules.
_convert_q_data_initializers_to_fp16is underscore-private tomodelopt/onnx/utils.pybut is consumed here. Since it's now part of the export pipeline contract, consider promoting it to a public name (and adding it to that module's__all__) so the dependency is explicit.Note the sequencing is load-bearing: this call raises if any Q scale is still FP32 after Line 670, which is why the stricter fold guard in
modelopt/onnx/utils.pymatters (flagged there).Also applies to: 669-673
🤖 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 `@modelopt/torch/_deploy/utils/torch_onnx.py` at line 48, Promote _convert_q_data_initializers_to_fp16 in modelopt/onnx/utils.py to a public helper name, add that name to the module’s __all__, and update the import and call sites in the ONNX export flow around _convert_q_data_initializers_to_fp16 accordingly. Preserve the existing sequencing and validation behavior after the fold guard.Source: Coding guidelines
🤖 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 `@examples/torch_onnx/torch_quant_to_onnx.py`:
- Around line 363-392: Update _finalize_resnet_quantizers so the int8/fp8 path
only accesses quantizer attributes when they exist, using safe getattr checks
for block.conv1, model.fc, model.global_pool, and related quantizers. Preserve
enabling/disabling behavior for present quantizers, and avoid AttributeError
when _prepare_resnet_quantizers skipped setup or a block variant lacks these
attributes.
In `@modelopt/onnx/export/fp8_exporter.py`:
- Line 34: Update the weight-channel threshold logic near the FP8
minimum-channel constant and the relevant convolution export check to account
for the convolution’s group count, matching the torch-side gate. Base the
decision on the effective total input channels rather than per-group
weight_input.values.shape[1], so grouped and depthwise convolutions consistently
enable or skip activation and weight Q/DQ paths.
In `@modelopt/onnx/utils.py`:
- Around line 1471-1481: The Cast-folding condition in the loop over
onnx_model.graph.node should not reject Cast-to-FLOAT nodes solely because
tensor_types lacks node.input[0]. Treat missing input type metadata as eligible
for folding, or resolve the producer output type before skipping; preserve the
existing FLOAT16 exclusion when the type is explicitly known.
- Around line 1524-1534: Replace the hard ValueError in the Q-consumer
validation with graceful skipping of that initializer, and emit a warning
identifying q_node.name and the observed scale dtype. Ensure
get_onnx_bytes_and_metadata continues exporting when a Q scale is a graph input
or remains non-FP16, while preserving FP16 conversion for valid Q scales.
---
Nitpick comments:
In `@examples/torch_onnx/torch_quant_to_onnx.py`:
- Around line 232-259: Update _calibrate_new_quantizers to track quantizers
based on their pre-existing _if_quant state rather than is_enabled when building
enabled_quantizers. Restore quantization only for quantizers whose _if_quant
state was originally active, preserving intentionally disabled quantizers
through the finally block.
In `@modelopt/torch/_deploy/utils/torch_onnx.py`:
- Line 48: Promote _convert_q_data_initializers_to_fp16 in
modelopt/onnx/utils.py to a public helper name, add that name to the module’s
__all__, and update the import and call sites in the ONNX export flow around
_convert_q_data_initializers_to_fp16 accordingly. Preserve the existing
sequencing and validation behavior after the fold guard.
In `@tests/examples/torch_onnx/test_torch_quant_to_onnx.py`:
- Around line 35-62: Harden _assert_residual_adds_are_quantized against missing
graph producers by replacing direct producers[...] accesses with
producers.get(...) and explicit assertions that identify the missing input or
node. Apply the same pattern to the MatMul/ReduceMean lookup paths around the
next(...) calls, asserting the expected producer exists before dereferencing it
while preserving the current validation behavior.
- Around line 64-116: Add concise comments immediately before each hard-coded
topology-count assertion in the quantization test, explaining how the expected
values derive from the ResNet-50 Q/DQ structure and mode-specific behavior.
Cover the activation quantizer counts, DQ fanout counts, and int8 weight
quantizer count, including details such as the number of convolution weights and
disabled fully connected weight quantization; leave the assertions unchanged.
In `@tests/unit/onnx/test_fold_casts.py`:
- Around line 86-97: In `_initializer_q_model`, replace the redundant
`input_dtype` and `output_dtype` calculations with one shared dtype variable
derived from `opset`, and use it when constructing the primary `y` output
instead of creating it as `FLOAT16` and mutating it afterward. Remove the
subsequent `outputs[0].type.tensor_type.elem_type` assignment while preserving
the existing opset-dependent dtype.
🪄 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: b4d1c1cd-df7c-44b1-be7d-a25ddb7f9ddb
📒 Files selected for processing (8)
CHANGELOG.rstexamples/torch_onnx/torch_quant_to_onnx.pymodelopt/onnx/export/fp8_exporter.pymodelopt/onnx/utils.pymodelopt/torch/_deploy/utils/torch_onnx.pytests/examples/torch_onnx/test_torch_quant_to_onnx.pytests/unit/onnx/quantization/test_fp8_mha_exporter.pytests/unit/onnx/test_fold_casts.py
🚧 Files skipped from review as they are similar to previous changes (1)
- CHANGELOG.rst
| # when using 1/448 as the Q scale (single fixed value — softmax range is data-independent). | ||
| _FP8_E4M3_MAX = 448.0 | ||
| _FP8_E4M3_SOFTMAX_SCALE = 1.0 / _FP8_E4M3_MAX | ||
| _FP8_MIN_CONV_CHANNELS = 16 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Are there grouped/depthwise convs in the FP8 export path's expected model set?
rg -nP -C3 '\bgroup\b|groups' modelopt/onnx/export/fp8_exporter.pyRepository: NVIDIA/Model-Optimizer
Length of output: 160
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant files and inspect the exporter + torch-side gate.
git ls-files 'modelopt/onnx/export/fp8_exporter.py' 'examples/torch_onnx/torch_quant_to_onnx.py'
wc -l modelopt/onnx/export/fp8_exporter.py examples/torch_onnx/torch_quant_to_onnx.py
ast-grep outline modelopt/onnx/export/fp8_exporter.py --view expanded
ast-grep outline examples/torch_onnx/torch_quant_to_onnx.py --view expandedRepository: NVIDIA/Model-Optimizer
Length of output: 2528
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the relevant implementation slices.
sed -n '150,230p' modelopt/onnx/export/fp8_exporter.py
printf '\n----\n'
sed -n '190,225p' examples/torch_onnx/torch_quant_to_onnx.pyRepository: NVIDIA/Model-Optimizer
Length of output: 4975
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Find every place the FP8 conv channel threshold is applied.
rg -n -C3 '_FP8_MIN_CONV_CHANNELS|shape\[:2\]|in_channels|out_channels|groups' modelopt/onnx/export/fp8_exporter.pyRepository: NVIDIA/Model-Optimizer
Length of output: 686
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the surrounding exporter logic to see how the skipped weight path is used.
sed -n '1,120p' modelopt/onnx/export/fp8_exporter.py
printf '\n----\n'
sed -n '120,190p' modelopt/onnx/export/fp8_exporter.py
printf '\n----\n'
sed -n '440,500p' modelopt/onnx/export/fp8_exporter.pyRepository: NVIDIA/Model-Optimizer
Length of output: 11674
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check whether grouped/depthwise convolutions are represented in tests/examples around FP8 export.
rg -n -C2 'depthwise|groups *= *[1-9]|grouped conv|grouped convolution|Conv2d\(.*groups' modelopt examples testsRepository: NVIDIA/Model-Optimizer
Length of output: 22860
Grouped convs need a group-aware channel check modelopt/onnx/export/fp8_exporter.py:193-196 — weight_input.values.shape[1] is per-group, so grouped/depthwise convs can be treated as “small channel” even when the module’s in_channels/out_channels are large. That can leave activation Q/DQ enabled while the weight DQ is skipped. Use group in the threshold check so the weight path matches the torch-side gate.
🤖 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 `@modelopt/onnx/export/fp8_exporter.py` at line 34, Update the weight-channel
threshold logic near the FP8 minimum-channel constant and the relevant
convolution export check to account for the convolution’s group count, matching
the torch-side gate. Base the decision on the effective total input channels
rather than per-group weight_input.values.shape[1], so grouped and depthwise
convolutions consistently enable or skip activation and weight Q/DQ paths.
Co-Authored-By: Codex <noreply@openai.com> Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
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 `@modelopt/onnx/export/fp8_exporter.py`:
- Around line 196-197: Update the input-stem skip around the weight
dequantization pass so it applies only when the export is identified as a
ResNet, rather than using a three-channel input as the model-family
discriminator. Propagate and check an explicit ResNet/export marker alongside
node.inputs[0].name and weight_input.values.shape, while preserving the skip for
ResNet input stems and normal FP8 DQ restoration for non-ResNet models.
🪄 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: 22b23dca-030c-4c44-891d-67056a837b68
📒 Files selected for processing (4)
examples/torch_onnx/torch_quant_to_onnx.pymodelopt/onnx/export/fp8_exporter.pytests/examples/torch_onnx/test_torch_quant_to_onnx.pytests/unit/onnx/quantization/test_fp8_mha_exporter.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/examples/torch_onnx/test_torch_quant_to_onnx.py
| if node.inputs[0].name in graph_inputs and weight_input.values.shape[1] == 3: | ||
| continue |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Keep the input-stem skip ResNet-specific.
This skips FP8 weight DQ for any direct RGB Conv, while the Torch-side exclusion only applies to timm.models.resnet.ResNet. A non-ResNet model with a direct three-channel input can therefore retain FP16 weights even though this exporter normally restores Conv weight FP8 DQ. Propagate an explicit ResNet/export marker into this pass, or otherwise avoid using RGB topology as the model-family discriminator.
🤖 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 `@modelopt/onnx/export/fp8_exporter.py` around lines 196 - 197, Update the
input-stem skip around the weight dequantization pass so it applies only when
the export is identified as a ResNet, rather than using a three-channel input as
the model-family discriminator. Propagate and check an explicit ResNet/export
marker alongside node.inputs[0].name and weight_input.values.shape, while
preserving the skip for ResNet input stems and normal FP8 DQ restoration for
non-ResNet models.
Co-Authored-By: Codex <noreply@openai.com> Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
Co-Authored-By: Codex <noreply@openai.com> Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
Co-Authored-By: Codex <noreply@openai.com> Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
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)
modelopt/torch/quantization/plugins/timm.py (1)
16-25: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd
__all__for this module's public API.Only
is_resnet_quantization_supportedis a non-underscore (public) symbol here; everything else is intentionally private. Declaring__all__ = ["is_resnet_quantization_supported"]makes that contract explicit and keepsfrom .timm import *inplugins/__init__.pypredictable if more public helpers are added later.As per coding guidelines, "Define each module's public API with
__all__ = [...]."🤖 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 `@modelopt/torch/quantization/plugins/timm.py` around lines 16 - 25, Add a module-level __all__ declaration in timm.py containing only is_resnet_quantization_supported, preserving the intended public API for wildcard imports.Source: Coding guidelines
🤖 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_recipes/README.md`:
- Around line 79-84: Update the “Choosing where to look” guidance in README.md
to add a timm-specific lookup step before the general fallback, directing
readers to timm/<architecture>/ for architecture-specific recipes. Preserve the
existing huggingface and general guidance while ensuring timm recipes are
included in the selection flow.
---
Nitpick comments:
In `@modelopt/torch/quantization/plugins/timm.py`:
- Around line 16-25: Add a module-level __all__ declaration in timm.py
containing only is_resnet_quantization_supported, preserving the intended public
API for wildcard imports.
🪄 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: 62647b7a-3e4f-40ce-bedb-56b81fdb85f9
📒 Files selected for processing (22)
CHANGELOG.rstexamples/torch_onnx/README.mdexamples/torch_onnx/torch_quant_to_onnx.pymodelopt/torch/opt/dynamic.pymodelopt/torch/quantization/algorithms.pymodelopt/torch/quantization/conversion.pymodelopt/torch/quantization/plugins/__init__.pymodelopt/torch/quantization/plugins/timm.pymodelopt_recipes/README.mdmodelopt_recipes/timm/resnet/ptq/README.mdmodelopt_recipes/timm/resnet/ptq/fp8.yamlmodelopt_recipes/timm/resnet/ptq/int8.yamlmodelopt_recipes/timm/resnet/ptq/mxfp8.yamlmodelopt_recipes/timm/resnet/ptq/nvfp4.yamlmodelopt_recipes/timm/resnet/ptq/nvfp4_awq_lite.yamlmodelopt_recipes/timm/resnet/ptq/static_fp8.quant_cfg.yamlmodelopt_recipes/timm/resnet/ptq/static_int8.quant_cfg.yamltests/unit/torch/nas/test_registry.pytests/unit/torch/quantization/plugins/test_timm.pytests/unit/torch/quantization/test_autoquant.pytests/unit/torch/quantization/test_config_validation.pytests/unit/torch/quantization/test_quantize_cpu.py
🚧 Files skipped from review as they are similar to previous changes (1)
- examples/torch_onnx/README.md
| ## `timm/` — architecture-specific recipes | ||
|
|
||
| Recipes under `timm/<architecture>/` capture quantization choices required by | ||
| vision architectures and their deployment backends. See | ||
| [`timm/resnet/ptq/`](timm/resnet/ptq/) for ResNet recipes. | ||
|
|
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Include timm recipes in the selection guidance.
The new section is discoverable, but the preceding “Choosing where to look” guidance still directs readers from huggingface/ directly to general/, skipping timm/<architecture>/. Add a timm-specific lookup step before the general fallback.
🤖 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 `@modelopt_recipes/README.md` around lines 79 - 84, Update the “Choosing where
to look” guidance in README.md to add a timm-specific lookup step before the
general fallback, directing readers to timm/<architecture>/ for
architecture-specific recipes. Preserve the existing huggingface and general
guidance while ensuring timm recipes are included in the selection flow.
Keep the change focused on shortcut QDQ placement with FP8, INT8, and AutoQuantize recipes. Co-Authored-By: Codex <noreply@openai.com> Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
Resolve the torch ONNX example conflicts while preserving the residual-only recipe scope. Co-Authored-By: Codex <noreply@openai.com> Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review of the reworked (recipe-driven) implementation. The approach changed substantially since the last round, and most of the earlier findings are genuinely resolved:
Resolved
- Duplicated calibration / second data pass (critical): the manual
enable_calib→ forward-loop →load_calib_amaxblock is gone._add_resnet_residual_quantizersnow runs as aCUSTOM_POST_CONVERSION_PLUGINScallback, i.e. after_replace_quant_modulebut beforeset_quantizer_by_cfg(seemodelopt/torch/quantization/conversion.py::replace_quant_module), so the recipe's'*residual_quantizer'entry configures it andmtq.quantize's own forward loop calibrates it. The quantizers are also visible toprint_quant_summaryand modelopt state. Good fix. - Strict
load_calib_amax()(critical): moot — standard calibration path is used now. - Auto-mode residual format inferred from the search space (critical): replaced by an explicit design — the residual quantizer is pinned to FP8 via
fixed_quantization_configin the AutoQuantize recipe, documented in the README and PR body. Acceptable as an intentional choice (see one follow-up below). - Test asserted "every
Addis residual" (minor): the test now selects residual adds structurally by their singleReluconsumer before the count assertion.
Still open / new
_disable_dead_quantizersstill only inspectsinput_quantizer/output_quantizer/weight_quantizer, so aresidual_quantizerthat calibrates toamax == 0/NaN reaches the FP8 exporter'sscale = 448 / amax— the exact case that guard was added for. The explicit NaN/non-positive check you described in the previous round no longer exists anywhere in the new code, so this half of the earlier comment regressed (one-line fix).tests/examples/torch_onnx/test_torch_quant_to_onnx.pyruns resnet50 +autowith--trt_build, but the PR body says the AutoQuantize FP8/INT8 mix only builds on Blackwell+. Please confirm how this test is expected to pass on the CI GPUs.- The recipe →
mtq.auto_quantizekwargs translation duplicatesexamples/hf_ptq/hf_ptq.py::_mtq_inputs_from_auto_quantize_config, and the copy silently dropsscore_size(the new recipe setsscore_size: 128but the example uses--num_score_steps),cost_excluded_layers, and top-levelcandidate_formats. --auto_quantization_formatsis silently overridden for ResNet inautomode with no warning and no opt-out (--recipeis rejected with--quantize_mode=auto).- Nits: the two new helpers are the only ones in this module without docstrings (the rationale docstring added in the previous round didn't survive the rewrite);
"residual_quantizer" in block.downsample._modulescould usehasattr;modelopt_recipes/README.mdstill doesn't mention the new top-leveltimm/tree (earlier bot comment unaddressed);timm/resnet/ptq/{fp8,int8}.yamlare verbatim copies ofconfigs/ptq/presets/model/{fp8,int8}.yamlplus one entry, so they will drift if the presets change.
No licensing concerns (new YAMLs carry the standard NVIDIA SPDX header); no prompt-injection content in the PR metadata.
| continue | ||
| residual_quantizer = TensorQuantizer() | ||
| residual_quantizer.disable() | ||
| block.downsample.add_module("residual_quantizer", residual_quantizer) |
There was a problem hiding this comment.
Bot comment.
The amax <= 0/NaN guard still doesn't cover these quantizers: _disable_dead_quantizers only looks at getattr(mod, attr) for attr in ("input_quantizer", "output_quantizer", "weight_quantizer"), and the new attribute is downsample.residual_quantizer. The previous round's explicit "disable shortcut quantizers with missing/NaN/non-positive amax" code was dropped in the rewrite, so a residual quantizer that calibrates to amax == 0 now reaches export_fp8's scale = 448 / amax — exactly the failure that helper exists to prevent (unlikely on pretrained weights, but the guard was added for random-init/zero-init cases and this test suite runs --no_pretrained).
One-line fix: add "residual_quantizer" to the attribute tuple in _disable_dead_quantizers (or have it iterate TensorQuantizer modules directly, which would make it name-agnostic).
Also, nit carried over from the last round: this and _get_resnet_recipe are now the only helpers in this module without docstrings. Worth restating why the quantizer goes at the end of downsample (so it lands immediately before the Add), why downsample=None is converted to an empty Sequential, and why only fp8/int8/auto get residual quantizers.
| constraints["effective_bits"] = effective_bits | ||
| format_configs = [] | ||
| fixed_quantization_config = recipe.quantize.model_dump() | ||
| module_search_spaces = [ |
There was a problem hiding this comment.
Bot comment.
This recipe → mtq.auto_quantize mapping duplicates examples/hf_ptq/hf_ptq.py::_mtq_inputs_from_auto_quantize_config, and the copy is lossy:
auto_config.score_sizeis ignored — the new recipe declaresscore_size: 128but the example passes--num_score_stepsinstead, so that recipe field is dead config (it happens to match the CLI default today, which hides the bug).auto_config.cost_excluded_layersis ignored (hf_ptq maps it intoconstraints["cost"]["excluded_module_name_patterns"]).auto_config.candidate_formats(top-level, nomodule_search_spaces) is ignored: for such a recipe this buildsquantization_formats=[]+module_search_spaces=[], andmtq.auto_quantizethen raises"fixed_quantization_config requires at least one explicit module_search_spaces entry". That's brittle for anything added to_RESNET_RECIPESlater.
Since this mapping is now needed by two examples, consider promoting it into modelopt/recipe (importable) and calling it from both, rather than maintaining a second partial copy.
| ) | ||
| print(f"Base Model - Top-1 Accuracy: {top1:.2f}%, Top-5 Accuracy: {top5:.2f}%") | ||
|
|
||
| resnet_recipe = None if args.recipe else _get_resnet_recipe(model, args.quantize_mode) |
There was a problem hiding this comment.
Bot comment.
For ResNet + --quantize_mode=auto, args.auto_quantization_formats is now silently ignored (the recipe's module_search_spaces replaces it), and since --recipe is rejected with --quantize_mode=auto there is no way to opt out of the ResNet recipe. Please either warn when a non-default --auto_quantization_formats is discarded, or allow --recipe (incl. an explicit "none") in auto mode. Same applies to the fp8/int8 paths: get_quant_config's overrides are bypassed for ResNet without any log line saying which recipe was chosen (load_recipe does print the path, so this is a smaller concern there).
| cmd_parts.extend(["--no_pretrained", "--trt_build"]) | ||
| run_example_command(cmd_parts, "torch_onnx") | ||
|
|
||
| if model_key == "resnet50" and quantize_mode in _RESNET_QUANT_MODES: |
There was a problem hiding this comment.
Bot comment.
This test runs every (model, mode) combination with --trt_build, including resnet50 + auto, which after this PR searches FP8 vs INT8 per block with the residual pinned to FP8. The PR body states that this INT8/FP8 mix is only supported by TensorRT on Blackwell and newer, and the listed validation is "TensorRT engine builds passed for FP8 and INT8 on Ada" — i.e. the auto engine build appears untested there. Note also that both candidates cost the same 8 effective bits under the effective_bits: 8.0 constraint, so the solver is free to mix formats arbitrarily.
How is this combination expected to pass on the CI GPUs? If they're pre-Blackwell, either skip --trt_build for resnet50+auto (with a comment explaining why) or constrain the recipe so a single format is selected.
2 participants
What does this PR do?
Type of change: Bug fix
Adds recipe-backed residual quantization for timm ResNet models in the torch ONNX example:
Add.Add.Usage
python examples/torch_onnx/torch_quant_to_onnx.py \ --timm_model_name=resnet50 \ --quantize_mode=<fp8|int8|auto> \ --onnx_save_path=resnet50.onnxTesting
Before your PR is "Ready for review"
CONTRIBUTING.md: N/A