[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>
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