[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 4 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. |
📝 WalkthroughWalkthroughThe Torch ONNX ResNet example now quantizes residual shortcut inputs before addition, calibrates those quantizers, validates exported ONNX graphs, and documents the behavior. ChangesResNet residual quantization
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Main
participant ResNetBlocks
participant DataLoader
participant TensorQuantizer
participant ONNXTest
Main->>ResNetBlocks: Install residual quantizers
Main->>DataLoader: Run calibration batches
DataLoader->>TensorQuantizer: Collect activation ranges
Main->>TensorQuantizer: Load amax and enable quantization
Main->>ONNXTest: Export quantized model
ONNXTest->>ONNXTest: Verify Add inputs originate from quantized paths
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 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 66.83% 66.65% -0.18%
==========================================
Files 519 519
Lines 58916 58916
==========================================
- Hits 39376 39272 -104
- Misses 19540 19644 +104
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.
| if not residual_quantizers: | ||
| return | ||
|
|
||
| for quantizer in residual_quantizers: |
There was a problem hiding this comment.
Bot comment.
This block duplicates _calibrate_uncalibrated_quantizers (a few dozen lines below) and costs a second full pass over the calibration set for ResNet, since quantize_model already ran one for the FP8 Conv overrides.
Two cheaper options:
- Attach the quantizer as
input_quantizeron the activation module and create it beforequantize_model— then_calibrate_uncalibrated_quantizersand_disable_dead_quantizersboth pick it up automatically, and this whole calibration block goes away. - Or use the framework path:
QuantModuleRegistry.register({nn.ReLU: "nn.ReLU"})(QuantInputBase)(seemodelopt/torch/quantization/nn/modules/quant_activations.py, which does exactly this fornn.LeakyReLU) plus a{"parent_class": "nn.ReLU", "quantizer_name": "*input_quantizer", "cfg": {...}}entry restricted to*act3*/*act2*, somtq.quantize's forward loop calibrates it andprint_quant_summaryreports it.
| model.train(was_training) | ||
|
|
||
| for quantizer in residual_quantizers: | ||
| quantizer.load_calib_amax() |
There was a problem hiding this comment.
Bot comment.
Two robustness gaps relative to the sibling helper in this file:
load_calib_amax()is strict here, whereas_calibrate_uncalibrated_quantizersusesload_calib_amax(strict=False); a block that saw no data will raise instead of being skipped.- These quantizers are created after
_disable_dead_quantizersruns (insidequantize_model), and that helper only looks atinput_quantizer/output_quantizer/weight_quantizer— so a residual quantizer that calibrates toamax == 0or NaN reaches the FP8 exporter and hits thescale = 448 / amaxblow-up the guard was added for.
Suggest either naming the attribute so the existing helpers cover it, or applying the same amax <= 0 / isnan disable check to residual_quantizers after load_calib_amax(strict=False).
| # Dynamic block quantizers do not support the residual path's 4D tensors. | ||
| num_bits = (4, 3) | ||
| elif quantize_mode == "auto": | ||
| activation_formats = set(auto_quantization_formats) - {"INT4_AWQ_CFG"} |
There was a problem hiding this comment.
Bot comment.
In auto mode this infers the residual format from the requested search space rather than what auto_quantize actually selected per block, so a block that ended up INT8 can get an FP8 residual Q/DQ (or vice versa) — the kind of precision mismatch the rest of this file works hard to avoid for TRT. Since this runs after the search, the per-block answer is available, e.g. getattr(block.conv3, "input_quantizer", None)._num_bits (fall back to the current heuristic if the quantizer is absent/disabled).
| consumers[input_name].append(node) | ||
|
|
||
| residual_adds = [node for node in model.graph.node if node.op_type == "Add"] | ||
| assert len(residual_adds) == 16 |
There was a problem hiding this comment.
Bot comment.
residual_adds is every Add node in the graph, so this asserts "ResNet-50 exports exactly 16 Adds and all of them are residual". That's incidentally true today, but it will fail confusingly if the exporter ever emits an unrelated Add (bias/fused pattern, autocast helper, etc.). Prefer identifying the residual adds structurally (e.g. Add whose single consumer is the Relu that follows the block) and then asserting the count of those is 16.
| return (module.residual_quantizer(inputs[0]),) | ||
|
|
||
|
|
||
| def _add_resnet_residual_quantizers(model, quantize_mode, auto_quantization_formats, data_loader): |
There was a problem hiding this comment.
Bot comment.
Nit: every other helper in this module has a docstring explaining the TRT/architecture reason it exists (see _disable_low_channel_conv_input_quantizers). Please add one here covering why only act2/act3 are hooked, why FP8 per-tensor is used instead of MXFP8/NVFP4 blocks for the 4D residual tensor, and why int4_awq is skipped.
![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
2 participants
What does this PR do?
Type of change: Bug fix
Quantizes shortcut inputs in timm ResNet
BasicBlockandBottleneckmodules immediately before residual addition. The new residual quantizers are calibrated with the example's calibration data and support INT8, FP8, MXFP8, NVFP4, and Auto activation-quantized modes.The ONNX example test verifies that all 16 ResNet-50 residual
Addnodes have a Q/DQ input and feed directly intoRelu. The README and changelog document the new behavior.Usage
Testing
Addnodes have Q/DQ immediately before the addition.Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/ASummary by CodeRabbit
Bug Fixes
Tests
Documentation