diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3a158bdb1ae..0ebfa76f3d3 100755 --- a/CHANGELOG.rst +++ b/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. diff --git a/examples/diffusers/quantization/ONNX-TRT-Deployment.md b/examples/diffusers/quantization/ONNX-TRT-Deployment.md index 57448b8a38e..b6933c88f72 100644 --- a/examples/diffusers/quantization/ONNX-TRT-Deployment.md +++ b/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} ``` diff --git a/examples/diffusers/quantization/config.py b/examples/diffusers/quantization/config.py index cb8fdf3a5da..3ab8f5db87c 100644 --- a/examples/diffusers/quantization/config.py +++ b/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): diff --git a/examples/diffusers/quantization/onnx_utils/export.py b/examples/diffusers/quantization/onnx_utils/export.py index 5da795f0f48..246dd1a9881 100644 --- a/examples/diffusers/quantization/onnx_utils/export.py +++ b/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) diff --git a/examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py b/examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py index 7194e672635..7904d27e0a3 100644 --- a/examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py +++ b/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 diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 1d71c088652..9262adbaac1 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/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) diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index c3cfdcd5cdd..b7a79e49e70 100644 --- a/examples/diffusers/quantization/utils.py +++ b/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: diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py index 338e2725b14..42598af5475 100644 --- a/modelopt/onnx/export/nvfp4_exporter.py +++ b/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 diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py index e5778c3c96b..2ad3383e572 100644 --- a/modelopt/torch/quantization/export_onnx.py +++ b/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( diff --git a/modelopt/torch/quantization/plugins/diffusion/diffusers.py b/modelopt/torch/quantization/plugins/diffusion/diffusers.py index f2f6a702479..e92c775c05b 100644 --- a/modelopt/torch/quantization/plugins/diffusion/diffusers.py +++ b/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, ) ) diff --git a/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_conv.yaml b/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_conv.yaml new file mode 100644 index 00000000000..16e2ebbebd2 --- /dev/null +++ b/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 diff --git a/tests/examples/diffusers/test_diffusers.py b/tests/examples/diffusers/test_diffusers.py index 15c5eb44934..979894819de 100644 --- a/tests/examples/diffusers/test_diffusers.py +++ b/tests/examples/diffusers/test_diffusers.py @@ -117,6 +117,17 @@ def inference(self, tmp_path: Path) -> None: quant_algo="smoothquant", collect_method="min-mean", ), + pytest.param( + DiffuserModel( + name="sd3-medium", + path=SD3_PATH, + dtype="Half", + format_type="fp8", + quant_algo="max", + collect_method="default", + ), + marks=minimum_sm(89), + ), pytest.param( DiffuserModel( name="sdxl-1.0", @@ -128,6 +139,17 @@ def inference(self, tmp_path: Path) -> None: ), marks=minimum_sm(89), ), + pytest.param( + DiffuserModel( + name="sdxl-1.0", + path=SDXL_PATH, + dtype="Half", + format_type="fp4", + quant_algo="max", + collect_method="default", + ), + marks=minimum_sm(100), + ), DiffuserModel( name="sdxl-1.0", path=SDXL_PATH, @@ -140,7 +162,9 @@ def inference(self, tmp_path: Path) -> None: ids=[ "flux_schnell_bf16_int8_smoothquant_3.0_min_mean", "sd3_medium_fp16_int8_smoothquant_3.0_min_mean", + "sd3_medium_fp16_fp8_max_3.0_default", "sdxl_1.0_fp16_fp8_max_3.0_default", + "sdxl_1.0_fp16_fp4_max_3.0_default", "sdxl_1.0_fp16_int8_smoothquant_3.0_min_mean", ], ) diff --git a/tests/gpu/torch/quantization/test_onnx_export_cuda.py b/tests/gpu/torch/quantization/test_onnx_export_cuda.py index 300abc52e9e..39f422c75a9 100644 --- a/tests/gpu/torch/quantization/test_onnx_export_cuda.py +++ b/tests/gpu/torch/quantization/test_onnx_export_cuda.py @@ -17,7 +17,6 @@ import pytest import torch -import torch.nn as nn from _test_utils.torch.quantization.onnx_export import TEST_MODELS, onnx_export_tester @@ -40,7 +39,4 @@ def test_onnx_export_cuda(model_cls, num_bits, per_channel_quantization, constan torch.manual_seed(0) model = model_cls() - for _, module in model.named_modules(): - if isinstance(module, nn.Conv2d) and num_bits == (4, 3): - pytest.skip("Conv2d with FP8 quantization is not supported yet") onnx_export_tester(model, "cuda", num_bits, per_channel_quantization, constant_folding, dtype) diff --git a/tests/unit/examples/test_diffusers_fp4.py b/tests/unit/examples/test_diffusers_fp4.py new file mode 100644 index 00000000000..16023e509ad --- /dev/null +++ b/tests/unit/examples/test_diffusers_fp4.py @@ -0,0 +1,234 @@ +# 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. + +import importlib.util +import logging +import sys +from pathlib import Path +from unittest.mock import Mock + +import pytest +import torch +from torch import nn + +pytest.importorskip("onnx") +pytest.importorskip("onnx_graphsurgeon") +pytest.importorskip("diffusers") + +import modelopt.torch.quantization as mtq +from examples.diffusers.quantization.onnx_utils import export as diffusion_export +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.plugins.diffusion import diffusers as diffusers_plugin + +_QUANTIZATION_EXAMPLE = ( + Path(__file__).resolve().parents[3] / "examples" / "diffusers" / "quantization" +) +_LOCAL_IMPORT_NAMES = ( + "calib.plugin_calib", + "calib", + "calibration", + "config", + "models_utils", + "pipeline_manager", + "quantize_config", + "utils", +) + + +def _load_quantize_example(): + spec = importlib.util.spec_from_file_location( + "diffusers_quantize_example", _QUANTIZATION_EXAMPLE / "quantize.py" + ) + assert spec is not None and spec.loader is not None + + original_modules = { + name: sys.modules.pop(name) for name in _LOCAL_IMPORT_NAMES if name in sys.modules + } + sys.path.insert(0, str(_QUANTIZATION_EXAMPLE)) + try: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + finally: + sys.path.pop(0) + for name in _LOCAL_IMPORT_NAMES: + sys.modules.pop(name, None) + sys.modules.update(original_modules) + return module + + +_quantize = _load_quantize_example() +ModelType = _quantize.ModelType +ModelConfig = _quantize.ModelConfig +QuantFormat = _quantize.QuantFormat +QuantizationConfig = _quantize.QuantizationConfig +Quantizer = _quantize.Quantizer +_infer_restored_quantization_format = _quantize._infer_restored_quantization_format + + +class _RecipeBackbone(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(16, 16, bias=False) + self.attn = nn.Module() + self.attn.to_q = nn.Linear(16, 16, bias=False) + self.attn.to_k = nn.Linear(16, 16, bias=False) + self.attn.to_v = nn.Linear(16, 16, bias=False) + self.conv = nn.Conv2d(4, 4, kernel_size=1, bias=False) + + +def _quantizer(*, num_bits, enabled=True, block_sizes=None): + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=num_bits, axis=None, block_sizes=block_sizes) + ) + quantizer.amax = torch.tensor(448.0) + if not enabled: + quantizer.disable() + return quantizer + + +_FP8_QUANTIZER_CONFIG = {"num_bits": (4, 3)} +_NVFP4_QUANTIZER_CONFIG = { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, +} + + +@pytest.mark.parametrize("model_type", [ModelType.SDXL_BASE, ModelType.SDXL_TURBO]) +def test_sdxl_fp4_recipe(model_type): + model = _RecipeBackbone() + config = Quantizer( + QuantizationConfig(format=QuantFormat.FP4), + ModelConfig(model_type=model_type), + logging.getLogger(__name__), + ).get_quant_config(n_steps=1, backbone=model) + + mtq.replace_quant_module(model) + mtq.set_quantizer_by_cfg(model, config["quant_cfg"]) + + for quantizer in (model.linear.input_quantizer, model.linear.weight_quantizer): + assert quantizer.is_enabled + assert quantizer.is_nvfp4_dynamic + assert quantizer.block_sizes[-1] == 16 + for projection in (model.attn.to_q, model.attn.to_k, model.attn.to_v): + assert not projection.input_quantizer.is_enabled + assert not projection.weight_quantizer.is_enabled + for quantizer in (model.conv.input_quantizer, model.conv.weight_quantizer): + assert quantizer.is_enabled + assert quantizer.is_fp8 + + +@pytest.mark.parametrize( + ("format_config", "mha_config", "expected_format", "disable_fp8_mha"), + [ + pytest.param( + _NVFP4_QUANTIZER_CONFIG, + _FP8_QUANTIZER_CONFIG, + QuantFormat.FP4, + False, + id="mixed-fp4", + ), + pytest.param( + _FP8_QUANTIZER_CONFIG, + _FP8_QUANTIZER_CONFIG, + QuantFormat.FP8, + False, + id="fp8", + ), + pytest.param( + {"num_bits": 8}, + {**_FP8_QUANTIZER_CONFIG, "enabled": False}, + QuantFormat.INT8, + True, + id="int8-disabled-fp8", + ), + pytest.param( + {"num_bits": 8}, + {"num_bits": 8}, + QuantFormat.INT8, + True, + id="int8-mha", + ), + ], +) +def test_restored_quantizer_state_drives_format_and_fp8_mha( + monkeypatch, format_config, mha_config, expected_format, disable_fp8_mha +): + backbone = nn.Module() + backbone.quantizer = _quantizer(**format_config) + backbone.attention = nn.Module() + for name in ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer"): + setattr(backbone.attention, name, _quantizer(**mha_config)) + backbone.attention.bmm2_output_quantizer = lambda output: output + + fp8_sdpa = Mock(return_value=torch.empty(0)) + monkeypatch.setattr(diffusers_plugin.FP8SDPA, "apply", fp8_sdpa) + monkeypatch.setattr(torch.onnx, "is_in_onnx_export", lambda: True) + + assert _infer_restored_quantization_format([("transformer", backbone)]) == expected_format + diffusers_plugin._quantized_sdpa(backbone.attention, *(torch.empty(1) for _ in range(3))) + assert fp8_sdpa.call_args.args[-1] is disable_fp8_mha + + +def test_restore_infers_checkpoint_format_for_export(monkeypatch, tmp_path): + backbone = nn.Module() + backbone.quantizer = _quantizer(**_NVFP4_QUANTIZER_CONFIG) + + pipeline_manager = Mock() + pipeline_manager.create_pipeline.return_value = object() + pipeline_manager.iter_backbones.return_value = [("transformer", backbone)] + export_manager = Mock() + monkeypatch.setattr(_quantize, "PipelineManager", lambda *args: pipeline_manager) + monkeypatch.setattr(_quantize, "ExportManager", lambda *args: export_manager) + monkeypatch.setattr( + sys, + "argv", + [ + "quantize.py", + "--model", + "flux-schnell", + "--restore-from", + str(tmp_path), + "--onnx-dir", + str(tmp_path / "onnx"), + ], + ) + + _quantize.main() + + export_manager.restore_checkpoint.assert_called_once_with() + assert export_manager.export_onnx.call_args.args[-1] == QuantFormat.FP4 + export_manager.export_hf_ckpt.assert_called_once() + + +def test_flux_fp8_export_saves_converted_rope_graph(monkeypatch, tmp_path): + original_model = Mock() + converted_model = Mock() + monkeypatch.setattr( + diffusion_export, + "generate_dummy_kwargs_and_dynamic_axes_and_shapes", + lambda *args: ({}, {}, None), + ) + monkeypatch.setattr(diffusion_export, "onnx_export", lambda *args, **kwargs: None) + monkeypatch.setattr(diffusion_export.onnx, "load", lambda *args, **kwargs: original_model) + convert_rope_weight_type = Mock(return_value=converted_model) + monkeypatch.setattr(diffusion_export, "flux_convert_rope_weight_type", convert_rope_weight_type) + save_onnx = Mock() + monkeypatch.setattr(diffusion_export, "save_onnx", save_onnx) + + diffusion_export.modelopt_export_sd(nn.Module(), tmp_path, "flux-dev", "fp8") + + convert_rope_weight_type.assert_called_once_with(original_model) + save_onnx.assert_called_once_with(converted_model, tmp_path / "model.onnx") diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py index 5cdaf9203d0..95944e1f7ab 100644 --- a/tests/unit/onnx/quantization/test_qdq_utils.py +++ b/tests/unit/onnx/quantization/test_qdq_utils.py @@ -42,6 +42,7 @@ replace_zero_scale_with_smallest_nonzero, ) from modelopt.onnx.quantization.quant_utils import pack_float32_to_4bit_cpp_based +from modelopt.onnx.utils import get_opset_version def create_test_model_with_int4_dq_reshape_transpose_matmul(constant_scale: bool = False): @@ -343,8 +344,7 @@ def create_test_model_with_nvfp4_qdq(with_transpose: bool = False): value_info=value_info, ) - model = helper.make_model(graph) - return model + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 20)]) class TestQuantizeWeightsToInt4: @@ -666,6 +666,8 @@ def test_fp4qdq_conversion(self, with_transpose): # Run FP4QDQ to 2DQ conversion converted_model = NVFP4QuantExporter.process_model(model) + assert get_opset_version(converted_model) == 23 + # Verify TRT_FP4QDQ node is removed fp4qdq_nodes = [node for node in converted_model.graph.node if node.op_type == "TRT_FP4QDQ"] assert len(fp4qdq_nodes) == 0 diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py index ce2ef626d63..9b470f7b019 100644 --- a/tests/unit/torch/quantization/test_onnx_export_cpu.py +++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py @@ -39,6 +39,15 @@ from modelopt.torch.quantization.utils import is_quantized_linear +def _export_to_onnx(model, sample_input, **kwargs): + buffer = io.BytesIO() + if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters: + kwargs["enable_onnx_checker"] = False + torch.onnx.export(model, sample_input, buffer, dynamo=False, **kwargs) + buffer.seek(0) + return onnx.load_model_from_string(buffer.read()) + + @pytest.mark.parametrize("model_cls", TEST_MODELS) @pytest.mark.parametrize( ("num_bits", "per_channel_quantization", "constant_folding"), @@ -59,6 +68,41 @@ def test_onnx_export_cpu(model_cls, num_bits, per_channel_quantization, constant ) +def test_fp8_conv_export_preserves_custom_qdq_and_kernel_shape(): + model = torch.nn.Conv2d(3, 4, 3, bias=False).eval() + sample_input = torch.randn(1, 3, 8, 8) + model = mtq.quantize( + model, + mtq.FP8_DEFAULT_CFG, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + + exported_model = _export_to_onnx(model, sample_input, opset_version=20) + producers = {output: node for node in exported_model.graph.node for output in node.output} + conv = next(node for node in exported_model.graph.node if node.op_type == "Conv") + + for conv_input in conv.input[:2]: + dequantize = producers[conv_input] + quantize = producers[dequantize.input[0]] + assert dequantize.op_type == "TRT_FP8DequantizeLinear" + assert quantize.op_type == "TRT_FP8QuantizeLinear" + + value_info = {value.name: value for value in exported_model.graph.value_info} + weight_dequantize = producers[conv.input[1]] + weight_quantize = producers[weight_dequantize.input[0]] + for value_name in (*weight_quantize.output, *weight_dequantize.output): + shape = [ + dimension.dim_value for dimension in value_info[value_name].type.tensor_type.shape.dim + ] + assert shape == [4, 3, 3, 3] + + kernel_shape = next( + attribute for attribute in conv.attribute if attribute.name == "kernel_shape" + ) + assert list(kernel_shape.ints) == [3, 3] + onnx.checker.check_model(exported_model) + + def test_nvfp4_exported_onnx_is_topologically_sorted(monkeypatch): def forward_loop(model): model(sample_input) @@ -78,26 +122,14 @@ def cpu_dynamic_block_quantize(inputs, *args): module.input_quantizer.disable() module.weight_quantizer._onnx_quantizer_type = "static" - buffer = io.BytesIO() - if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters: - kwargs = {"enable_onnx_checker": False} - else: - kwargs = {} - - torch.onnx.export( + exported_model = _export_to_onnx( model, sample_input, - buffer, input_names=["input"], output_names=["output"], export_params=True, opset_version=21, - dynamo=False, - **kwargs, ) - - buffer.seek(0) - exported_model = onnx.load_model_from_string(buffer.read()) assert any(node.op_type == "TRT_FP4QDQ" for node in exported_model.graph.node) converted_model = NVFP4QuantExporter.process_model(exported_model)