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

CHANGELOG.rst

@@ -63,6 +63,8 @@ Experimental
 
 **Bug Fixes**
 
+- Add FP8, INT8, and AutoQuantize recipes that quantize timm ResNet shortcut inputs
+  immediately before residual adds.
 - 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.
+- Uses ResNet FP8, INT8, and AutoQuantize recipes to quantize shortcut inputs immediately before residual adds.
 - Exports the quantized model to ONNX.
 - Postprocesses the ONNX model to be compatible with TensorRT.
 - Saves the final ONNX model.
@@ -65,7 +66,7 @@ The `torch_quant_to_onnx.py` script quantizes [timm](https://github.com/huggingf
 ```bash
 python torch_quant_to_onnx.py \
     --timm_model_name=<timm model name> \
-    --quantize_mode=<fp8|mxfp8|int8|nvfp4|int4_awq> \
+    --quantize_mode=<fp8|mxfp8|int8|nvfp4|int4_awq|auto> \
     --onnx_save_path=<path to save the exported ONNX model>
 ```
 
@@ -74,6 +75,9 @@ Quantization configs are loaded from the YAML preset recipes under
 `--recipe=<preset basename or path to a QuantizeConfig YAML>` to use a different
 recipe (e.g. `--recipe=nvfp4_awq_lite` or `--recipe=/path/to/my_quant_cfg.yaml`).
 
+ResNet AutoQuantize searches FP8 and INT8 at eight effective bits and keeps its
+residual connections in FP8.
+
 ### Conv2d Quantization Override
 
 TensorRT only supports FP8 and INT8 for convolution operations. When quantizing models with Conv2d layers (like SwinTransformer), the script automatically applies the following overrides:

examples/torch_onnx/torch_quant_to_onnx.py

@@ -34,9 +34,11 @@
 from evaluation import evaluate
 
 import modelopt.torch.quantization as mtq
-from modelopt.recipe import load_config
+from modelopt.recipe import load_config, load_recipe
 from modelopt.recipe.presets import MODEL_QUANT_PRESET_DIR
 from modelopt.torch.quantization.config import QuantizeConfig
+from modelopt.torch.quantization.nn import TensorQuantizer
+from modelopt.torch.quantization.plugins.custom import CUSTOM_POST_CONVERSION_PLUGINS
 
 """
 Quantize a timm vision model and export to ONNX for TensorRT deployment.
@@ -70,6 +72,13 @@ def load_quant_config(recipe: str) -> dict:
     return load_config(recipe, schema_type=QuantizeConfig).model_dump()
 
 
+_RESNET_RECIPES = {
+    "fp8": "timm/resnet/ptq/fp8",
+    "int8": "timm/resnet/ptq/int8",
+    "auto": "timm/resnet/auto_quantize/fp8_int8_at_8p0bits",
+}
+
+
 _FP8_CONV_OVERRIDE: list = [
     {
         "parent_class": "nn.Conv2d",
@@ -151,6 +160,27 @@ def get_quant_config(quantize_mode, recipe=None):
     return config
 
 
+def _get_resnet_recipe(model, quantize_mode):
+    if not isinstance(model, timm.models.resnet.ResNet):
+        return None
+    recipe_path = _RESNET_RECIPES.get(quantize_mode)
+    return load_recipe(recipe_path) if recipe_path is not None else None
+
+
+def _add_resnet_residual_quantizers(model):
+    block_types = (timm.models.resnet.BasicBlock, timm.models.resnet.Bottleneck)
+    for block in (module for module in model.modules() if isinstance(module, block_types)):
+        if block.downsample is None:
+            block.downsample = torch.nn.Sequential()
+        elif not isinstance(block.downsample, torch.nn.Sequential):
+            block.downsample = torch.nn.Sequential(block.downsample)
+        if "residual_quantizer" in block.downsample._modules:
+            continue
+        residual_quantizer = TensorQuantizer()
+        residual_quantizer.disable()
+        block.downsample.add_module("residual_quantizer", residual_quantizer)
+
+
 def filter_func(name):
     """Filter function to exclude certain layers from quantization.
 
@@ -366,9 +396,10 @@ def auto_quantize_model(
     model,
     data_loader,
     quantization_formats,
-    effective_bits=4.8,
+    effective_bits=None,
     num_calib_steps=512,
     num_score_steps=128,
+    recipe=None,
 ):
     """Auto-quantize the model using optimal per-layer quantization search.
 
@@ -385,26 +416,48 @@ def auto_quantize_model(
         Tuple of (quantized_model, search_state_dict)
     """
     _disable_inplace_relu(model)
-    constraints = {"effective_bits": effective_bits}
-
-    # Convert string format names to config objects, incorporating Conv2d TRT overrides.
-    # TRT DynamicQuantize requires 2D/3D input, but Conv2d operates on 4D tensors.
-    # By including the overrides in the format configs, the auto_quantize search
-    # correctly accounts for Conv2d being FP8/INT8 in the effective_bits budget.
-    format_configs: list[dict[str, Any] | str] = []
-    for fmt in quantization_formats:
-        if isinstance(fmt, str):
-            config = load_quant_config(fmt)
-            if fmt in _NEEDS_FP8_CONV_OVERRIDE:
-                config["quant_cfg"].extend(_FP8_CONV_OVERRIDE)
-            elif fmt in _NEEDS_INT8_CONV_OVERRIDE:
-                config["quant_cfg"].extend(_INT8_CONV_OVERRIDE)
-            format_configs.append(config)
-        else:
-            format_configs.append(fmt)
-
-    print(f"Starting auto-quantization search with {len(format_configs)} formats...")
-    print(f"Effective bits constraint: {effective_bits}")
+    if recipe is None:
+        constraints = {"effective_bits": 4.8 if effective_bits is None else effective_bits}
+        format_configs: list[dict[str, Any] | str] = []
+        for fmt in quantization_formats:
+            if isinstance(fmt, str):
+                config = load_quant_config(fmt)
+                if fmt in _NEEDS_FP8_CONV_OVERRIDE:
+                    config["quant_cfg"].extend(_FP8_CONV_OVERRIDE)
+                elif fmt in _NEEDS_INT8_CONV_OVERRIDE:
+                    config["quant_cfg"].extend(_INT8_CONV_OVERRIDE)
+                format_configs.append(config)
+            else:
+                format_configs.append(fmt)
+        fixed_quantization_config = None
+        module_search_spaces = None
+        disabled_layers = None
+        method = "gradient"
+    else:
+        auto_config = recipe.auto_quantize
+        constraints = auto_config.constraints.model_dump(exclude_none=True)
+        if effective_bits is not None:
+            constraints["effective_bits"] = effective_bits
+        format_configs = []
+        fixed_quantization_config = recipe.quantize.model_dump()
+        module_search_spaces = [
+            {
+                "module_name_patterns": search_space.module_name_patterns,
+                "quantization_formats": [
+                    candidate.model_dump() for candidate in search_space.candidate_formats
+                ],
+                "allow_no_quant": search_space.allow_no_quant,
+            }
+            for search_space in auto_config.module_search_spaces
+        ]
+        disabled_layers = auto_config.disabled_layers
+        method = auto_config.auto_quantize_method
+
+    format_count = len(format_configs) or sum(
+        len(search_space["quantization_formats"]) for search_space in module_search_spaces or []
+    )
+    print(f"Starting auto-quantization search with {format_count} formats...")
+    print(f"Effective bits constraint: {constraints['effective_bits']}")
     print(f"Calibration steps: {num_calib_steps}, Scoring steps: {num_score_steps}")
 
     quantized_model, search_state = mtq.auto_quantize(
@@ -417,6 +470,10 @@ def auto_quantize_model(
         num_calib_steps=num_calib_steps,
         num_score_steps=num_score_steps,
         verbose=True,
+        fixed_quantization_config=fixed_quantization_config,
+        module_search_spaces=module_search_spaces,
+        disabled_layers=disabled_layers,
+        method=method,
     )
 
     # Disable quantization for specified layers
@@ -513,8 +570,11 @@ def main():
     parser.add_argument(
         "--effective_bits",
         type=float,
-        default=4.8,
-        help="Target effective bits for auto quantization constraint. Default is 4.8.",
+        default=None,
+        help=(
+            "Target effective bits for auto quantization. Defaults to 4.8 without a recipe "
+            "and overrides the ResNet recipe when provided."
+        ),
     )
     parser.add_argument(
         "--num_score_steps",
@@ -571,6 +631,10 @@ def main():
         )
         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)
+    if resnet_recipe is not None:
+        CUSTOM_POST_CONVERSION_PLUGINS.add(_add_resnet_residual_quantizers)
+
     # Quantize model based on mode
     if args.quantize_mode == "auto":
         # Auto quantization requires labels for loss computation
@@ -589,13 +653,18 @@ def main():
             args.effective_bits,
             args.calibration_data_size,
             args.num_score_steps,
+            recipe=resnet_recipe,
         )
     else:
         # Standard quantization - load calibration data
         # Note: MXFP8 is dynamic and does not need calibration itself, but when
         # Conv2d layers are overridden to FP8 (for TRT compatibility), those FP8
         # quantizers require calibration data.
-        config = get_quant_config(args.quantize_mode, args.recipe)
+        config = (
+            resnet_recipe.quantize.model_dump()
+            if resnet_recipe is not None
+            else get_quant_config(args.quantize_mode, args.recipe)
+        )
 
         data_loader = load_calibration_data(
             model,
@@ -612,6 +681,7 @@ def main():
     # Auto mode also needs this when an MXFP8/NVFP4 candidate format is in the search set.
     uses_dynamic_quantize = args.quantize_mode in ("mxfp8", "nvfp4") or (
         args.quantize_mode == "auto"
+        and resnet_recipe is None
         and any(fmt in _NEEDS_FP8_CONV_OVERRIDE for fmt in args.auto_quantization_formats)
     )
     if uses_dynamic_quantize:
@@ -621,7 +691,10 @@ def main():
     # no tactic for that 3-channel Q→Conv fusion. Skip for pure INT8 (unaffected).
     uses_fp8_conv_input = args.quantize_mode in ("fp8", "mxfp8", "nvfp4") or (
         args.quantize_mode == "auto"
-        and any(fmt not in {"int8", "int4_awq"} for fmt in args.auto_quantization_formats)
+        and (
+            resnet_recipe is not None
+            or any(fmt not in {"int8", "int4_awq"} for fmt in args.auto_quantization_formats)
+        )
     )
     if uses_fp8_conv_input:
         _disable_low_channel_conv_input_quantizers(quantized_model)

modelopt_recipes/timm/resnet/auto_quantize/fp8_int8_at_8p0bits.yaml

@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe
+imports:
+  base_disable_all: configs/ptq/units/base_disable_all
+  fp8: configs/numerics/fp8
+  fp8_model: configs/ptq/presets/model/fp8
+  int8_model: configs/ptq/presets/model/int8
+
+metadata:
+  recipe_type: auto_quantize
+  description: FP8 and INT8 ResNet AutoQuantize with FP8 residual connections.
+
+quantize:
+  algorithm: max
+  quant_cfg:
+    - $import: base_disable_all
+    - quantizer_name: '*residual_quantizer'
+      cfg:
+        $import: fp8
+
+auto_quantize:
+  constraints:
+    effective_bits: 8.0
+  module_search_spaces:
+    - module_name_patterns:
+        - '*'
+      candidate_formats:
+        - $import: fp8_model
+        - $import: int8_model
+      allow_no_quant: false
+  auto_quantize_method: gradient
+  score_size: 128

modelopt_recipes/timm/resnet/ptq/fp8.yaml

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe
+imports:
+  base_disable_all: configs/ptq/units/base_disable_all
+  default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers
+  fp8: configs/numerics/fp8
+  w8a8_fp8_fp8: configs/ptq/units/w8a8_fp8_fp8
+
+metadata:
+  recipe_type: ptq
+  description: FP8 ResNet PTQ with quantized residual connections.
+
+quantize:
+  algorithm: max
+  quant_cfg:
+    - $import: base_disable_all
+    - $import: w8a8_fp8_fp8
+    - $import: default_disabled_quantizers
+    - quantizer_name: '*residual_quantizer'
+      cfg:
+        $import: fp8

modelopt_recipes/timm/resnet/ptq/int8.yaml

@@ -0,0 +1,28 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe
+imports:
+  base_disable_all: configs/ptq/units/base_disable_all
+  default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers
+  int8: configs/numerics/int8
+  int8_per_channel: configs/numerics/int8_per_channel
+
+metadata:
+  recipe_type: ptq
+  description: INT8 ResNet PTQ with quantized residual connections.
+
+quantize:
+  algorithm: max
+  quant_cfg:
+    - $import: base_disable_all
+    - quantizer_name: '*weight_quantizer'
+      cfg:
+        $import: int8_per_channel
+    - quantizer_name: '*input_quantizer'
+      cfg:
+        $import: int8
+    - $import: default_disabled_quantizers
+    - quantizer_name: '*residual_quantizer'
+      cfg:
+        $import: int8

tests/examples/torch_onnx/test_torch_quant_to_onnx.py

@@ -14,13 +14,16 @@
 # 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
 
 # TODO: Add int4_awq once the INT4 exporter supports non-MatMul/Gemm consumer patterns
 # (e.g., DQ -> Reshape -> Slice in small ViT / SwinTransformer ONNX graphs).
 _QUANT_MODES = ["fp8", "int8", "mxfp8", "nvfp4", "auto"]
+_RESNET_QUANT_MODES = {"fp8", "int8", "auto"}
 
 _MODELS = {
     "vit_tiny": ("vit_tiny_patch16_224", '{"depth": 1}'),
@@ -30,6 +33,36 @@
 }
 
 
+def _assert_residual_inputs_are_quantized(onnx_save_path):
+    model = onnx.load(onnx_save_path)
+    consumers = defaultdict(list)
+    producers = {}
+    for node in model.graph.node:
+        for input_name in node.input:
+            consumers[input_name].append(node)
+        for output_name in node.output:
+            producers[output_name] = node
+
+    residual_adds = [
+        node
+        for node in model.graph.node
+        if node.op_type == "Add"
+        and [consumer.op_type for consumer in consumers[node.output[0]]] == ["Relu"]
+    ]
+    assert len(residual_adds) == 16
+    for add in residual_adds:
+        input_producers = [producers[input_name] for input_name in add.input]
+        input_producers = [
+            producers[node.input[0]] if node.op_type == "Cast" else node for node in input_producers
+        ]
+        assert any(
+            node.op_type.endswith("DequantizeLinear")
+            and "/downsample/residual_quantizer/" in node.name
+            and producers[node.input[0]].op_type.endswith("QuantizeLinear")
+            for node in input_producers
+        )
+
+
 @pytest.mark.parametrize("quantize_mode", _QUANT_MODES)
 @pytest.mark.parametrize("model_key", list(_MODELS))
 def test_torch_onnx(model_key, quantize_mode):
@@ -48,6 +81,9 @@ 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" and quantize_mode in _RESNET_QUANT_MODES:
+        _assert_residual_inputs_are_quantized(onnx_save_path)
+
 
 def test_torch_onnx_recipe_flag(tmp_path):
     timm_model_name, model_kwargs = _MODELS["vit_tiny"]