Add ResNet50 support for torch_onnx quantization workflow

Add end-to-end support for ResNet50 (Conv2d-heavy model) in the
torch_onnx quantization → ONNX export → TRT engine pipeline.

Key fixes for Conv2d-heavy models:
- Disable FP8 Conv2d weight quantizers during ONNX export to avoid
  TorchScript exporter's "kernel of unknown shape" error (FP8
  DequantizeLinear produces dynamic-shape outputs incompatible with
  Conv2d's static kernel requirement)
- Disable autocast for FP8/INT8 quantized models during export (prevents
  dynamic-shape kernels from autocast-induced FP16 casting)
- Fix configure_linear_module_onnx_quantizers to handle all modules with
  block quantization (not just nn.Linear), fixing NVFP4/MXFP8 export for
  models with quantized non-Linear modules like MaxPool2d
- Add calibration step for FP8 override quantizers that aren't calibrated
  by mtq.quantize() in MXFP8/NVFP4 modes
- Override Conv2d block quantizers to FP8 in auto mode for TRT compat
- Add maxpool and global_pool to filter_func (TRT DynamicQuantize
  requires 2D/3D input, but pooling layers operate on 4D tensors)
- Always load calibration data (MXFP8 Conv2d FP8 overrides need it)

Signed-off-by: ajrasane <arasane@nvidia.com>

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>

Author: ajrasane

examples/torch_onnx/README.md

@@ -307,6 +307,7 @@ python torch_quant_to_onnx.py \
 | [vit_base_patch16_224](https://huggingface.co/timm/vit_base_patch16_224.augreg_in21k_ft_in1k) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | [swin_tiny_patch4_window7_224](https://huggingface.co/timm/swin_tiny_patch4_window7_224.ms_in1k) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | [swinv2_tiny_window8_256](https://huggingface.co/timm/swinv2_tiny_window8_256.ms_in1k) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| [resnet50](https://huggingface.co/timm/resnet50.a1_in1k) | ✅ | ✅ | ✅ | ✅ | | ✅ |
 
 ## Resources
 

examples/torch_onnx/torch_quant_to_onnx.py

@@ -109,7 +109,8 @@ def filter_func(name):
     """Filter function to exclude certain layers from quantization."""
     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|downsample).*"
+        r"pos_embed|time_text_embed|context_embedder|norm_out|x_embedder|patch_embed|cpb_mlp|"
+        r"downsample|maxpool|global_pool).*"
     )
     return pattern.match(name) is not None
 
@@ -147,6 +148,36 @@ def load_calibration_data(model_name, data_size, batch_size, device, with_labels
         )
 
 
+def _calibrate_uncalibrated_quantizers(model, data_loader):
+    """Calibrate FP8 quantizers that weren't calibrated by mtq.quantize().
+
+    When MXFP8/NVFP4 modes override Conv2d to FP8, the FP8 quantizers may not
+    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 not hasattr(quantizer, "_amax"):
+                quantizer.enable_calib()
+                uncalibrated.append(quantizer)
+
+    if not uncalibrated:
+        return
+
+    model.eval()
+    with torch.no_grad():
+        for batch in data_loader:
+            model(batch)
+
+    for quantizer in uncalibrated:
+        quantizer.disable_calib()
+        quantizer.load_calib_amax()
+
+
 def quantize_model(model, config, data_loader=None):
     """Quantize the model using the given config and calibration data."""
     if data_loader is not None:
@@ -159,6 +190,10 @@ def forward_loop(model):
     else:
         quantized_model = mtq.quantize(model, config)
 
+    # Calibrate any FP8 override quantizers that weren't calibrated by mtq.quantize()
+    if data_loader is not None:
+        _calibrate_uncalibrated_quantizers(quantized_model, data_loader)
+
     mtq.disable_quantizer(quantized_model, filter_func)
     return quantized_model
 
@@ -185,6 +220,38 @@ def _disable_inplace_relu(model):
             module.inplace = False
 
 
+def _override_conv2d_to_fp8(model, data_loader):
+    """Override Conv2d layers with NVFP4/MXFP8 block quantization to FP8.
+
+    TRT DynamicQuantize requires 2D/3D input, but Conv2d operates on 4D tensors.
+    This overrides Conv2d block quantizers to FP8 per-tensor and calibrates them.
+    """
+    overridden = []
+    for _, module in model.named_modules():
+        if not isinstance(module, torch.nn.Conv2d):
+            continue
+        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 quantizer.block_sizes:
+                # Override to FP8 per-tensor
+                quantizer.block_sizes = None
+                quantizer._num_bits = (4, 3)
+                quantizer._axis = None
+                quantizer.enable_calib()
+                overridden.append(quantizer)
+
+    if overridden:
+        model.eval()
+        with torch.no_grad():
+            for batch in data_loader:
+                model(batch["image"])
+        for quantizer in overridden:
+            quantizer.disable_calib()
+            quantizer.load_calib_amax()
+
+
 def auto_quantize_model(
     model,
     data_loader,
@@ -233,6 +300,10 @@ def auto_quantize_model(
         verbose=True,
     )
 
+    # Override Conv2d layers that got NVFP4/MXFP8 to FP8 for TRT compatibility.
+    # TRT DynamicQuantize requires 2D/3D input, but Conv2d operates on 4D tensors.
+    _override_conv2d_to_fp8(quantized_model, data_loader)
+
     # Disable quantization for specified layers
     mtq.disable_quantizer(quantized_model, filter_func)
 
@@ -378,18 +449,18 @@ def main():
             args.num_score_steps,
         )
     else:
-        # Standard quantization - only load calibration data if needed
+        # 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)
-        if args.quantize_mode == "mxfp8":
-            data_loader = None
-        else:
-            data_loader = load_calibration_data(
-                args.timm_model_name,
-                args.calibration_data_size,
-                args.batch_size,
-                device,
-                with_labels=False,
-            )
+        data_loader = load_calibration_data(
+            args.timm_model_name,
+            args.calibration_data_size,
+            args.batch_size,
+            device,
+            with_labels=False,
+        )
 
         quantized_model = quantize_model(model, config, data_loader)
 

modelopt/torch/_deploy/utils/torch_onnx.py

@@ -16,6 +16,7 @@
 """Utility functions related to Onnx."""
 
 import base64
+import contextlib
 import inspect
 import json
 import logging
@@ -402,6 +403,29 @@ def is_fp8_quantized(model: nn.Module) -> bool:
     return False
 
 
+@contextlib.contextmanager
+def _disable_fp8_conv_weight_quantizers(model: nn.Module):
+    """Temporarily disable FP8 weight quantizers on Conv layers during ONNX export.
+
+    The TorchScript ONNX exporter requires static kernel shapes for Conv operations,
+    but FP8 weight quantization (TRT_FP8QuantizeLinear -> TRT_FP8DequantizeLinear)
+    produces dynamic-shape outputs that break this requirement. Disabling Conv weight
+    quantizers during export allows the Conv to export with static-shape FP16/FP32
+    weights. Conv activations still have FP8 QDQ nodes (input quantizers remain enabled).
+    """
+    disabled = []
+    for _, module in model.named_modules():
+        if isinstance(module, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+            if hasattr(module, "weight_quantizer") and module.weight_quantizer.is_enabled:
+                module.weight_quantizer.disable()
+                disabled.append(module)
+    try:
+        yield
+    finally:
+        for module in disabled:
+            module.weight_quantizer.enable()
+
+
 def quantize_weights(model: nn.Module, onnx_model: onnx.ModelProto) -> onnx.ModelProto:
     """Real quantizes the weights in the onnx model.
 
@@ -522,7 +546,11 @@ def get_onnx_bytes_and_metadata(
     input_none_names = list(set(tree_spec_input.names) - set(input_names))
 
     use_torch_autocast = not (
-        is_fp4_quantized(model) or is_mxfp8_quantized(model) or weights_dtype == "fp32"
+        is_fp4_quantized(model)
+        or is_mxfp8_quantized(model)
+        or is_fp8_quantized(model)
+        or is_int8_quantized(model)
+        or weights_dtype == "fp32"
     )
     autocast = torch.autocast("cuda") if use_torch_autocast else nullcontext()
 
@@ -556,7 +584,14 @@ def get_onnx_bytes_and_metadata(
         if is_fp4_quantized(model) or is_mxfp8_quantized(model)
         else nullcontext()
     )
-    with torch.inference_mode(), autocast, quantizer_context:
+    # Disable FP8 Conv weight quantizers: TorchScript ONNX exporter requires static
+    # kernel shapes, but FP8 DequantizeLinear produces dynamic shapes.
+    conv_wq_context = (
+        _disable_fp8_conv_weight_quantizers(model)
+        if is_fp8_quantized(model)
+        else nullcontext()
+    )
+    with torch.inference_mode(), autocast, quantizer_context, conv_wq_context:
         additional_kwargs = {}
         if not dynamo_export:
             additional_kwargs["dynamic_axes"] = dynamic_axes
@@ -598,7 +633,12 @@ def get_onnx_bytes_and_metadata(
         onnx_opt_graph = qdq_to_dq(onnx_opt_graph)
 
     if weights_dtype in ["fp16", "bf16"]:
-        if is_int4_quantized(model) or is_mxfp8_quantized(model) or is_fp8_quantized(model):
+        if (
+            is_int4_quantized(model)
+            or is_mxfp8_quantized(model)
+            or is_fp8_quantized(model)
+            or is_int8_quantized(model)
+        ):
             assert weights_dtype == "fp16", "BF16 + MXFP8/INT4 mixed precision is not supported yet"
             onnx_opt_graph = convert_float_to_float16(
                 onnx_opt_graph,

modelopt/torch/quantization/export_onnx.py

@@ -656,9 +656,20 @@ def export_fp4(
 
 @contextlib.contextmanager
 def configure_linear_module_onnx_quantizers(model):
-    """Sets the onnx export attributes for the given model."""
+    """Sets the onnx export attributes for the given model.
+
+    For modules with block quantization (NVFP4/MXFP8):
+    - Weight quantizers use "static" export (TRT_FP4QDQ for NVFP4, DQ-only for MXFP8)
+    - Input/activation quantizers use "dynamic" export (TRT_FP4DynamicQuantize, etc.)
+
+    This must be set for ALL modules with block quantization, not just nn.Linear,
+    because models like ResNet have non-Linear modules (e.g., MaxPool2d) with NVFP4/MXFP8
+    input quantizers that would otherwise default to the static path and produce
+    TRT_FP4QDQ nodes on activations (which the NVFP4 exporter cannot handle).
+    """
     for _, module in model.named_modules():
-        if isinstance(module, torch.nn.Linear):
+        if hasattr(module, "input_quantizer") and module.input_quantizer.block_sizes:
             module.input_quantizer._onnx_quantizer_type = "dynamic"
+        if hasattr(module, "weight_quantizer") and module.weight_quantizer.block_sizes:
             module.weight_quantizer._onnx_quantizer_type = "static"
     yield

tests/_test_utils/torch/vision_models.py

@@ -117,6 +117,7 @@ def get_model_and_input(on_gpu: bool = False):
             # "dm_nfnet_f0",
             "efficientnet_b0",
             "swin_tiny_patch4_window7_224",
+            "resnet50",
         ],
         _create_timm_fn,
     ),

tests/examples/torch_onnx/test_torch_quant_to_onnx.py

@@ -28,6 +28,7 @@
     "vit_tiny": ("vit_tiny_patch16_224", '{"depth": 1}'),
     "swin_tiny": ("swin_tiny_patch4_window7_224", '{"depths": [1, 1, 1, 1]}'),
     "swinv2_tiny": ("swinv2_tiny_window8_256", '{"depths": [1, 1, 1, 1]}'),
+    "resnet50": ("resnet50", None),
 }
 
 # Builder optimization level: 4 for low-bit modes, 3 otherwise