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

CHANGELOG.rst

@@ -81,6 +81,10 @@ Changelog
 
 **Bug Fixes**
 
+- Add timm ResNet PTQ recipes for FP8, INT8, MXFP8, and NVFP4. The recipes
+  keep the input stem convolution unquantized, select TensorRT-compatible convolution formats,
+  and place shared block-input and projection-shortcut Q/DQ before residual adds, including
+  during AutoQuantize.
 - 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,8 @@ 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 the [timm ResNet PTQ recipes](../../modelopt_recipes/timm/resnet/ptq/) to
+  quantize shortcut inputs before residual addition.
 - Exports the quantized model to ONNX.
 - Postprocesses the ONNX model to be compatible with TensorRT.
 - Saves the final ONNX model.
@@ -273,7 +275,7 @@ The `auto` mode enables mixed precision quantization by searching for the optima
 
 | Parameter | Default | Description |
 | :--- | :---: | :--- |
-| `--effective_bits` | 4.8 | Target average bits per weight across the model. Lower values = more compression but potentially lower accuracy. The search algorithm finds the optimal per-layer format assignment that meets this constraint while minimizing accuracy loss. For example, 4.8 means an average of 4.8 bits per weight (mix of FP4 and FP8 layers). |
+| `--effective_bits` | 4.8 (8.0 for ResNet) | Target average bits per weight across the model. Lower values = more compression but potentially lower accuracy. The ResNet default remains feasible when Conv2d candidates use FP8 or INT8. |
 | `--num_score_steps` | 128 | Number of forward/backward passes used to compute per-layer sensitivity scores via gradient-based analysis. Higher values provide more accurate sensitivity estimates but increase search time. Recommended range: 64-256. |
 | `--calibration_data_size` | 512 | Number of calibration samples used for both sensitivity scoring and calibration. For auto mode, labels are required for loss computation. |
 

examples/torch_onnx/torch_quant_to_onnx.py

@@ -34,6 +34,9 @@
 from evaluation import evaluate
 
 import modelopt.torch.quantization as mtq
+from modelopt.recipe import load_recipe
+from modelopt.torch.quantization.nn import TensorQuantizer
+from modelopt.torch.quantization.plugins.timm import is_resnet_quantization_supported
 
 """
 Quantize a timm vision model and export to ONNX for TensorRT deployment.
@@ -61,6 +64,15 @@
     "nvfp4": mtq.NVFP4_DEFAULT_CFG,
     "int4_awq": mtq.INT4_AWQ_CFG,
 }
+_RESNET_RECIPE_DIR = "timm/resnet/ptq"
+_RESNET_QUANTIZE_MODES = {"fp8", "int8", "mxfp8", "nvfp4"}
+_RESNET_AUTO_RECIPE_NAMES = {
+    "FP8_DEFAULT_CFG": "fp8",
+    "INT8_DEFAULT_CFG": "int8",
+    "MXFP8_DEFAULT_CFG": "mxfp8",
+    "NVFP4_DEFAULT_CFG": "nvfp4",
+    "NVFP4_AWQ_LITE_CFG": "nvfp4_awq_lite",
+}
 
 _FP8_CONV_OVERRIDE: list = [
     {
@@ -114,7 +126,7 @@
 _NEEDS_INT8_CONV_OVERRIDE: set[str] = {"INT4_AWQ_CFG"}
 
 
-def get_quant_config(quantize_mode):
+def get_quant_config(quantize_mode, model=None):
     """Get quantization config, overriding Conv2d for TRT compatibility.
 
     TensorRT only supports FP8 and INT8 for Conv layers.
@@ -124,6 +136,9 @@ def get_quant_config(quantize_mode):
     - For MXFP8, NVFP4: override Conv2d to FP8
     - For INT4_AWQ: override Conv2d to INT8
     """
+    if is_resnet_quantization_supported(model) and quantize_mode in _RESNET_QUANTIZE_MODES:
+        return load_recipe(f"{_RESNET_RECIPE_DIR}/{quantize_mode}").quantize.model_dump()
+
     config: dict = copy.deepcopy(QUANT_CONFIG_DICT[quantize_mode])
     if quantize_mode == "fp8":
         config["quant_cfg"].extend(_FP8_MHA_OVERRIDE)
@@ -142,6 +157,19 @@ def get_quant_config(quantize_mode):
     return config
 
 
+def get_auto_quant_config(format_name, model):
+    recipe_name = _RESNET_AUTO_RECIPE_NAMES.get(format_name)
+    if is_resnet_quantization_supported(model) and recipe_name is not None:
+        return load_recipe(f"{_RESNET_RECIPE_DIR}/{recipe_name}").quantize.model_dump()
+
+    config = copy.deepcopy(getattr(mtq, format_name))
+    if format_name in _NEEDS_FP8_CONV_OVERRIDE:
+        config["quant_cfg"].extend(_FP8_CONV_OVERRIDE)
+    elif format_name in _NEEDS_INT8_CONV_OVERRIDE:
+        config["quant_cfg"].extend(_INT8_CONV_OVERRIDE)
+    return config
+
+
 def filter_func(name):
     """Filter function to exclude certain layers from quantization.
 
@@ -194,30 +222,6 @@ 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``.
-
-    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]
-
-    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.
-    """
-    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()
-
-
 def load_calibration_data(model, data_size, batch_size, device, with_labels=False):
     """Load and prepare calibration data.
 
@@ -262,16 +266,14 @@ def _disable_dead_quantizers(model):
     calibrates to ``amax == 0``. Disable such dead quantizers — they have nothing
     meaningful to quantize and would otherwise break ONNX export.
     """
-    for _, mod in model.named_modules():
-        for attr in ("input_quantizer", "output_quantizer", "weight_quantizer"):
-            q = getattr(mod, attr, None)
-            if q is None or not q.is_enabled:
-                continue
-            amax = q.amax
-            if amax is None or not torch.is_tensor(amax):
-                continue
-            if torch.any(torch.isnan(amax)) or torch.all(amax <= 0):
-                q.disable()
+    for quantizer in model.modules():
+        if not isinstance(quantizer, TensorQuantizer) or not quantizer.is_enabled:
+            continue
+        amax = quantizer.amax
+        if amax is None or not torch.is_tensor(amax):
+            continue
+        if torch.any(torch.isnan(amax)) or torch.all(amax <= 0):
+            quantizer.disable()
 
 
 def _calibrate_uncalibrated_quantizers(model, data_loader):
@@ -281,15 +283,16 @@ def _calibrate_uncalibrated_quantizers(model, data_loader):
     be calibrated because the MXFP8/NVFP4 quantization pipeline skips standard
     calibration. This function explicitly calibrates those uncalibrated quantizers.
     """
-    uncalibrated = []
-    for _, module in model.named_modules():
-        for attr_name in ("input_quantizer", "weight_quantizer"):
-            if not hasattr(module, attr_name):
-                continue
-            quantizer = getattr(module, attr_name)
-            if quantizer.is_enabled and not quantizer.block_sizes and quantizer.amax is None:
-                quantizer.enable_calib()
-                uncalibrated.append(quantizer)
+    uncalibrated = [
+        quantizer
+        for quantizer in model.modules()
+        if isinstance(quantizer, TensorQuantizer)
+        and quantizer.is_enabled
+        and not quantizer.block_sizes
+        and quantizer.amax is None
+    ]
+    for quantizer in uncalibrated:
+        quantizer.enable_calib()
 
     if not uncalibrated:
         return
@@ -316,9 +319,8 @@ def forward_loop(model):
     else:
         quantized_model = mtq.quantize(model, config)
 
-    # 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)
+    if not is_resnet_quantization_supported(quantized_model):
+        mtq.disable_quantizer(quantized_model, filter_func)
 
     # Calibrate any FP8 override quantizers that weren't calibrated by mtq.quantize().
     if data_loader is not None:
@@ -357,7 +359,7 @@ def auto_quantize_model(
     model,
     data_loader,
     quantization_formats,
-    effective_bits=4.8,
+    effective_bits=None,
     num_calib_steps=512,
     num_score_steps=128,
 ):
@@ -375,23 +377,18 @@ def auto_quantize_model(
         Tuple of (quantized_model, search_state_dict)
     """
     _disable_inplace_relu(model)
+    if effective_bits is None:
+        effective_bits = 8.0 if is_resnet_quantization_supported(model) else 4.8
     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 = []
-    for fmt in quantization_formats:
-        if isinstance(fmt, str):
-            config = copy.deepcopy(getattr(mtq, 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)
+    format_configs = [
+        get_auto_quant_config(fmt, model) if isinstance(fmt, str) else fmt
+        for fmt in quantization_formats
+    ]
 
     print(f"Starting auto-quantization search with {len(format_configs)} formats...")
     print(f"Effective bits constraint: {effective_bits}")
@@ -409,8 +406,8 @@ def auto_quantize_model(
         verbose=True,
     )
 
-    # Disable quantization for specified layers
-    mtq.disable_quantizer(quantized_model, filter_func)
+    if not is_resnet_quantization_supported(quantized_model):
+        mtq.disable_quantizer(quantized_model, filter_func)
 
     _disable_dead_quantizers(quantized_model)
 
@@ -481,6 +478,7 @@ def main():
         nargs="+",
         choices=[
             "NVFP4_AWQ_LITE_CFG",
+            "NVFP4_DEFAULT_CFG",
             "FP8_DEFAULT_CFG",
             "MXFP8_DEFAULT_CFG",
             "INT8_DEFAULT_CFG",
@@ -492,8 +490,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 8.0 for ResNet "
+            "and 4.8 for other models."
+        ),
     )
     parser.add_argument(
         "--num_score_steps",
@@ -568,7 +569,7 @@ def main():
         # 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)
+        config = get_quant_config(args.quantize_mode, model)
 
         data_loader = load_calibration_data(
             model,
@@ -590,15 +591,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

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"