[5565357] Fix SDXL NVFP4 export and performance

CHANGELOG.rst

@@ -29,6 +29,7 @@ Changelog
 
 **Bug Fixes**
 
+- Fix shared ONNX export metadata and Diffusers attention policy: every ``NVFP4QuantExporter`` post-process now upgrades the default-domain opset to at least 23, all FP8 custom-op exports re-run ONNX shape/type inference after setting output metadata, and quantized SDPA derives FP8 MHA enablement from the live Q/K/V quantizers instead of honoring a caller-set ``_disable_fp8_mha`` attribute.
 - Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations.
 - Fix ``--use_fsdp2`` HuggingFace checkpoint export gathering the whole model onto rank 0, which made export the dominant phase of a PTQ run and could exhaust host memory on large models. The model is now split into per-decoder-layer units dealt round-robin across ranks; each rank gathers every unit but keeps, packs, and writes only the ones it owns, so a rank buffers roughly ``model / world_size`` instead of the whole checkpoint, and rank 0 writes the combined index. Export configurations that cannot be split this way now raise instead of producing a mismatched checkpoint: FSDP2 combined with another DTensor parallelism (for example FSDP2 + tensor parallel on a 2-D mesh; HSDP is supported), models whose decoder layers cannot be discovered, a decoder layer object reused across layers, and a module that holds the decoder layers while owning parameters of its own.
 - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration.

examples/diffusers/quantization/ONNX-TRT-Deployment.md

@@ -28,12 +28,12 @@ python quantize.py \
 
 #### FLUX-Dev|SDXL|SDXL-Turbo|LTX-Video FP8/FP4 [Script](./quantize.py)
 
-*In our example code, FP4 is only supported for Flux. However, you can modify our script to enable FP4 format support for your own model.*
+FP4 ONNX export is supported for Flux and SDXL.
 
 ```sh
 python quantize.py \
     --model {flux-dev|sdxl-1.0|sdxl-turbo|ltx-video-dev} --model-dtype {Half|BFloat16} --trt-high-precision-dtype {Half|BFloat16} \
-    --format {fp8|fp4} --batch-size 2 --calib-size {128|256} --quantize-mha \
+    --format {fp8|fp4} --batch-size 2 --calib-size {128|256} \
     --n-steps 20 --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt --collect-method default \
     --onnx-dir {ONNX_DIR}
 ```

examples/diffusers/quantization/config.py

@@ -31,6 +31,9 @@
 NVFP4_FP8_MHA_CONFIG = load_config(
     "configs/ptq/presets/diffusers/nvfp4_fp8_mha", schema_type=QuantizeConfig
 ).model_dump(exclude_unset=True)
+NVFP4_FP8_CONV_CONFIG = load_config(
+    "configs/ptq/presets/diffusers/nvfp4_fp8_conv", schema_type=QuantizeConfig
+).model_dump(exclude_unset=True)
 
 
 def set_quant_config_attr(quant_config, trt_high_precision_dtype, quant_algo, **kwargs):

examples/diffusers/quantization/onnx_utils/export.py

@@ -51,8 +51,6 @@
 from modelopt.torch.quantization.export_onnx import configure_linear_module_onnx_quantizers
 from modelopt.torch.utils import torch_to
 
-from .fp8_onnx_graphsurgeon import convert_zp_fp8
-
 MODEL_ID_TO_DYNAMIC_AXES = {
     "sdxl-1.0": {
         "sample": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"},
@@ -124,18 +122,6 @@ def flux_convert_rope_weight_type(onnx_graph):
     return gs.export_onnx(graph)
 
 
-def generate_fp8_scales(backbone):
-    # temporary solution due to a known bug in torch.onnx._dynamo_export
-    for _, module in backbone.named_modules():
-        if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)) and (
-            hasattr(module.input_quantizer, "_amax") and module.input_quantizer is not None
-        ):
-            module.input_quantizer._num_bits = 8
-            module.weight_quantizer._num_bits = 8
-            module.input_quantizer._amax = module.input_quantizer._amax * (127 / 448.0)
-            module.weight_quantizer._amax = module.weight_quantizer._amax * (127 / 448.0)
-
-
 def _gen_dummy_inp_and_dyn_shapes_sdxl(backbone, min_bs=1, opt_bs=1):
     assert isinstance(backbone, UNet2DConditionModel) or isinstance(
         backbone._orig_mod, UNet2DConditionModel
@@ -469,7 +455,6 @@ def modelopt_export_sd(backbone, onnx_dir, model_name, precision):
     tmp_subfolder = tempfile.mkdtemp(prefix="myapp_")
     tmp_output = Path(f"{tmp_subfolder}/{model_file_name}")
     q_output = Path(f"{onnx_dir}/{model_file_name}")
-
     quantizer_context = (
         configure_linear_module_onnx_quantizers(backbone) if precision == "fp4" else nullcontext()
     )
@@ -536,16 +521,8 @@ def modelopt_export_sd(backbone, onnx_dir, model_name, precision):
         )
     print(f"Saved at {tmp_output}")
     onnx_model = onnx.load(str(tmp_output), load_external_data=True)
-    if precision == "fp8":
-        if not model_name.startswith("flux"):
-            graph = gs.import_onnx(onnx_model)
-            graph.cleanup().toposort()
-            onnx_model = gs.export_onnx(graph)
-            onnx_model = convert_zp_fp8(onnx_model)
-            graph = gs.import_onnx(onnx_model)
-            onnx_model = gs.export_onnx(graph.cleanup())
-        else:
-            flux_convert_rope_weight_type(onnx_model)
+    if precision == "fp8" and model_name.startswith("flux"):
+        onnx_model = flux_convert_rope_weight_type(onnx_model)
     if precision == "fp4":
         onnx_model = NVFP4QuantExporter.process_model(onnx_model)
     save_onnx(onnx_model, q_output)

examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py

@@ -97,29 +97,6 @@ def insert_cast(graph, input_tensor, attrs):
                 next_node.inputs[idx] = output_tensor
 
 
-def convert_zp_fp8(onnx_graph):
-    """
-    Convert Q/DQ zero datatype from INT8 to FP8.
-    We use this WAR because FP8 Conv cannot be exported to ONNX directly.
-    The workaround is to first convert the FP8 QDQs into INT8 QDQs,
-    then modify the ONNX model afterward to change those INT8 QDQs back into FP8 QDQs.
-    """
-    # Find all zero constant nodes
-    qdq_zero_nodes = set()
-    for node in onnx_graph.graph.node:
-        if node.op_type == "QuantizeLinear" and len(node.input) > 2:
-            qdq_zero_nodes.add(node.input[2])
-
-    print(f"[WAR], found {len(qdq_zero_nodes)} INT8 QDQ pairs, you can ignore this message..")
-
-    # Convert zero point datatype from INT8 to FP8.
-    for node in onnx_graph.graph.node:
-        if node.output[0] in qdq_zero_nodes:
-            node.attribute[0].t.data_type = onnx.TensorProto.FLOAT8E4M3FN
-
-    return onnx_graph
-
-
 def cast_resize_io(graph):
     """
     After all activations and weights are converted to fp16, we will

examples/diffusers/quantization/quantize.py

@@ -27,6 +27,7 @@
     FP8_DEFAULT_CONFIG,
     INT8_DEFAULT_CONFIG,
     NVFP4_DEFAULT_CONFIG,
+    NVFP4_FP8_CONV_CONFIG,
     NVFP4_FP8_MHA_CONFIG,
     reset_set_int8_config,
     set_quant_config_attr,
@@ -55,6 +56,9 @@
 import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export import export_hf_checkpoint
+from modelopt.torch.quantization.nn import TensorQuantizer
+
+_SDXL_MODEL_TYPES = (ModelType.SDXL_BASE, ModelType.SDXL_TURBO)
 
 
 def setup_logging(verbose: bool = False) -> logging.Logger:
@@ -130,7 +134,9 @@ def get_quant_config(self, n_steps: int, backbone: torch.nn.Module) -> Any:
         elif self.config.format == QuantFormat.FP8:
             base_cfg = FP8_DEFAULT_CONFIG
         elif self.config.format == QuantFormat.FP4:
-            if self.model_config.model_type.value.startswith("flux"):
+            if self.model_config.model_type in _SDXL_MODEL_TYPES:
+                base_cfg = NVFP4_FP8_CONV_CONFIG
+            elif self.model_config.model_type.value.startswith("flux"):
                 base_cfg = NVFP4_FP8_MHA_CONFIG
             else:
                 base_cfg = NVFP4_DEFAULT_CONFIG
@@ -271,23 +277,6 @@ def __init__(
         self.logger = logger
         self.pipeline_manager = pipeline_manager
 
-    def _has_conv_layers(self, model: torch.nn.Module) -> bool:
-        """
-        Check if the model contains any convolutional layers.
-
-        Args:
-            model: Model to check
-
-        Returns:
-            True if model contains Conv layers, False otherwise
-        """
-        for module in model.modules():
-            if isinstance(module, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)) and (
-                module.input_quantizer.is_enabled or module.weight_quantizer.is_enabled
-            ):
-                return True
-        return False
-
     def save_checkpoint(
         self,
         backbone: torch.nn.Module,
@@ -335,15 +324,10 @@ def export_onnx(
 
         # Deferred: the ONNX stack (onnx, onnx_graphsurgeon, ...) is only needed
         # for --onnx-dir exports; HF-checkpoint-only runs must not require it.
-        from onnx_utils.export import generate_fp8_scales, modelopt_export_sd
+        from onnx_utils.export import modelopt_export_sd
 
         self.logger.info(f"Starting ONNX export to {self.config.onnx_dir}")
 
-        if quant_format == QuantFormat.FP8 and self._has_conv_layers(backbone):
-            self.logger.info(
-                "Detected quantizing conv layers in backbone. Generating FP8 scales..."
-            )
-            generate_fp8_scales(backbone)
         self.logger.info("Preparing models for export...")
         pipe.to("cpu")
         torch.cuda.empty_cache()
@@ -457,7 +441,7 @@ def create_argument_parser() -> argparse.ArgumentParser:
             %(prog)s --model ltx-video-dev --format fp8 --batch-size 1 --calib-size 32 --ltx-skip-upsampler
 
             # Restore and export a previously quantized model
-            %(prog)s --model flux-schnell --restore-from checkpoint.pt --onnx-dir ./exports/
+            %(prog)s --model flux-schnell --restore-from ./checkpoints/ --onnx-dir ./exports/
         """,
     )
     model_group = parser.add_argument_group("Model Configuration")
@@ -586,7 +570,9 @@ def create_argument_parser() -> argparse.ArgumentParser:
         help="Directory for HuggingFace checkpoint export",
     )
     export_group.add_argument(
-        "--restore-from", type=str, help="Path to restore from previous checkpoint"
+        "--restore-from",
+        type=str,
+        help="Checkpoint directory; quantization format and MHA policy are restored automatically",
     )
     export_group.add_argument(
         "--trt-high-precision-dtype",
@@ -600,6 +586,25 @@ def create_argument_parser() -> argparse.ArgumentParser:
     return parser
 
 
+def _infer_restored_quantization_format(
+    backbones: list[tuple[str, torch.nn.Module]],
+) -> QuantFormat:
+    has_nvfp4 = False
+    has_fp8 = False
+
+    for _, backbone in backbones:
+        for module in backbone.modules():
+            if isinstance(module, TensorQuantizer) and module.is_enabled:
+                has_nvfp4 |= module.is_nvfp4_dynamic or module.is_nvfp4_static
+                has_fp8 |= module.is_fp8
+
+    if has_nvfp4:
+        return QuantFormat.FP4
+    if has_fp8:
+        return QuantFormat.FP8
+    return QuantFormat.INT8
+
+
 def main() -> None:
     from diffusers.models.normalization import RMSNorm as DiffuserRMSNorm
 
@@ -674,9 +679,9 @@ def main() -> None:
         )
 
         logger.info("Validating configurations...")
-        quant_config.validate()
         export_config.validate()
         if not export_config.restore_from:
+            quant_config.validate()
             calib_config.validate()
 
         pipeline_manager = PipelineManager(model_config, logger)
@@ -685,8 +690,12 @@ def main() -> None:
 
         export_manager = ExportManager(export_config, logger, pipeline_manager)
 
-        if export_config.restore_from and export_config.restore_from.exists():
+        if export_config.restore_from:
             export_manager.restore_checkpoint()
+            quant_config.format = _infer_restored_quantization_format(
+                list(pipeline_manager.iter_backbones())
+            )
+            logger.info(f"Detected restored quantization format: {quant_config.format.value}")
 
         else:
             logger.info("Initializing calibration...")
@@ -716,11 +725,12 @@ def forward_loop(mod):
                     mtq.compress(backbone)
                     logger.info(f"{backbone_name} compression completed")
 
-                # For VAE backbones, skip check_conv_and_mha — the whole point
-                # of VAE quantization is to quantize Conv layers.
                 if backbone_name not in ("video_decoder", "vae"):
                     check_conv_and_mha(
-                        backbone, quant_config.format == QuantFormat.FP4, quant_config.quantize_mha
+                        backbone,
+                        quant_config.format == QuantFormat.FP4
+                        and model_config.model_type not in _SDXL_MODEL_TYPES,
+                        quant_config.quantize_mha,
                     )
 
                 export_manager.save_checkpoint(backbone, backbone_name)

examples/diffusers/quantization/utils.py

@@ -64,11 +64,8 @@ def check_conv_and_mha(backbone, if_fp4, quantize_mha):
                 ):
                     if hasattr(module, attr):
                         getattr(module, attr).disable()
-                setattr(module, "_disable_fp8_mha", True)
 
                 print(f"Disabled Attention layer quantization for layer {name}")
-            else:
-                setattr(module, "_disable_fp8_mha", False)
 
 
 def filter_func_ltx_video(name: str) -> bool:

modelopt/onnx/export/nvfp4_exporter.py

@@ -430,4 +430,10 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str):
 
         utils.topologically_sort_graph_nodes(graph)
 
+        if fp4_qdq_nodes:
+            default_opset = next(
+                opset for opset in onnx_model.opset_import if opset.domain in {"", "ai.onnx"}
+            )
+            default_opset.version = max(default_opset.version, 23)
+
         return onnx_model

modelopt/torch/quantization/export_onnx.py

@@ -225,9 +225,12 @@ def _fp8_quantize(
         "Constant",
         value_t=torch.tensor(scale_inv).to(torch_dtype_map[inputs.type().scalarType()]),
     )
-    return g.op("trt::TRT_FP8QuantizeLinear", inputs, scale).setType(
+    quantized = g.op("trt::TRT_FP8QuantizeLinear", inputs, scale).setType(
         inputs.type().with_dtype(torch.uint8).with_sizes(output_shape)
     )
+    # PyTorch runs shape inference before setType for custom ops, so refresh its reliability state.
+    torch._C._jit_pass_onnx_node_shape_type_inference(quantized.node(), g.params_dict, g.opset)
+    return quantized
 
 
 def _fp8_dequantize(

modelopt/torch/quantization/plugins/diffusion/diffusers.py

@@ -141,6 +141,14 @@ def _quantized_sdpa(self, *args, **kwargs):
     q_quantized_scale = self.q_bmm_quantizer._get_amax(query)
     k_quantized_scale = self.k_bmm_quantizer._get_amax(key)
     v_quantized_scale = self.v_bmm_quantizer._get_amax(value)
+    disable_fp8_mha = not all(
+        quantizer.is_enabled and quantizer.is_fp8
+        for quantizer in (
+            self.q_bmm_quantizer,
+            self.k_bmm_quantizer,
+            self.v_bmm_quantizer,
+        )
+    )
 
     # We don't need to calibrate the output of softmax
     return self.bmm2_output_quantizer(
@@ -155,7 +163,7 @@ def _quantized_sdpa(self, *args, **kwargs):
             self.q_bmm_quantizer.trt_high_precision_dtype
             if hasattr(self.q_bmm_quantizer, "trt_high_precision_dtype")
             else "Half",
-            self._disable_fp8_mha if hasattr(self, "_disable_fp8_mha") else True,
+            disable_fp8_mha,
         )
     )
 

modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_conv.yaml

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Diffusers SDXL preset with dynamic NVFP4 Linears and per-tensor FP8 Conv2d layers.
+
+# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig
+imports:
+  base_disable_all: configs/ptq/units/base_disable_all
+  fp8: configs/numerics/fp8
+  nvfp4: configs/numerics/nvfp4
+
+algorithm: max
+quant_cfg:
+  - $import: base_disable_all
+  - parent_class: nn.Linear
+    quantizer_name: '*weight_quantizer'
+    cfg:
+      $import: nvfp4
+  - parent_class: nn.Linear
+    quantizer_name: '*input_quantizer'
+    cfg:
+      $import: nvfp4
+  - parent_class: nn.Linear
+    quantizer_name: '*to_[qkv].input_quantizer'
+    enable: false
+  - parent_class: nn.Linear
+    quantizer_name: '*to_[qkv].weight_quantizer'
+    enable: false
+  - parent_class: nn.Conv2d
+    quantizer_name: '*weight_quantizer'
+    cfg:
+      $import: fp8
+  - parent_class: nn.Conv2d
+    quantizer_name: '*input_quantizer'
+    cfg:
+      $import: fp8
+  - quantizer_name: '*output_quantizer'
+    enable: false
+  - quantizer_name: '*softmax_quantizer'
+    cfg:
+      $import: fp8