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

CHANGELOG.rst

@@ -81,6 +81,8 @@ Changelog
 
 **Bug Fixes**
 
+- Share each ResNet block-input Q/DQ between the main and shortcut paths in the torch ONNX
+  example, restoring TensorRT residual fusion for INT8 and FP8 exports.
 - 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 shortcut inputs before residual addition 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.
@@ -88,6 +90,13 @@
     },
 ]
 
+_FP8_RESIDUAL_OVERRIDE: list = [
+    {
+        "quantizer_name": "*residual_quantizer.input_quantizer",
+        "cfg": {"num_bits": (4, 3), "axis": None},
+    },
+]
+
 # FP8 MHA-aware config entries: quantize LayerNorm output so TRT can fuse the shared
 # Q/DQ across all downstream Q/K/V/FC consumers. Softmax-output Q/DQ is handled by the
 # FP8 ONNX exporter's post-processing pass (fixed 1/448 scale, data-independent).
@@ -133,6 +142,7 @@ def get_quant_config(quantize_mode):
             f"Overriding Conv2d quantization to FP8 for '{quantize_mode}' mode."
         )
         config["quant_cfg"].extend(_FP8_CONV_OVERRIDE)
+        config["quant_cfg"].extend(_FP8_RESIDUAL_OVERRIDE)
     elif quantize_mode == "int4_awq":
         warnings.warn(
             "TensorRT only supports FP8/INT8 for Conv layers. "
@@ -142,15 +152,19 @@ def get_quant_config(quantize_mode):
     return config
 
 
-def filter_func(name):
+def filter_func(name, model=None):
     """Filter function to exclude certain layers from quantization.
 
+    A ResNet's top-level ``conv1`` consumes three-channel image input, for which TensorRT
+    cannot find an FP8 Q→Conv tactic on Blackwell.
     ``downsample.reduction`` (Swin/SwinV2) is excluded because it operates on 4D tensors
     and TRT's DynamicQuantize layer (used for MXFP8/NVFP4) requires 2D/3D input.
     Other 4D-input layers (e.g. Swin's ``norm1``, ``downsample.norm``, top-level ``norm``)
     are handled dynamically by ``_disable_high_rank_input_quantizers`` via a forward-pass
     rank probe — that avoids false positives on ViT, whose same-named ``norm`` sees 3D input.
     """
+    if isinstance(model, timm.models.resnet.ResNet) and name.startswith("conv1."):
+        return True
     pattern = re.compile(
         r".*(time_emb_proj|time_embedding|conv_in|conv_out|conv_shortcut|add_embedding|"
         r"pos_embed|time_text_embed|context_embedder|norm_out|x_embedder|patch_embed|cpb_mlp|"
@@ -194,28 +208,171 @@ def hook(m, inp, out, _n=name):
     mtq.disable_quantizer(model, lambda n: n.startswith(prefixes))
 
 
-def _disable_low_channel_conv_input_quantizers(model):
-    """Disable ``input_quantizer`` on Conv2d modules whose ``in_channels <= 3``.
+def _quantize_module_input(module, inputs):
+    return (module.input_quantizer(inputs[0]), *inputs[1:])
 
-    The first Conv2d of an image backbone (e.g. ResNet50's ``conv1``) consumes raw
-    RGB input, so ``in_channels == 3``. On Blackwell (compute capability 12.0) TRT
-    fails to find an FP8/MXFP8/NVFP4 tactic for this first-layer Q→Conv fusion:
 
-        Error Code 10: Could not find any implementation for node
-        /conv1/input_quantizer/TRT_FP8QuantizeLinear ... [ElementWise]
+def _calibrate_new_quantizers(model, quantizers, data_loader):
+    enabled_quantizers = [
+        module
+        for module in model.modules()
+        if isinstance(module, TensorQuantizer) and module.is_enabled
+    ]
+    for quantizer in enabled_quantizers:
+        quantizer.disable_quant()
+    for quantizer in quantizers:
+        quantizer.enable_calib()
 
-    Ada (8.9) happens to have a tactic, which is why local runs pass. Disabling the
-    input quantizer on the raw-RGB conv is also standard quantization practice —
-    first/last layers are typically left in higher precision. Weight quantization
-    still applies. Swin/ViT's ``patch_embed.proj`` is already excluded via
-    ``filter_func``'s ``patch_embed`` pattern, so this helper is effectively the
-    ResNet-shaped analogue.
+    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)
+        for quantizer in quantizers:
+            quantizer.load_calib_amax(strict=False)
+    finally:
+        model.train(was_training)
+        for quantizer in quantizers:
+            quantizer.disable_calib()
+        for quantizer in enabled_quantizers:
+            quantizer.enable_quant()
+
+    _disable_invalid_quantizers(quantizers)
+
+
+def _append_resnet_residual_quantizer(block, num_bits):
+    device = next(block.parameters()).device
+    quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=num_bits, axis=None)).to(device)
+    residual_quantizer = torch.nn.Sequential()
+    residual_quantizer.add_module("input_quantizer", quantizer)
+    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)
+    block.downsample.add_module("residual_quantizer", residual_quantizer)
+    return quantizer
+
+
+def _prepare_resnet_quantizers(model, quantize_mode):
+    """Install ResNet shortcut quantizers before the standard calibration pass.
+
+    INT8 and FP8 share one block-input Q/DQ between ``conv1`` and an identity shortcut;
+    projection blocks additionally quantize the downsample output immediately before ``Add``.
+    MXFP8 and NVFP4 use per-tensor FP8 on every 4D shortcut because their dynamic block
+    quantizers only support 2D/3D tensors. INT4-AWQ is skipped because it is weight-only.
+    Auto mode is configured after its per-block format search.
     """
-    for _, mod in model.named_modules():
-        if isinstance(mod, torch.nn.Conv2d) and mod.in_channels <= 3:
-            q = getattr(mod, "input_quantizer", None)
-            if q is not None and q.is_enabled:
-                q.disable()
+    block_types = (timm.models.resnet.BasicBlock, timm.models.resnet.Bottleneck)
+    blocks = [block for block in model.modules() if isinstance(block, block_types)]
+    if (
+        not blocks
+        or quantize_mode in ("auto", "int4_awq")
+        or any(
+            hasattr(block, "input_quantizer")
+            or (block.downsample is not None and hasattr(block.downsample, "residual_quantizer"))
+            for block in blocks
+        )
+    ):
+        return []
+
+    num_bits = 8 if quantize_mode == "int8" else (4, 3)
+    device = next(model.parameters()).device
+    residual_quantizers = []
+    for block in blocks:
+        if quantize_mode in ("int8", "fp8"):
+            quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=num_bits, axis=None)).to(
+                device
+            )
+            block.add_module("input_quantizer", quantizer)
+            block.register_forward_pre_hook(_quantize_module_input)
+            residual_quantizers.append(quantizer)
+
+            if block.downsample is None:
+                continue
+
+        residual_quantizers.append(_append_resnet_residual_quantizer(block, num_bits))
+
+    if quantize_mode == "int8":
+        quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=num_bits, axis=None)).to(
+            device
+        )
+        model.global_pool.add_module("input_quantizer", quantizer)
+        model.global_pool.register_forward_pre_hook(_quantize_module_input)
+        residual_quantizers.append(quantizer)
+
+    return residual_quantizers
+
+
+def _disable_invalid_quantizers(quantizers):
+    for quantizer in quantizers:
+        if not quantizer.is_enabled:
+            continue
+        amax = quantizer.amax
+        if (
+            amax is None
+            or not torch.is_tensor(amax)
+            or torch.any(torch.isnan(amax))
+            or torch.all(amax <= 0)
+        ):
+            quantizer.disable()
+
+
+def _add_auto_resnet_residual_quantizers(model, auto_quantization_formats, data_loader):
+    """Calibrate shortcuts with each block's AutoQuantize-selected activation format."""
+    activation_formats = set(auto_quantization_formats) - {"INT4_AWQ_CFG"}
+    if not activation_formats:
+        return
+    fallback_num_bits = 8 if activation_formats == {"INT8_DEFAULT_CFG"} else (4, 3)
+    residual_quantizers = []
+    block_types = (timm.models.resnet.BasicBlock, timm.models.resnet.Bottleneck)
+    for block in model.modules():
+        if not isinstance(block, block_types):
+            continue
+        last_conv = block.conv3 if isinstance(block, timm.models.resnet.Bottleneck) else block.conv2
+        selected_quantizer = getattr(last_conv, "input_quantizer", None)
+        num_bits = (
+            selected_quantizer.num_bits
+            if selected_quantizer is not None
+            and selected_quantizer.is_enabled
+            and selected_quantizer.num_bits in (8, (4, 3))
+            else fallback_num_bits
+        )
+        residual_quantizers.append(_append_resnet_residual_quantizer(block, num_bits))
+
+    _calibrate_new_quantizers(model, residual_quantizers, data_loader)
+
+
+def _finalize_resnet_quantizers(
+    model, quantize_mode, auto_quantization_formats, data_loader, residual_quantizers
+):
+    block_types = (timm.models.resnet.BasicBlock, timm.models.resnet.Bottleneck)
+    blocks = [block for block in model.modules() if isinstance(block, block_types)]
+    if not blocks:
+        return
+
+    if quantize_mode in ("int8", "fp8"):
+        for block in blocks:
+            block.conv1.input_quantizer.disable()
+            if block.downsample is not None:
+                downsample_conv = next(
+                    (
+                        module
+                        for module in block.downsample.modules()
+                        if isinstance(module, torch.nn.Conv2d)
+                    ),
+                    None,
+                )
+                if downsample_conv is not None:
+                    downsample_conv.input_quantizer.disable()
+        model.fc.input_quantizer.disable()
+        model.fc.weight_quantizer.disable()
+        if quantize_mode == "int8":
+            model.global_pool.input_quantizer.enable()
+    elif quantize_mode == "auto":
+        _add_auto_resnet_residual_quantizers(model, auto_quantization_formats, data_loader)
+
+    _disable_invalid_quantizers(residual_quantizers)
 
 
 def load_calibration_data(model, data_size, batch_size, device, with_labels=False):
@@ -318,7 +475,7 @@ def forward_loop(model):
 
     # Disable filtered quantizers BEFORE calibrating override quantizers so we don't
     # waste time calibrating quantizers that are about to be turned off.
-    mtq.disable_quantizer(quantized_model, filter_func)
+    mtq.disable_quantizer(quantized_model, lambda name: filter_func(name, quantized_model))
 
     # Calibrate any FP8 override quantizers that weren't calibrated by mtq.quantize().
     if data_loader is not None:
@@ -410,7 +567,7 @@ def auto_quantize_model(
     )
 
     # Disable quantization for specified layers
-    mtq.disable_quantizer(quantized_model, filter_func)
+    mtq.disable_quantizer(quantized_model, lambda name: filter_func(name, quantized_model))
 
     _disable_dead_quantizers(quantized_model)
 
@@ -544,6 +701,8 @@ def main():
         )
         print(f"Base Model - Top-1 Accuracy: {top1:.2f}%, Top-5 Accuracy: {top5:.2f}%")
 
+    residual_quantizers = []
+
     # Quantize model based on mode
     if args.quantize_mode == "auto":
         # Auto quantization requires labels for loss computation
@@ -569,6 +728,14 @@ def main():
         # Conv2d layers are overridden to FP8 (for TRT compatibility), those FP8
         # quantizers require calibration data.
         config = get_quant_config(args.quantize_mode)
+        block_types = (timm.models.resnet.BasicBlock, timm.models.resnet.Bottleneck)
+        if args.quantize_mode != "int4_awq" and any(
+            isinstance(module, block_types) for module in model.modules()
+        ):
+            conversion_config = copy.deepcopy(config)
+            conversion_config["algorithm"] = None
+            model = mtq.quantize(model, conversion_config)
+            residual_quantizers = _prepare_resnet_quantizers(model, args.quantize_mode)
 
         data_loader = load_calibration_data(
             model,
@@ -580,6 +747,14 @@ def main():
 
         quantized_model = quantize_model(model, config, data_loader)
 
+    _finalize_resnet_quantizers(
+        quantized_model,
+        args.quantize_mode,
+        args.auto_quantization_formats,
+        data_loader,
+        residual_quantizers,
+    )
+
     # 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.
@@ -590,15 +765,6 @@ def main():
     if uses_dynamic_quantize:
         _disable_high_rank_input_quantizers(quantized_model, input_shape, device)
 
-    # FP8-family modes emit TRT_FP8QuantizeLinear on the first-layer conv; Blackwell has
-    # 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 != "INT8_DEFAULT_CFG" for fmt in args.auto_quantization_formats)
-    )
-    if uses_fp8_conv_input:
-        _disable_low_channel_conv_input_quantizers(quantized_model)
-
     # Print quantization summary
     print("\nQuantization Summary:")
     mtq.print_quant_summary(quantized_model)

modelopt/onnx/export/fp8_exporter.py

@@ -172,13 +172,17 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int:
         2. Quantize weights to FP8E4M3FN
         3. Insert a DequantizeLinear(fp8_weights, scale) before the Conv weight input
 
+        An RGB Conv that directly consumes an unquantized graph input is treated as a
+        filtered input stem and left entirely in high precision.
+
         Args:
             graph: The onnx-graphsurgeon graph to modify in-place.
 
         Returns:
             Number of Conv weight DQ nodes inserted.
         """
         count = 0
+        graph_inputs = {tensor.name for tensor in graph.inputs}
 
         for node in list(graph.nodes):
             if node.op != "Conv":
@@ -189,7 +193,8 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int:
             weight_input = node.inputs[1]
             if not isinstance(weight_input, gs.Constant):
                 continue
-
+            if node.inputs[0].name in graph_inputs and weight_input.values.shape[1] == 3:
+                continue
             # Skip if weight already has a DQ producer
             if any(out.op == "DequantizeLinear" for out in weight_input.outputs):
                 continue