[OMNIML-5613] Quantize ResNet residual adds in torch ONNX example

CHANGELOG.rst

@@ -81,6 +81,7 @@ Changelog
 
 **Bug Fixes**
 
+- Quantize residual-add outputs in the torch ONNX ResNet example.
 - Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. ``clear_stale_value_info`` now reconciles stale output shapes via symbolic shape inference (keeping every output's shape field populated), and AutoCast runs ONNX shape inference in strict mode and falls back to schema-based standalone type inference when it fails, so unresolved ops no longer leave tensors untyped.
 - Fused MoE expert auto-detection (``register_fused_experts_on_the_fly``) no longer requires an ``act_fn`` attribute. Some fused-expert modules (e.g. ``MiniMaxM3VLExperts``) apply a custom gated activation between the two ``F.linear`` calls instead of exposing ``act_fn``; they were silently skipped, leaving routed experts unquantized (an experts-only recipe matched nothing) and failing HF export with ``NotImplementedError``. ``_QuantFusedExperts`` is activation-agnostic (it only intercepts the two ``F.linear`` calls), so the requirement was unnecessary. This enables NVFP4/FP8 quantization and export for MiniMax-M2 / MiniMax-M3.
 - Fix unified HF export emitting transformers' *in-memory* (post-``conversion_mapping``) tensor names instead of the original model-hub names, breaking the unified-checkpoint contract (observed on MiniMax-M3: exported ``model.language_model.*`` / ``mlp.experts.*.gate_proj`` instead of hub ``language_model.model.*`` / ``block_sparse_moe.experts.*.w{1,2,3}``). transformers' own save-side ``revert_weight_conversion`` is disabled by ModelOpt because it raises ``RuntimeError`` on 0-d scalar scale tensors, so a new quant-aware reverse conversion (``modelopt/torch/export/quant_aware_conversion.py``) derives rename/split rules from the model's conversion mapping via transformers' ``reverse_transform()`` and carries each weight's companion scale tensors (``weight_scale``, ``weight_scale_2``, ``input_scale``, ``weight_scale_inv``, ``bias``) through the renames and un-fusions, so quantized exports round-trip to the hub names. Any mapping op that cannot be reversed quant-aware yet (e.g. still-stacked fused experts) falls back to the previous in-memory names instead of aborting the export.

examples/torch_onnx/README.md

@@ -54,6 +54,7 @@ The `torch_quant_to_onnx.py` script quantizes [timm](https://github.com/huggingf
 - Loads a pretrained timm torch model (default: ViT-Base).
 - Quantizes the torch model to FP8, MXFP8, INT8, NVFP4, or INT4_AWQ using ModelOpt.
 - For models with Conv2d layers (e.g., SwinTransformer), automatically overrides Conv2d quantization to FP8 (for MXFP8/NVFP4 modes) or INT8 (for INT4_AWQ mode) for TensorRT compatibility.
+- Quantizes ResNet residual-add outputs before activation for activation-quantized modes.
 - Exports the quantized model to ONNX.
 - Postprocesses the ONNX model to be compatible with TensorRT.
 - Saves the final ONNX model.

examples/torch_onnx/torch_quant_to_onnx.py

@@ -34,6 +34,8 @@
 from evaluation import evaluate
 
 import modelopt.torch.quantization as mtq
+from modelopt.torch.quantization.config import QuantizerAttributeConfig
+from modelopt.torch.quantization.nn import TensorQuantizer
 
 """
 Quantize a timm vision model and export to ONNX for TensorRT deployment.
@@ -218,6 +220,58 @@ def _disable_low_channel_conv_input_quantizers(model):
                 q.disable()
 
 
+def _quantize_residual_input(module, inputs):
+    return (module.residual_quantizer(inputs[0]),)
+
+
+def _add_resnet_residual_quantizers(model, quantize_mode, auto_quantization_formats, data_loader):
+    if quantize_mode == "int8":
+        num_bits = 8
+    elif quantize_mode in ("fp8", "mxfp8", "nvfp4"):
+        # 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"}
+        if not activation_formats:
+            return
+        num_bits = 8 if activation_formats == {"INT8_DEFAULT_CFG"} else (4, 3)
+    else:
+        return
+
+    residual_quantizers = []
+    block_types = (timm.models.resnet.BasicBlock, timm.models.resnet.Bottleneck)
+    for block in model.modules():
+        if not isinstance(block, block_types):
+            continue
+        activation = block.act3 if isinstance(block, timm.models.resnet.Bottleneck) else block.act2
+        activation.residual_quantizer = TensorQuantizer(
+            QuantizerAttributeConfig(num_bits=num_bits, axis=None)
+        ).to(next(block.parameters()).device)
+        activation.register_forward_pre_hook(_quantize_residual_input)
+        residual_quantizers.append(activation.residual_quantizer)
+
+    if not residual_quantizers:
+        return
+
+    for quantizer in residual_quantizers:
+        quantizer.disable_quant()
+        quantizer.enable_calib()
+
+    was_training = model.training
+    model.eval()
+    try:
+        with torch.no_grad():
+            for batch in data_loader:
+                model(batch["image"] if isinstance(batch, dict) else batch)
+    finally:
+        model.train(was_training)
+
+    for quantizer in residual_quantizers:
+        quantizer.load_calib_amax()
+        quantizer.disable_calib()
+        quantizer.enable_quant()
+
+
 def load_calibration_data(model, data_size, batch_size, device, with_labels=False):
     """Load and prepare calibration data.
 
@@ -580,6 +634,13 @@ def main():
 
         quantized_model = quantize_model(model, config, data_loader)
 
+    _add_resnet_residual_quantizers(
+        quantized_model,
+        args.quantize_mode,
+        args.auto_quantization_formats,
+        data_loader,
+    )
+
     # MXFP8/NVFP4 lower their input quantizers to TRT DynamicQuantize (2D/3D only).
     # Disable quantizers on 4D-input layers (Swin's norm1 / downsample.norm / top-level norm).
     # Auto mode also needs this when an MXFP8/NVFP4 candidate format is in the search set.

tests/examples/torch_onnx/test_torch_quant_to_onnx.py

@@ -14,6 +14,9 @@
 # limitations under the License.
 
 
+from collections import defaultdict
+
+import onnx
 import pytest
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
 
@@ -29,11 +32,30 @@
 }
 
 
+def _assert_residual_adds_are_quantized(onnx_save_path):
+    model = onnx.load(onnx_save_path)
+    consumers = defaultdict(list)
+    for node in model.graph.node:
+        for input_name in node.input:
+            consumers[input_name].append(node)
+
+    residual_adds = [node for node in model.graph.node if node.op_type == "Add"]
+    assert len(residual_adds) == 16
+    for add in residual_adds:
+        add_consumers = consumers[add.output[0]]
+        assert len(add_consumers) == 1
+        quantizer_input = add_consumers[0]
+        if quantizer_input.op_type == "Cast":
+            add_consumers = consumers[quantizer_input.output[0]]
+        assert len(add_consumers) == 1
+        assert add_consumers[0].op_type.endswith("QuantizeLinear")
+
+
 @pytest.mark.parametrize("quantize_mode", _QUANT_MODES)
 @pytest.mark.parametrize("model_key", list(_MODELS))
-def test_torch_onnx(model_key, quantize_mode):
+def test_torch_onnx(tmp_path, model_key, quantize_mode):
     timm_model_name, model_kwargs = _MODELS[model_key]
-    onnx_save_path = f"{model_key}.{quantize_mode}.onnx"
+    onnx_save_path = tmp_path / f"{model_key}.{quantize_mode}.onnx"
 
     cmd_parts = extend_cmd_parts(
         ["python", "torch_quant_to_onnx.py"],
@@ -46,3 +68,6 @@ def test_torch_onnx(model_key, quantize_mode):
     )
     cmd_parts.extend(["--no_pretrained", "--trt_build"])
     run_example_command(cmd_parts, "torch_onnx")
+
+    if model_key == "resnet50":
+        _assert_residual_adds_are_quantized(onnx_save_path)