[OMNIML-3349][6114983] Add FP8 MHA quantization support for HuggingFace ViT

CHANGELOG.rst

@@ -15,6 +15,10 @@ Changelog
   - ``from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import ...`` → ``from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import ...``
   - ``from modelopt.torch.sparsity.attention_sparsity.kernels import ...`` → ``from modelopt.torch.kernels.sparsity.attention import ...``
 
+**New Features**
+
+- Add FP8 MHA quantization support for vision transformers. Adds an attention-aware ONNX post-processing pass (scale Mul / K-transpose move before Q, Q→DQ insertion on softmax output) in :class:`FP8QuantExporter <modelopt.onnx.export.fp8_exporter.FP8QuantExporter>`, per-instance nested-attention-wrapper skipping in the HF plugin, and ``nn.LayerNorm`` registration in ``QuantModuleRegistry`` so BMM input quantizers and LayerNorm output quantizers defined in FP8_DEFAULT_CFG are honored end-to-end. See `examples/torch_onnx/torch_quant_to_onnx.py <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/torch_onnx/torch_quant_to_onnx.py>`_ for the general timm-model quantize→ONNX workflow.
+
 0.44 (2026-05-xx)
 ^^^^^^^^^^^^^^^^^
 

examples/torch_onnx/torch_quant_to_onnx.py

@@ -88,6 +88,22 @@
     },
 ]
 
+# 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).
+_FP8_MHA_OVERRIDE: list = [
+    {
+        "parent_class": "nn.LayerNorm",
+        "quantizer_name": "*output_quantizer",
+        "cfg": {"num_bits": (4, 3), "axis": None},
+    },
+    {
+        "parent_class": "nn.LayerNorm",
+        "quantizer_name": "*input_quantizer",
+        "enable": False,
+    },
+]
+
 # 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] = {
@@ -102,11 +118,16 @@ def get_quant_config(quantize_mode):
     """Get quantization config, overriding Conv2d for TRT compatibility.
 
     TensorRT only supports FP8 and INT8 for Conv layers.
+    - For FP8: add MHA-aware LayerNorm output quantizer so TRT fuses shared Q/DQ into
+      downstream attention matmuls. Softmax-output Q/DQ is inserted by the FP8 ONNX
+      exporter's post-processing (fixed 1/448 scale, no calibration needed).
     - For MXFP8, NVFP4: override Conv2d to FP8
     - For INT4_AWQ: override Conv2d to INT8
     """
     config: dict = copy.deepcopy(QUANT_CONFIG_DICT[quantize_mode])
-    if quantize_mode in ("mxfp8", "nvfp4"):
+    if quantize_mode == "fp8":
+        config["quant_cfg"].extend(_FP8_MHA_OVERRIDE)
+    elif quantize_mode in ("mxfp8", "nvfp4"):
         warnings.warn(
             f"TensorRT only supports FP8/INT8 for Conv layers. "
             f"Overriding Conv2d quantization to FP8 for '{quantize_mode}' mode."
@@ -126,6 +147,9 @@ def filter_func(name):
 
     ``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.
     """
     pattern = re.compile(
         r".*(time_emb_proj|time_embedding|conv_in|conv_out|conv_shortcut|add_embedding|"
@@ -135,6 +159,65 @@ def filter_func(name):
     return pattern.match(name) is not None
 
 
+def _disable_high_rank_input_quantizers(model, input_shape, device):
+    """Disable quantizers on Linear/LayerNorm modules that receive 4D+ input.
+
+    TRT's MXFP8/NVFP4 ``DynamicQuantize`` op only supports 2D/3D input, so Swin's
+    per-block ``norm1``, ``downsample.norm``, and top-level ``norm`` (all 4D in Swin
+    but 3D in ViT) must be skipped. A forward pass with hooks identifies them at
+    runtime, so this works across architectures without hardcoded paths.
+    """
+    high_rank: set[str] = set()
+    handles = []
+    for name, mod in model.named_modules():
+        if isinstance(mod, (torch.nn.Linear, torch.nn.LayerNorm)):
+
+            def hook(m, inp, out, _n=name):
+                if inp and hasattr(inp[0], "ndim") and inp[0].ndim > 3:
+                    high_rank.add(_n)
+
+            handles.append(mod.register_forward_hook(hook))
+
+    was_training = model.training
+    model.eval()
+    try:
+        with torch.no_grad():
+            model(torch.randn(input_shape, device=device))
+    finally:
+        for h in handles:
+            h.remove()
+        model.train(was_training)
+
+    if not high_rank:
+        return
+    prefixes = tuple(n + "." for n in high_rank)
+    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.
 
@@ -169,6 +252,28 @@ def load_calibration_data(model, data_size, batch_size, device, with_labels=Fals
         )
 
 
+def _disable_dead_quantizers(model):
+    """Disable quantizers whose calibrated ``amax`` is non-positive or NaN.
+
+    ``export_fp8`` computes ``scale = 448 / amax`` and blows up on ``amax == 0``.
+    This shows up on SwinV2 with ``--no_pretrained``: timm's ``res-post-norm`` scheme
+    zero-inits each block's ``norm1``/``norm2`` weight and bias, so those LayerNorm
+    outputs are exactly zero at init and the MHA override's output_quantizer
+    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()
+
+
 def _calibrate_uncalibrated_quantizers(model, data_loader):
     """Calibrate FP8 quantizers that weren't calibrated by mtq.quantize().
 
@@ -219,6 +324,10 @@ def forward_loop(model):
     if data_loader is not None:
         _calibrate_uncalibrated_quantizers(quantized_model, data_loader)
 
+    # Drop quantizers whose calibration saw only zeros (e.g. SwinV2 zero-init norm1/norm2)
+    # so ``export_fp8`` doesn't divide by zero.
+    _disable_dead_quantizers(quantized_model)
+
     return quantized_model
 
 
@@ -303,6 +412,8 @@ def auto_quantize_model(
     # Disable quantization for specified layers
     mtq.disable_quantizer(quantized_model, filter_func)
 
+    _disable_dead_quantizers(quantized_model)
+
     return quantized_model, search_state
 
 
@@ -458,6 +569,7 @@ def main():
         # Conv2d layers are overridden to FP8 (for TRT compatibility), those FP8
         # quantizers require calibration data.
         config = get_quant_config(args.quantize_mode)
+
         data_loader = load_calibration_data(
             model,
             args.calibration_data_size,
@@ -468,6 +580,25 @@ def main():
 
         quantized_model = quantize_model(model, config, data_loader)
 
+    # 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.
+    uses_dynamic_quantize = args.quantize_mode in ("mxfp8", "nvfp4") or (
+        args.quantize_mode == "auto"
+        and any(fmt in _NEEDS_FP8_CONV_OVERRIDE for fmt in args.auto_quantization_formats)
+    )
+    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)