[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. "
@@ -195,29 +205,193 @@ def hook(m, inp, out, _n=name):
 
 
 def _disable_low_channel_conv_input_quantizers(model):
-    """Disable ``input_quantizer`` on Conv2d modules whose ``in_channels <= 3``.
+    """Disable input quantization on Conv2d modules with at most 16 input or output channels.
 
     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:
+    RGB input. TensorRT does not reliably accelerate FP8 convolutions with such small
+    channel dimensions, so ONNX PTQ leaves the entire Conv in high precision. On
+    Blackwell, TRT can also fail to find a tactic for the first-layer Q→Conv fusion:
 
         Error Code 10: Could not find any implementation for node
         /conv1/input_quantizer/TRT_FP8QuantizeLinear ... [ElementWise]
 
-    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.
+    Ada (8.9) happens to have a tactic, which is why local runs pass. Swin/ViT's
+    ``patch_embed.proj`` is already excluded via ``filter_func``.
     """
     for _, mod in model.named_modules():
-        if isinstance(mod, torch.nn.Conv2d) and mod.in_channels <= 3:
+        if isinstance(mod, torch.nn.Conv2d) and min(mod.in_channels, mod.out_channels) <= 16:
             q = getattr(mod, "input_quantizer", None)
             if q is not None and q.is_enabled:
                 q.disable()
 
 
+def _quantize_module_input(module, inputs):
+    return (module.input_quantizer(inputs[0]), *inputs[1:])
+
+
+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()
+
+    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.
+    """
+    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):
     """Load and prepare calibration data.
 
@@ -544,6 +718,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 +745,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 +764,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.

modelopt/onnx/export/fp8_exporter.py

@@ -31,6 +31,7 @@
 # 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
 
 
 class FP8QuantExporter(ONNXQuantExporter):
@@ -189,6 +190,10 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int:
             weight_input = node.inputs[1]
             if not isinstance(weight_input, gs.Constant):
                 continue
+            if any(
+                channels <= _FP8_MIN_CONV_CHANNELS for channels in weight_input.values.shape[:2]
+            ):
+                continue
 
             # Skip if weight already has a DQ producer
             if any(out.op == "DequantizeLinear" for out in weight_input.outputs):

modelopt/onnx/utils.py

@@ -28,6 +28,7 @@
 import onnx_graphsurgeon as gs
 from onnx.helper import get_attribute_value
 from onnx_graphsurgeon import Constant, Node, Variable
+from onnxconverter_common.float16 import convert_np_to_float16
 
 from modelopt.onnx.logging_config import logger
 
@@ -1467,13 +1468,17 @@ def fold_q_fp16_to_fp32_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
         for inp in node.input:
             consumer_map.setdefault(inp, []).append(node)
     initializers = {init.name: init for init in onnx_model.graph.initializer}
+    tensor_types = _build_tensor_type_map(onnx_model)
 
     to_remove = []
     for node in onnx_model.graph.node:
         if node.op_type != "Cast":
             continue
         cast_to = next((a.i for a in node.attribute if a.name == "to"), None)
-        if cast_to != onnx.TensorProto.FLOAT:
+        if (
+            cast_to != onnx.TensorProto.FLOAT
+            or tensor_types.get(node.input[0]) != onnx.TensorProto.FLOAT16
+        ):
             continue
         consumers = consumer_map.get(node.output[0], [])
         if not consumers or not all(c.op_type in _Q_OPS for c in consumers):
@@ -1492,6 +1497,53 @@ def fold_q_fp16_to_fp32_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
     return onnx_model
 
 
+def _convert_q_data_initializers_to_fp16(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
+    """Convert FP32 initializer data inputs after Q/DQ scales have been normalized to FP16."""
+    if get_opset_version(onnx_model) < BASE_MIN_OPSET:
+        return onnx_model
+
+    consumers: dict[str, list[tuple[onnx.NodeProto, int]]] = defaultdict(list)
+    for node in onnx_model.graph.node:
+        for index, input_name in enumerate(node.input):
+            consumers[input_name].append((node, index))
+
+    initializers = {initializer.name: initializer for initializer in onnx_model.graph.initializer}
+    tensor_types = _build_tensor_type_map(onnx_model)
+
+    for name, initializer in list(initializers.items()):
+        if initializer.data_type != onnx.TensorProto.FLOAT:
+            continue
+
+        initializer_consumers = consumers.get(name, [])
+        q_consumers = [
+            node for node, index in initializer_consumers if index == 0 and node.op_type in _Q_OPS
+        ]
+        if not q_consumers:
+            continue
+
+        for q_node in q_consumers:
+            scale_type = tensor_types.get(q_node.input[1]) if len(q_node.input) >= 2 else None
+            if scale_type != onnx.TensorProto.FLOAT16:
+                raise ValueError("Q scales must be FP16 before converting Q data initializers")
+
+        fp16_initializer = onnx.numpy_helper.from_array(
+            convert_np_to_float16(onnx.numpy_helper.to_array(initializer)), initializer.name
+        )
+        if len(q_consumers) == len(initializer_consumers):
+            initializer.CopyFrom(fp16_initializer)
+            continue
+
+        fp16_initializer.name = f"{initializer.name}_fp16_q"
+        while fp16_initializer.name in initializers:
+            fp16_initializer.name += "_"
+        onnx_model.graph.initializer.append(fp16_initializer)
+        initializers[fp16_initializer.name] = fp16_initializer
+        for node in q_consumers:
+            node.input[0] = fp16_initializer.name
+
+    return onnx_model
+
+
 def _is_foldable_constant_cast_pattern(model: onnx.ModelProto, node: onnx.NodeProto) -> bool:
     """Check if a Constant -> Cast pattern can be folded."""
     assert node.op_type == "Cast"