[OMNIML-3349] Add ResNet50 support for torch_onnx quantization workflow

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

@@ -17,6 +17,7 @@
 import copy
 import json
 import re
+import subprocess
 import sys
 import warnings
 from pathlib import Path
@@ -35,13 +36,19 @@
 import modelopt.torch.quantization as mtq
 
 """
-This script is used to quantize a timm model using dynamic quantization like MXFP8 or NVFP4,
-or using auto quantization for optimal per-layer quantization.
+Quantize a timm vision model and export to ONNX for TensorRT deployment.
+
+Supports FP8, INT8, MXFP8, NVFP4, and AUTO (mixed-precision) quantization modes end-to-end
+(quantize + ONNX export + TRT build). INT4_AWQ is quantize/export-only; it is not compatible
+with ``--trt_build``.
 
 The script will:
-1. Given the model name, create a timm torch model.
-2. Quantize the torch model in MXFP8, NVFP4, INT4_AWQ, or AUTO mode.
-3. Export the quantized torch model to ONNX format.
+1. Load a pretrained timm model (e.g., ViT, Swin, ResNet).
+2. Quantize the model using the specified mode. For models with Conv2d layers,
+   Conv2d quantization is automatically overridden for TensorRT compatibility
+   (FP8 for MXFP8/NVFP4, INT8 for INT4_AWQ).
+3. Export the quantized model to ONNX with FP16 weights.
+4. Optionally evaluate accuracy on ImageNet-1k before and after quantization.
 """
 
 
@@ -81,6 +88,15 @@
     },
 ]
 
+# Auto-quantize format configs that use block quantization and need Conv2d overrides for TRT.
+# TRT DynamicQuantize requires 2D/3D input, but Conv2d operates on 4D tensors.
+_NEEDS_FP8_CONV_OVERRIDE: set[str] = {
+    "NVFP4_AWQ_LITE_CFG",
+    "NVFP4_DEFAULT_CFG",
+    "MXFP8_DEFAULT_CFG",
+}
+_NEEDS_INT8_CONV_OVERRIDE: set[str] = {"INT4_AWQ_CFG"}
+
 
 def get_quant_config(quantize_mode):
     """Get quantization config, overriding Conv2d for TRT compatibility.
@@ -106,27 +122,33 @@ def get_quant_config(quantize_mode):
 
 
 def filter_func(name):
-    """Filter function to exclude certain layers from quantization."""
+    """Filter function to exclude certain layers from quantization.
+
+    ``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.
+    """
     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"maxpool|global_pool|downsample\.reduction).*"
     )
     return pattern.match(name) is not None
 
 
-def load_calibration_data(model_name, data_size, batch_size, device, with_labels=False):
+def load_calibration_data(model, data_size, batch_size, device, with_labels=False):
     """Load and prepare calibration data.
 
     Args:
-        model_name: Name of the timm model
+        model: The timm model being quantized; used to derive the calibration transforms so the
+               data pipeline matches the exact model config (respects --no_pretrained and
+               --model_kwargs).
         data_size: Number of samples to load
         batch_size: Batch size for data loader
         device: Device to load data to
         with_labels: If True, return dict with 'image' and 'label' keys (for auto_quantize)
                     If False, return just the images (for standard quantize)
     """
     dataset = load_dataset("zh-plus/tiny-imagenet")
-    model = timm.create_model(model_name, pretrained=True, num_classes=1000)
     data_config = timm.data.resolve_model_data_config(model)
     transforms = timm.data.create_transform(**data_config, is_training=False)
 
@@ -147,6 +169,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 quantizer.amax is None:
+                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(strict=False)
+
+
 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,7 +211,14 @@ 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)
+
+    # 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)
+
     return quantized_model
 
 
@@ -209,11 +268,19 @@ def auto_quantize_model(
     _disable_inplace_relu(model)
     constraints = {"effective_bits": effective_bits}
 
-    # Convert string format names to actual config objects
+    # 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):
-            format_configs.append(getattr(mtq, fmt))
+            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)
 
@@ -248,7 +315,10 @@ def get_model_input_shape(model):
 
 def main():
     parser = argparse.ArgumentParser(
-        description="Quantize timm models to FP8, MXFP8, INT8, NVFP4, INT4_AWQ, or use AUTO quantization"
+        description=(
+            "Quantize timm models to FP8, MXFP8, INT8, NVFP4, or use AUTO quantization. "
+            "INT4_AWQ is supported for quantize/export only and is not compatible with --trt_build."
+        )
     )
 
     # Model hyperparameters
@@ -320,6 +390,11 @@ def main():
         default=128,
         help="Number of scoring steps for auto quantization. Default is 128.",
     )
+    parser.add_argument(
+        "--trt_build",
+        action="store_true",
+        help="Build a TensorRT engine from the exported ONNX model using trtexec.",
+    )
     parser.add_argument(
         "--no_pretrained",
         action="store_true",
@@ -362,7 +437,7 @@ def main():
     if args.quantize_mode == "auto":
         # Auto quantization requires labels for loss computation
         data_loader = load_calibration_data(
-            args.timm_model_name,
+            model,
             args.calibration_data_size,
             args.batch_size,
             device,
@@ -378,18 +453,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(
+            model,
+            args.calibration_data_size,
+            args.batch_size,
+            device,
+            with_labels=False,
+        )
 
         quantized_model = quantize_model(model, config, data_loader)
 
@@ -421,6 +496,33 @@ def main():
 
     print(f"Quantized ONNX model is saved to {args.onnx_save_path}")
 
+    if args.trt_build:
+        build_trt_engine(args.onnx_save_path)
+
+
+def build_trt_engine(onnx_path):
+    """Build a TensorRT engine from the exported ONNX model using trtexec."""
+    cmd = [
+        "trtexec",
+        f"--onnx={onnx_path}",
+        "--stronglyTyped",
+        "--builderOptimizationLevel=4",
+    ]
+    print(f"\nBuilding TensorRT engine: {' '.join(cmd)}")
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
+    except FileNotFoundError as e:
+        raise RuntimeError(
+            "trtexec not found on PATH; install TensorRT or drop --trt_build."
+        ) from e
+    except subprocess.TimeoutExpired as e:
+        raise RuntimeError(f"trtexec timed out building {onnx_path} after 600s.") from e
+    if result.returncode != 0:
+        raise RuntimeError(
+            f"TensorRT engine build failed for {onnx_path}:\n{result.stdout}\n{result.stderr}"
+        )
+    print("TensorRT engine build succeeded.")
+
 
 if __name__ == "__main__":
     main()

modelopt/onnx/export/fp8_exporter.py

@@ -101,13 +101,87 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
 
         return gs.export_onnx(graph)
 
+    @staticmethod
+    def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int:
+        """Add FP8 weight DequantizeLinear for Conv layers with unquantized weights.
+
+        Conv weight quantizers are disabled during TorchScript ONNX export because the
+        TRT_FP8DequantizeLinear custom op produces outputs with unknown shapes, causing
+        the _convolution symbolic to fail. This method restores FP8 weight quantization
+        by inserting DQ nodes in the ONNX graph, mirroring the compress_weights logic.
+
+        For each Conv node with an unquantized constant weight:
+        1. Compute per-tensor scale = max(abs(weight)) / 448.0
+        2. Quantize weights to FP8E4M3FN
+        3. Insert a DequantizeLinear(fp8_weights, scale) before the Conv weight input
+
+        Args:
+            graph: The onnx-graphsurgeon graph to modify in-place.
+
+        Returns:
+            Number of Conv weight DQ nodes inserted.
+        """
+        fp8_max = 448.0
+        count = 0
+
+        for node in list(graph.nodes):
+            if node.op != "Conv":
+                continue
+            if len(node.inputs) < 2:
+                continue
+
+            weight_input = node.inputs[1]
+            if not isinstance(weight_input, gs.Constant):
+                continue
+
+            # Skip if weight already has a DQ producer
+            if any(out.op == "DequantizeLinear" for out in weight_input.outputs):
+                continue
+
+            torch_weights = torch.from_numpy(weight_input.values.copy())
+            amax = torch_weights.abs().max().float()
+            if amax == 0:
+                continue
+            scale_val = (amax / fp8_max).item()
+
+            # Quantize weights to FP8 (WAR: numpy doesn't support fp8)
+            fp8_data = (torch_weights / scale_val).to(torch.float8_e4m3fn).view(torch.uint8).numpy()
+            fp8_tensor = onnx.TensorProto()
+            fp8_tensor.data_type = onnx.TensorProto.FLOAT8E4M3FN
+            fp8_tensor.dims.extend(fp8_data.shape)
+            fp8_tensor.raw_data = fp8_data.tobytes()
+            fp8_constant = gs.Constant(
+                node.name + "/weight_quantizer/fp8_weights", LazyValues(fp8_tensor)
+            )
+
+            # Scale in FP16 — DQ output type matches scale dtype, must match activation type
+            import numpy as np
+
+            scale_constant = gs.Constant(
+                node.name + "/weight_quantizer/scale",
+                np.array(scale_val, dtype=np.float16),
+            )
+
+            dq_output = gs.Variable(node.name + "/weight_quantizer/dq_output")
+            dq_node = gs.Node(
+                op="DequantizeLinear",
+                name=node.name + "/weight_quantizer/DequantizeLinear",
+                inputs=[fp8_constant, scale_constant],
+                outputs=[dq_output],
+            )
+            graph.nodes.append(dq_node)
+            node.inputs[1] = dq_output
+            count += 1
+
+        return count
+
     @staticmethod
     def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
         """Post-processes the ONNX model for FP8 quantization.
 
-        Converts TRT_FP8 QDQ ops to native ONNX QuantizeLinear/DequantizeLinear:
-        - TRT_FP8QuantizeLinear -> QuantizeLinear with FP8E4M3FN zero_point and saturate=1
-        - TRT_FP8DequantizeLinear -> DequantizeLinear
+        Converts TRT_FP8 QDQ ops to native ONNX QuantizeLinear/DequantizeLinear and
+        adds FP8 weight DQ for Conv layers whose weight quantizers were disabled during
+        TorchScript export.
 
         Args:
             onnx_model: The ONNX model containing TRT_FP8 quantization nodes.
@@ -144,5 +218,10 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
                     f"Converted {node.name} from TRT_FP8DequantizeLinear to DequantizeLinear"
                 )
 
+        # Add FP8 weight DQ for Conv layers that had weight quantizers disabled during export
+        count = FP8QuantExporter._quantize_conv_weights_to_fp8(graph)
+        if count > 0:
+            logger.info(f"Inserted FP8 weight DequantizeLinear for {count} Conv nodes")
+
         graph.cleanup().toposort()
         return gs.export_onnx(graph)